diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 1380b07de53..fb91bcce77f 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -128,6 +128,8 @@ fi sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt" awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt" +export UV_OFFLINE=true + if [ "$mode" = parity ]; then export INTEGRATION_WORKERS=1 else diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index aafd968ee0f..4f1a18e925e 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -6,12 +6,13 @@ ## Handwritten 1-3 sentence description - diff --git a/.github/workflows/compat-matrix-image.yml b/.github/workflows/compat-matrix-image.yml index f08792c904c..ddd2aa1a50b 100644 --- a/.github/workflows/compat-matrix-image.yml +++ b/.github/workflows/compat-matrix-image.yml @@ -4,7 +4,7 @@ on: pull_request: paths: - tests/e2e/claude_code/cron_vm/** - - tests/e2e/claude_code/pr_gate_version_resolver.py + - tests/e2e/cli_harness/npm_version_resolver.py - .github/workflows/compat-matrix-image.yml workflow_dispatch: @@ -36,7 +36,7 @@ jobs: whoami gh --version uv --version - version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/claude_code/pr_gate_version_resolver.py)" + version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)" /opt/litellm/tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" /tmp/claude-cli /tmp/claude-cli/claude --version ' diff --git a/.github/workflows/deploy-dev-on-main.yml b/.github/workflows/deploy-dev-on-main.yml deleted file mode 100644 index 1208a2dbbdb..00000000000 --- a/.github/workflows/deploy-dev-on-main.yml +++ /dev/null @@ -1,30 +0,0 @@ -name: Deploy main to gateway-dev - -on: - push: - branches: [main] - -permissions: {} - -concurrency: - group: deploy-dev-dispatch - cancel-in-progress: true - -jobs: - dispatch: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 5 - steps: - - uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2 - id: token - with: - app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }} - private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }} - owner: BerriAI - repositories: litellm-ops - permission-actions: write - - name: Dispatch the litellm-ops dev deploy - env: - GH_TOKEN: ${{ steps.token.outputs.token }} - run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index c3388430589..cbeea71ad0f 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -208,7 +208,7 @@ jobs: exit 0 fi retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; } - CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)" + CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)" tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli" PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \ diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2c05ee827a3..0a490bb8f45 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -441,6 +441,7 @@ jobs: tests/unit/ocr tests/unit/passthrough tests/unit/rag + tests/unit/cookbook tests/unit/realtime_api tests/unit/rerank_api tests/unit/search diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml deleted file mode 100644 index 0ba4df89cc4..00000000000 --- a/.github/workflows/weekly_load_anomaly.yml +++ /dev/null @@ -1,85 +0,0 @@ -name: "Weekly Load Anomaly Check" - -on: - schedule: - - cron: "0 12 * * 6" - workflow_dispatch: - -permissions: - contents: read - -jobs: - weekly-load-anomaly: - if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 45 - services: - postgres: - image: postgres:16.6 - env: - POSTGRES_USER: llmproxy - POSTGRES_PASSWORD: dbpassword9090 - POSTGRES_DB: litellm - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready -U llmproxy" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - env: - DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm - LITELLM_MASTER_KEY: sk-weekly-anomaly-check - ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} - AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }} - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache the Rust build - uses: ./.github/actions/cache-cargo-build - - - name: Install dependencies - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy - - - name: Cache Prisma binaries - uses: ./.github/actions/cache-prisma-binaries - - - name: Generate Prisma client - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Start the proxy - run: | - nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 & - for _ in $(seq 1 90); do - if curl -fs http://localhost:4000/health/liveliness > /dev/null; then - exit 0 - fi - sleep 2 - done - echo "proxy never became live" - tail -n 100 proxy.log - exit 1 - - - name: Run the weekly session anomaly test - env: - E2E_WEEKLY_ANOMALY: "1" - run: | - uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA - - - name: Show proxy log on failure - if: failure() - run: tail -n 300 proxy.log diff --git a/README.md b/README.md index 3b5710fa0c5..c4aeee80fb4 100644 --- a/README.md +++ b/README.md @@ -444,6 +444,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [OpenAI-like (`openai_like`)](https://docs.litellm.ai/docs/providers/openai_compatible) | | | | ✅ | | | | | | | | [OpenRouter (`openrouter`)](https://docs.litellm.ai/docs/providers/openrouter) | ✅ | ✅ | ✅ | | | | | | | | | [OVHCloud AI Endpoints (`ovhcloud`)](https://docs.litellm.ai/docs/providers/ovhcloud) | ✅ | ✅ | ✅ | | | | | | | | +| [Parallel AI (`parallel_ai`)](https://docs.litellm.ai/docs/search/parallel_ai) | | | | | | | | | | | | [Perplexity AI (`perplexity`)](https://docs.litellm.ai/docs/providers/perplexity) | ✅ | ✅ | ✅ | | | | | | | | | [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | | | [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | | diff --git a/cookbook/gpt_realtime_translate.py b/cookbook/gpt_realtime_translate.py new file mode 100644 index 00000000000..827c65cdbfe --- /dev/null +++ b/cookbook/gpt_realtime_translate.py @@ -0,0 +1,245 @@ +#!/usr/bin/env python3 + +import argparse +import asyncio +import base64 +import json +import os +import sys +import wave +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from pathlib import Path +from urllib.parse import urlencode, urlsplit, urlunsplit + +import websockets +from websockets.asyncio.client import ClientConnection + +SAMPLE_RATE = 24_000 +CHANNELS = 1 +SAMPLE_WIDTH = 2 +CHUNK_DURATION_SECONDS = 0.1 +CHUNK_BYTES = int(SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH * CHUNK_DURATION_SECONDS) +RESPONSE_TIMEOUT_SECONDS = 30.0 +AUDIO_EVENT_TYPES = frozenset( + { + "session.output_audio.delta", + "response.audio.delta", + "response.output_audio.delta", + } +) +TRANSCRIPT_EVENT_TYPES = frozenset( + { + "session.output_transcript.delta", + "response.text.delta", + "response.output_audio_transcript.delta", + } +) + + +@dataclass(frozen=True, slots=True) +class Settings: + input_wav: Path + output_wav: Path + base_url: str + model: str + target_language: str + trailing_silence_seconds: float + api_key: str + + +def write_stdout(message: str = "", *, end: str = "\n", flush: bool = False) -> None: + sys.stdout.write(f"{message}{end}") + if flush: + sys.stdout.flush() + + +def write_stderr(message: str) -> None: + sys.stderr.write(f"{message}\n") + + +def parse_args(argv: Sequence[str] | None = None) -> Settings | str: + parser = argparse.ArgumentParser( + description="Stream a 24 kHz PCM16 WAV through gpt-realtime-translate and save the translated audio", + ) + parser.add_argument("input_wav", type=Path) + parser.add_argument("--output", type=Path, default=Path("translated.wav")) + parser.add_argument("--base-url", default=os.getenv("LITELLM_BASE_URL", "http://localhost:4000")) + parser.add_argument("--model", default=os.getenv("REALTIME_TRANSLATE_MODEL", "gpt-realtime-translate")) + parser.add_argument("--target-language", default="fr") + parser.add_argument("--trailing-silence", type=float, default=1.5) + parsed = parser.parse_args(argv) + api_key = os.getenv("LITELLM_API_KEY") or os.getenv("OPENAI_API_KEY") + if not api_key: + return "Set LITELLM_API_KEY or OPENAI_API_KEY before running the script" + if parsed.trailing_silence < 0: + return "--trailing-silence must be zero or greater" + return Settings( + input_wav=parsed.input_wav, + output_wav=parsed.output, + base_url=parsed.base_url, + model=parsed.model, + target_language=parsed.target_language, + trailing_silence_seconds=parsed.trailing_silence, + api_key=api_key, + ) + + +def translation_url(base_url: str, model: str) -> str | None: + parsed = urlsplit(base_url.rstrip("/")) + scheme = {"http": "ws", "https": "wss", "ws": "ws", "wss": "wss"}.get(parsed.scheme) + if not scheme or not parsed.netloc: + return None + base_path = parsed.path.rstrip("/") + realtime_path = ( + f"{base_path}/realtime/translations" if base_path.endswith("/v1") else f"{base_path}/v1/realtime/translations" + ) + return urlunsplit((scheme, parsed.netloc, realtime_path, urlencode({"model": model}), "")) + + +def read_pcm16_wav(path: Path) -> bytes | str: + try: + with wave.open(str(path), "rb") as source: + actual_format = ( + source.getnchannels(), + source.getsampwidth(), + source.getframerate(), + source.getcomptype(), + ) + expected_format = (CHANNELS, SAMPLE_WIDTH, SAMPLE_RATE, "NONE") + if actual_format != expected_format: + return ( + f"{path} must be mono, 16-bit PCM, 24 kHz WAV; received " + f"channels={actual_format[0]}, sample_width={actual_format[1]}, " + f"sample_rate={actual_format[2]}, compression={actual_format[3]}" + ) + return source.readframes(source.getnframes()) + except (OSError, EOFError, wave.Error) as exc: + return f"Unable to read {path}: {exc}" + + +def audio_chunks(audio: bytes) -> Iterator[bytes]: + return (audio[offset : offset + CHUNK_BYTES] for offset in range(0, len(audio), CHUNK_BYTES)) + + +def audio_message(audio: bytes) -> str: + return json.dumps( + { + "type": "session.input_audio_buffer.append", + "audio": base64.b64encode(audio).decode("ascii"), + } + ) + + +async def configure_session(connection: ClientConnection, target_language: str) -> str | None: + await connection.send( + json.dumps( + { + "type": "session.update", + "session": {"audio": {"output": {"language": target_language}}}, + } + ) + ) + while True: + raw_event = await asyncio.wait_for(connection.recv(), timeout=20) + event = json.loads(raw_event) + event_type = event.get("type") + if event_type == "session.created": + write_stdout(f"Session: {event.get('session', {}).get('id', 'created')}") + if event_type == "session.updated": + return None + if event_type == "error": + return f"Realtime API error: {json.dumps(event.get('error', event), ensure_ascii=False)}" + + +async def send_audio( + connection: ClientConnection, pcm: bytes, trailing_silence_seconds: float, finished: asyncio.Event +) -> None: + silence = bytes(round(trailing_silence_seconds * SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH)) + try: + for chunk in audio_chunks(pcm + silence): + await connection.send(audio_message(chunk)) + await asyncio.sleep(len(chunk) / (SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH)) + await connection.send(json.dumps({"type": "session.close"})) + finally: + finished.set() + + +async def receive_translation( + connection: ClientConnection, output_path: Path, sender_finished: asyncio.Event +) -> str | None: + audio_received = asyncio.Event() + try: + with wave.open(str(output_path), "wb") as output: + output.setnchannels(CHANNELS) + output.setsampwidth(SAMPLE_WIDTH) + output.setframerate(SAMPLE_RATE) + write_stdout("Translation: ", end="", flush=True) + while True: + try: + raw_event = await asyncio.wait_for(connection.recv(), timeout=RESPONSE_TIMEOUT_SECONDS) + except TimeoutError: + if not sender_finished.is_set(): + continue + return ( + f"The translation stream went quiet for {RESPONSE_TIMEOUT_SECONDS:g} s without closing the session" + ) + event = json.loads(raw_event) + event_type = event.get("type") + if event_type in AUDIO_EVENT_TYPES: + output.writeframes(base64.b64decode(event.get("delta", ""), validate=True)) + audio_received.set() + elif event_type in TRANSCRIPT_EVENT_TYPES: + write_stdout(event.get("delta", event.get("text", "")), end="", flush=True) + elif event_type == "error": + return f"Realtime API error: {json.dumps(event.get('error', event), ensure_ascii=False)}" + elif event_type == "session.closed": + write_stdout() + return None if audio_received.is_set() else "The translation stream ended without translated audio" + except (OSError, wave.Error) as exc: + return f"Unable to write {output_path}: {exc}" + + +async def translate(settings: Settings, pcm: bytes) -> str | None: + url = translation_url(settings.base_url, settings.model) + if not url: + return f"Invalid --base-url: {settings.base_url}" + sender_finished = asyncio.Event() + try: + async with websockets.connect( + url, + additional_headers={"Authorization": f"Bearer {settings.api_key}"}, + proxy=None, + open_timeout=20, + close_timeout=5, + ) as connection: + configuration_error = await configure_session(connection, settings.target_language) + if configuration_error: + return configuration_error + async with asyncio.TaskGroup() as tasks: + receiver = tasks.create_task(receive_translation(connection, settings.output_wav, sender_finished)) + tasks.create_task(send_audio(connection, pcm, settings.trailing_silence_seconds, sender_finished)) + return receiver.result() + except Exception as exc: + return f"Translation failed: {type(exc).__name__}: {exc}" + + +def main(argv: Sequence[str] | None = None) -> int: + settings = parse_args(argv) + if isinstance(settings, str): + write_stderr(settings) + return 2 + pcm = read_pcm16_wav(settings.input_wav) + if isinstance(pcm, str): + write_stderr(pcm) + return 2 + error = asyncio.run(translate(settings, pcm)) + if error: + write_stderr(error) + return 1 + write_stdout(f"Translated audio: {settings.output_wav}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 0d46efbbb08..91c5e2c5f39 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,9 +3,11 @@ import asyncio import base64 +import copy import json import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta from itertools import chain from types import MappingProxyType @@ -34,7 +36,7 @@ from typing_extensions import ReadOnly, Unpack import litellm from litellm import Router, verbose_logger -from litellm._internal_context import with_service_target +from litellm._internal_context import internal_sub_call, post_response_phase, with_service_target from litellm._uuid import uuid from litellm.caching.caching import DualCache from litellm.constants import ( @@ -95,6 +97,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_metadata, ) from litellm.repositories.managed_file_repository import ManagedFileRepository +from litellm.responses.streaming_iterator import spawn_logging_task from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue] AllMessageValues, @@ -395,6 +398,134 @@ class _RouterFileCallKwargs(TypedDict, total=False): custom_llm_provider: ReadOnly[str | None] +_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str]) +_FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"}) +_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id", "_litellm_internal_model_credentials"}) + + +@dataclass(frozen=True, slots=True) +class _ProviderUpload: + model_id: str + provider_file_id: str + upload_call_id: str | None = None + + +def _provider_uploads(file_object: OpenAIFileObject) -> tuple[_ProviderUpload, ...]: + file_hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic file metadata + dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR) + ) + model_file_id_mapping: Final = _MODEL_FILE_ID_MAPPING.validate_python( + file_hidden_params.get("model_file_id_mapping") or {} + ) + upload_call_id: Final = file_hidden_params.get("litellm_call_id") + return tuple( + _ProviderUpload( + model_id=model_id, + provider_file_id=provider_file_id, + upload_call_id=upload_call_id if isinstance(upload_call_id, str) else None, + ) + for model_id, provider_file_id in model_file_id_mapping.items() + ) + + +def _keeps_one_upload_per_target(llm_router: Router, model_id: str) -> bool: + credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id) + return (credentials or {}).get("custom_llm_provider") == LlmProviders.LITELLM_PROXY.value + + +def _is_kept_upload(uploads: Sequence[_ProviderUpload], index: int, per_target_model_ids: frozenset[str]) -> bool: + upload: Final = uploads[index] + other_uploads: Final = uploads[index + 1 :] if upload.model_id in per_target_model_ids else uploads[:index] + return all(other.model_id != upload.model_id for other in other_uploads) + + +def _upload_from_trailing_id_fields(decoded_unified_file_id: str) -> _ProviderUpload: + trailing_fields: Final = decoded_unified_file_id.rpartition(";llm_output_file_id,")[2] + provider_file_id, _, model_id = trailing_fields.partition(";llm_output_file_model_id,") + return _ProviderUpload(model_id=model_id, provider_file_id=provider_file_id) + + +def _carried_upload_missing_from(unified_file_id: str, model_file_ids: Mapping[str, str]) -> _ProviderUpload | None: + decoded: Final = _is_base64_encoded_unified_file_id(unified_file_id) + carried: Final = _upload_from_trailing_id_fields(decoded) if decoded else None + if carried is None or model_file_ids.get(carried.model_id) in (None, carried.provider_file_id): + return None + return carried + + +async def _delete_provider_file( + llm_router: Router, + model_id: str, + provider_file_id: str, + request_data: Mapping[str, object], +) -> None: + credentials: Final = llm_router.registry().credentials(model_id, form="with_provider") + delete_data: Final = { + **{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS}, + **( + {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))} + if credentials is not None + else {} + ), + } + await llm_router.afile_delete( + model=model_id, + file_id=provider_file_id, + **cast(_RouterFileCallKwargs, delete_data), + ) + + +def _with_own_nested_dicts(request_data: Mapping[str, object]) -> MappingProxyType[str, object]: + return MappingProxyType({k: copy.copy(v) if isinstance(v, dict) else v for k, v in request_data.items()}) + + +async def _delete_redundant_upload( + llm_router: Router, + upload: _ProviderUpload, + delete_request: Mapping[str, object], + parent_call_id: str | None, +) -> None: + try: + with internal_sub_call(parent_call_id): + await _delete_provider_file( + llm_router, upload.model_id, upload.provider_file_id, _with_own_nested_dicts(delete_request) + ) + except Exception as error: + verbose_logger.warning( + f"Could not delete the redundant upload provider_file_id={sanitize_for_log(upload.provider_file_id)} " + f"at deployment model_id={sanitize_for_log(upload.model_id)}: " + f"{type(error).__name__} {sanitize_for_log(error)}" + ) + + +async def _delete_redundant_uploads( + llm_router: Router, + uploads: Sequence[_ProviderUpload], + delete_request: Mapping[str, object], +) -> None: + await asyncio.gather( + *(_delete_redundant_upload(llm_router, upload, delete_request, upload.upload_call_id) for upload in uploads) + ) + + +def _start_deleting_redundant_uploads( + llm_router: Router, + uploads: Sequence[_ProviderUpload], + create_file_request: CreateFileRequest, +) -> None: + if not uploads: + return + delete_request: Final = MappingProxyType( + {k: v for k, v in create_file_request.items() if k not in _FILE_UPLOAD_ONLY_FIELDS} + ) + with post_response_phase(): + spawn_logging_task( + asyncio.get_running_loop(), + _delete_redundant_uploads(llm_router, uploads, delete_request), + task_name="managed file repeat upload cleanup", + ) + + def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase: return prisma_client.db @@ -2034,16 +2165,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): target_model_names_list=target_model_names_list, ) - ## STORE MODEL MAPPINGS IN DB - model_mappings: Dict[str, str] = {} - - for file_object in responses: - file_hidden_params = cast( # cast-ok: preserve mapping operations on dynamic file metadata - dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR) - ) - model_file_id_mapping = file_hidden_params.get("model_file_id_mapping") - if model_file_id_mapping and isinstance(model_file_id_mapping, dict): - model_mappings.update(model_file_id_mapping) + uploads: Final = tuple(chain.from_iterable(map(_provider_uploads, responses))) + per_target_model_ids: Final = frozenset( + model_id + for model_id in {upload.model_id for upload in uploads} + if _keeps_one_upload_per_target(llm_router, model_id) + ) + kept_uploads: Final = tuple( + upload for index, upload in enumerate(uploads) if _is_kept_upload(uploads, index, per_target_model_ids) + ) + model_mappings: Final = {upload.model_id: upload.provider_file_id for upload in kept_uploads} await self.store_unified_file_id( file_id=response.id, @@ -2052,6 +2183,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings=model_mappings, user_api_key_dict=user_api_key_dict, ) + _start_deleting_redundant_uploads( + llm_router, + tuple( + upload + for upload in uploads + if upload not in kept_uploads and upload.model_id not in per_target_model_ids + ), + create_file_request, + ) # Emit Prometheus metrics for managed file creation prom_logger = self._get_prometheus_logger() @@ -2692,23 +2832,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id) if not specific_model_file_id_mapping: return - filtered_data: Final = { - k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials") - } for model_id, model_file_id in specific_model_file_id_mapping.items(): - credentials = llm_router.registry().credentials(model_id, form="with_provider") - delete_data = { - **filtered_data, - **( - {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))} - if credentials is not None - else {} - ), - } - await llm_router.afile_delete( - model=model_id, - file_id=model_file_id, - **cast(_RouterFileCallKwargs, delete_data), + await _delete_provider_file(llm_router, model_id, model_file_id, data) + unrecorded_upload: Final = _carried_upload_missing_from(file_id, specific_model_file_id_mapping) + delete_call_id: Final = data.get("litellm_call_id") + if unrecorded_upload is not None: + await _delete_redundant_upload( + llm_router, unrecorded_upload, data, delete_call_id if isinstance(delete_call_id, str) else None ) async def afile_content( diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index f65516265c5..ab887bbc014 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -33,6 +33,17 @@ "api_key", "additional_args" ], + "record_response": [ + "logger", + "status", + "headers", + "body" + ], + "record_api_window": [ + "logger", + "start", + "end" + ], "defers_async_logging": [ "logger" ], @@ -91,6 +102,13 @@ "logger", "head" ], + "stamp_stream_head": [ + "head", + "logger", + "kwargs", + "start_time", + "end_time" + ], "stream_success": [ "logger", "url_route", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 9431b9d1326..1bb84184e4c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -57,6 +57,8 @@ pub struct LegacyLogging { asynchronous: bool, internal: bool, cache_key: Option, + /// When the provider request went out, until its window is recorded (Messages only). + api_started: Option, } fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult> { @@ -87,6 +89,60 @@ impl LegacyLogging { asynchronous, internal: false, cache_key: None, + api_started: None, + } + } + + /// The kwargs `Logging.update_from_kwargs` sees. Python's Messages handler takes `metadata` as + /// its own argument, so the logger never sees the caller's raw metadata, only the validated + /// request metadata merged in through `litellm_params`. Python updates the logger once, before + /// the call; the refresh once the result is ready leaves out `litellm_metadata` too, so the + /// logged metadata is not rebuilt from what the call wrote into it since. + fn logging_kwargs<'py>(&self, py: Python<'py>, refresh: bool) -> PyResult> { + let kwargs = self.call.kwargs().bind(py).copy()?; + if self.operation != LoggingOperation::Messages { + return Ok(kwargs); + } + let consumed: &[&str] = if refresh { + &["metadata", "litellm_metadata"] + } else { + &["metadata"] + }; + for name in consumed { + if kwargs.contains(name)? { + kwargs.del_item(name)?; + } + } + Ok(kwargs) + } + + /// Python's Messages handler hands the logger its request params, validated `metadata` + /// included, as `litellm_params`, beside the `aanthropic_messages` call flag. + fn logged_request_params( + &self, + context: &RequestContext, + ) -> Option> { + if self.operation != LoggingOperation::Messages { + return None; + } + let mut params = context + .optional_params + .as_object() + .cloned() + .unwrap_or_default(); + params.insert( + "aanthropic_messages".to_owned(), + Value::Bool(self.asynchronous), + ); + Some(params) + } + + /// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync + /// client is untimed. + fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> { + match self.api_started.take() { + Some(start) => self.logger()?.record_api_window(py, start, end), + None => Ok(()), } } @@ -111,6 +167,16 @@ impl LegacyLogging { } } + /// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key. + fn logged_api_key<'a>(&self, context: &'a RequestContext) -> Option<&'a str> { + match self.operation { + LoggingOperation::Messages => Some(""), + LoggingOperation::Completion | LoggingOperation::Responses | LoggingOperation::Ocr => { + context.api_key.as_ref().map(|api_key| api_key.expose()) + } + } + } + fn stream_billing(&self) -> Option { match self.operation { LoggingOperation::Messages => Some(PassThroughStream { @@ -377,7 +443,13 @@ impl LegacyLogging { context: &RequestContext, ) -> PyResult>> { let logger = self.logger()?; - logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?; + logger.update_from_kwargs( + py, + &self.logging_kwargs(py, false)?, + &wire, + context, + self.logged_request_params(context).as_ref(), + )?; let body = to_py(py, &wire.body)? .into_bound(py) .cast_into::()?; @@ -400,11 +472,14 @@ impl LegacyLogging { self.logger()?.pre_call( py, self.input_description(), - context.api_key.as_ref().map(|api_key| api_key.expose()), + self.logged_api_key(context), &body, &headers, &wire.url, )?; + if self.operation == LoggingOperation::Messages && self.asynchronous { + self.api_started = Some(epoch_seconds()); + } let headers = headers .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) @@ -455,7 +530,7 @@ impl LegacyLogging { py, ( &logger, - self.call.kwargs(), + self.logging_kwargs(py, true)?, &facts.provider.model, logger.getattr("optional_params")?, params, @@ -483,6 +558,16 @@ impl LegacyLogging { py: Python<'_>, raw: &RawResponse, ) -> PyResult> { + if self.operation == LoggingOperation::Messages { + // Python's Messages handler never calls `post_call` at receipt: it stores the raw + // response, and logging replays it through `AnthropicConfig.transform_response`. + self.close_api_window(py, epoch_seconds())?; + if let Some(head) = &raw.head { + self.logger()? + .record_response(py, head.status, &head.headers, &raw.body)?; + } + return Ok(HookStep::Ready(())); + } let api_key = self .request .as_ref() @@ -522,6 +607,9 @@ impl LegacyLogging { ) -> PyResult> { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); + if self.logger.is_some() { + self.close_api_window(py, timing.end_time)?; + } if self.stream.is_some() { return self.stream_failure(py); } @@ -543,6 +631,20 @@ impl LegacyLogging { head.bind(py).set_item("cache_key", key)?; head.bind(py).set_item("cache_hit", true)?; } + self.close_api_window(py, epoch_seconds())?; + if self.operation == LoggingOperation::Messages { + // Python's `@client` stamps the call metadata onto the stream it hands back. + Streaming::StampHead.call( + py, + ( + head.bind(py), + self.logger()?.object(py), + self.call.kwargs(), + &self.start, + datetime(py, epoch_seconds())?, + ), + )?; + } Streaming::Opened.call(py, (self.logger()?.object(py), head.bind(py)))?; self.stream = Some(DeliveredStream { chunks: PyList::empty(py).unbind(), @@ -1117,9 +1219,7 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update let step = logging .before_provider_request(py, Box::new(wire), context) .unwrap(); - let raw = RawResponse { - body: "raw response".into(), - }; + let raw = RawResponse::new("raw response".into()); assert!(matches!( logging .on_event( @@ -1869,9 +1969,10 @@ assert logger.calls[1][1] is response py, &locals, c" -assert logger.names() == ['stream_opened', 'stream_success'], logger.calls +assert logger.names() == ['stamp_stream_head', 'stream_opened', 'stream_success'], logger.calls assert logger.calls[0][1] is head -chunks = logger.calls[1][1] +assert logger.calls[1][1] is head +chunks = logger.calls[2][1] assert len(chunks) == 2 assert chunks[0] is first assert chunks[1] is last diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index f185f21a9c3..1951c9ebfb4 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -5,6 +5,7 @@ use litellm_host::interceptors::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; +use serde_json::Value; use crate::logger::PythonLogger; use crate::python::{Logging, Wrapper}; @@ -15,9 +16,10 @@ pub trait LegacyCallbacks { fn update_from_kwargs( &self, py: Python<'_>, - kwargs: &Py, + kwargs: &Bound<'_, PyDict>, wire: &WireRequest, context: &RequestContext, + request_params: Option<&serde_json::Map>, ) -> PyResult<()>; /// `Logging.pre_call`. @@ -31,6 +33,20 @@ pub trait LegacyCallbacks { url: &str, ) -> PyResult<()>; + /// Records the provider's raw HTTP response on the logger, as Python's Messages handler + /// does with `httpx_response` and `response_headers`. + fn record_response( + &self, + py: Python<'_>, + status: u16, + headers: &[(String, String)], + body: &str, + ) -> PyResult<()>; + + /// Records one provider HTTP exchange's window, as Python's timed `AsyncHTTPHandler.post` + /// does with `llm_api_duration_ms` and the metadata's `llm_api_timing_windows`. + fn record_api_window(&self, py: Python<'_>, start: f64, end: f64) -> PyResult<()>; + /// `Logging.post_call`. fn post_call( &self, @@ -83,12 +99,13 @@ impl LegacyCallbacks for PythonLogger { fn update_from_kwargs( &self, py: Python<'_>, - kwargs: &Py, + kwargs: &Bound<'_, PyDict>, wire: &WireRequest, context: &RequestContext, + request_params: Option<&serde_json::Map>, ) -> PyResult<()> { let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect(); - let redacted_kwargs = redact(py, kwargs.bind(py), &secret_fields)?; + let redacted_kwargs = redact(py, kwargs, &secret_fields)?; let optional_params = redact( py, &to_py(py, &context.optional_params)? @@ -97,23 +114,23 @@ impl LegacyCallbacks for PythonLogger { &secret_fields, )?; let params = PyDict::new(py); - params.set_item( - "litellm_call_id", - kwargs.bind(py).get_item("litellm_call_id")?, - )?; + params.set_item("litellm_call_id", kwargs.get_item("litellm_call_id")?)?; params.set_item("api_base", &wire.url)?; for name in ["logger_fn", "litellm_request_debug"] { - if let Some(value) = kwargs.bind(py).get_item(name)? { + if let Some(value) = kwargs.get_item(name)? { params.set_item(name, value)?; } } for name in custom_pricing_fields(py)? { - if let Some(value) = kwargs.bind(py).get_item(&name)? + if let Some(value) = kwargs.get_item(&name)? && !value.is_none() { params.set_item(name, value)?; } } + for (name, value) in request_params.into_iter().flatten() { + params.set_item(name, to_py(py, value)?)?; + } Logging::Update.call( py, ( @@ -145,6 +162,22 @@ impl LegacyCallbacks for PythonLogger { Ok(()) } + fn record_response( + &self, + py: Python<'_>, + status: u16, + headers: &[(String, String)], + body: &str, + ) -> PyResult<()> { + Logging::RecordResponse.call(py, (self.object(py), status, headers.to_vec(), body))?; + Ok(()) + } + + fn record_api_window(&self, py: Python<'_>, start: f64, end: f64) -> PyResult<()> { + Logging::RecordApiWindow.call(py, (self.object(py), start, end))?; + Ok(()) + } + fn post_call( &self, py: Python<'_>, diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6f1ee9dfd78..6d910d35f53 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -44,6 +44,10 @@ pub(crate) enum Logging { PreCall, #[strum(serialize = "post_call")] PostCall, + #[strum(serialize = "record_response")] + RecordResponse, + #[strum(serialize = "record_api_window")] + RecordApiWindow, #[strum(serialize = "defers_async_logging")] DefersAsync, #[strum(serialize = "defer_success")] @@ -77,6 +81,8 @@ pub(crate) enum DeploymentHooks { pub(crate) enum Streaming { #[strum(serialize = "stream_opened")] Opened, + #[strum(serialize = "stamp_stream_head")] + StampHead, #[strum(serialize = "stream_success")] Success, #[strum(serialize = "stream_failure")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 093ce60e510..8acfa7a898b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -65,6 +65,12 @@ FAKES = { 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( original_response, api_key, additional_args ), + 'record_response': lambda logger, status, headers, body: logger.model_call_details.update( + response_headers=dict(headers) + ), + 'record_api_window': lambda logger, start, end: logger.model_call_details.update( + llm_api_duration_ms=(end - start) * 1000 + ), 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( @@ -85,6 +91,7 @@ FAKES = { ), 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), 'stream_opened': lambda logger, head: logger.record('stream_opened', head), + 'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head), 'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record( 'stream_success', list(chunks) ), diff --git a/litellm-rust/crates/host-http/tests/serve.rs b/litellm-rust/crates/host-http/tests/serve.rs index 2ba39ad4cd2..fd1a8940b33 100644 --- a/litellm-rust/crates/host-http/tests/serve.rs +++ b/litellm-rust/crates/host-http/tests/serve.rs @@ -215,9 +215,7 @@ async fn projection_custom_operations_and_hooks_feed_the_http_response(intercept ) .await?; route_hooks - .after_provider_response(RawResponse { - body: wire.url.clone(), - }) + .after_provider_response(RawResponse::new(wire.url.clone())) .await?; Ok(CallOutput::Complete(Bytes::from(wire.url))) }, @@ -501,9 +499,7 @@ async fn hook_rejection_stops_execution_and_is_reported_once( move |_, _, route_hooks, _observations| async move { if event { route_hooks - .after_provider_response(RawResponse { - body: "response".into(), - }) + .after_provider_response(RawResponse::new("response".into())) .await?; } else { route_hooks @@ -637,9 +633,7 @@ async fn unary_calls_use_into_response_after_hooks_and_before_success(intercepto ) .await?; route_hooks - .after_provider_response(RawResponse { - body: wire.url.clone(), - }) + .after_provider_response(RawResponse::new(wire.url.clone())) .await?; Ok(CallOutput::Complete(json!({"url": wire.url}))) }, @@ -689,7 +683,7 @@ async fn unary_failure_preserves_the_error_without_converting( None, |_, _, route_hooks, _observations| async move { route_hooks - .after_provider_response(RawResponse { body: "raw".into() }) + .after_provider_response(RawResponse::new("raw".into())) .await?; Err(TestError::Provider) }, diff --git a/litellm-rust/crates/host-native/tests/driver.rs b/litellm-rust/crates/host-native/tests/driver.rs index c41c1e5a43b..44c30234328 100644 --- a/litellm-rust/crates/host-native/tests/driver.rs +++ b/litellm-rust/crates/host-native/tests/driver.rs @@ -170,9 +170,7 @@ fn dispatching_call() -> TestMachine { ) .await?; route_hooks - .after_provider_response(RawResponse { - body: wire.url.clone(), - }) + .after_provider_response(RawResponse::new(wire.url.clone())) .await?; Ok(CallOutput::Complete(wire.url)) }, @@ -660,9 +658,7 @@ async fn response_interception_waits_and_can_reject_after_observation(#[case] re "input", Some(sender), |_, _, interceptors, observers| async move { - let raw = RawResponse { - body: "provider response".into(), - }; + let raw = RawResponse::new("provider response".into()); observers.unwrap().emit(CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, )); diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index 003fce74309..15e864c0718 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -467,7 +467,7 @@ pub fn success_machine() -> MachineFactory { .before_provider_request(wire(), context()) .await?; host.interceptors - .after_provider_response(RawResponse { body: "raw".into() }) + .after_provider_response(RawResponse::new("raw".into())) .await?; Ok(format!("{projected}|{signed}|{}", wire.url)) }) diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index 39af7c410e0..9398ea4bad6 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -28,6 +28,28 @@ pub struct RequestContext { #[derive(Clone, Debug, PartialEq, Eq)] pub struct RawResponse { pub body: String, + /// The provider's HTTP status and headers, when the route forwards them to logging. + pub head: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ResponseHead { + pub status: u16, + pub headers: Vec<(String, String)>, +} + +impl RawResponse { + pub fn new(body: String) -> Self { + Self { body, head: None } + } + + #[must_use] + pub fn with_head(self, head: ResponseHead) -> Self { + Self { + head: Some(head), + ..self + } + } } #[derive(Clone, Debug, PartialEq, Eq)] @@ -186,7 +208,7 @@ mod tests { .await?; Interceptors::after_provider_response( &channel.interceptors, - RawResponse { body: "raw".into() }, + RawResponse::new("raw".into()), ) .await?; Ok((sent, ())) diff --git a/litellm-rust/crates/host/tests/observation.rs b/litellm-rust/crates/host/tests/observation.rs index 22872047d81..fba85b63bea 100644 --- a/litellm-rust/crates/host/tests/observation.rs +++ b/litellm-rust/crates/host/tests/observation.rs @@ -74,9 +74,7 @@ fn snapshots_do_not_retain_runtime_objects(capacity: NonZeroUsize, #[case] faile #[rstest] fn provider_snapshots_own_the_response_body(capacity: NonZeroUsize) { - let mut raw = RawResponse { - body: "provider response".into(), - }; + let mut raw = RawResponse::new("provider response".into()); let event: CallEvent<(), (), &RawResponse> = CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw: &raw }); let expected = raw.clone(); diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 29c8b988c0b..7078d4aba2a 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -58,8 +58,20 @@ pub(crate) enum Repr { Network(#[source] reqwest::Error), #[error("upstream request timed out")] Timeout(#[source] reqwest::Error), + /// No bytes arrived within the read timeout, which the caller enforces per read; `body` says + /// the response head had already arrived. + #[error("upstream request timed out")] + ReadTimeout { + #[source] + elapsed: Elapsed, + body: bool, + }, #[error("upstream request failed with status {status}: {body}")] - Status { status: u16, body: UpstreamBody }, + Status { + status: u16, + body: UpstreamBody, + headers: Vec<(String, String)>, + }, #[error("invalid WebSocket URL")] WebSocketUrl(#[source] url::ParseError), @@ -136,13 +148,59 @@ impl Repr { impl Error { pub fn upstream(status: u16, body: String) -> Self { + Self::upstream_with_headers(status, body, Vec::new()) + } + + /// An upstream failure that keeps the provider's response headers (`retry-after`, + /// `request-id`) for the caller. + pub fn upstream_with_headers( + status: u16, + body: String, + headers: Vec<(String, String)>, + ) -> Self { Repr::Status { status, body: UpstreamBody(body), + headers, } .into() } + /// The provider's response headers on an upstream failure; empty for every other error. + pub fn upstream_headers(&self) -> &[(String, String)] { + match &*self.0 { + Repr::Status { headers, .. } => headers, + _ => &[], + } + } + + /// A timeout reached before the connection was established, as opposed to while reading. + pub fn is_connect_timeout(&self) -> bool { + matches!(&*self.0, Repr::Timeout(error) if error.is_connect()) + } + + /// A read the caller timed itself (httpx's per-read timeout) ran out. + pub fn read_timeout(elapsed: Elapsed) -> Self { + Repr::ReadTimeout { + elapsed, + body: false, + } + .into() + } + + /// A read of the response body the caller timed itself ran out, after the head arrived. + pub fn body_read_timeout(elapsed: Elapsed) -> Self { + Repr::ReadTimeout { + elapsed, + body: true, + } + .into() + } + + pub fn is_body_read_timeout(&self) -> bool { + matches!(&*self.0, Repr::ReadTimeout { body: true, .. }) + } + pub fn computed_header(name: impl Into) -> Self { Repr::ComputedHeader(name.into()).into() } @@ -169,6 +227,7 @@ impl Error { match &*self.0 { Repr::Network(_) | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::Handshake(_) | Repr::HandshakeTimeout(_) @@ -227,6 +286,7 @@ impl Error { | Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::Handshake(_) | Repr::HandshakeTimeout(_) @@ -259,6 +319,7 @@ impl Error { | Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -295,6 +356,7 @@ impl Error { | Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -333,6 +395,7 @@ impl Error { | Repr::Signature(_) | Repr::Network(_) | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -356,7 +419,10 @@ impl Error { pub fn is_timeout(&self) -> bool { match &*self.0 { - Repr::Timeout(_) | Repr::HandshakeTimeout(_) | Repr::DownloadTimeout(_) => true, + Repr::Timeout(_) + | Repr::ReadTimeout { .. } + | Repr::HandshakeTimeout(_) + | Repr::DownloadTimeout(_) => true, Repr::Read { .. } | Repr::InvalidPem { .. } | Repr::Client(_) @@ -392,7 +458,7 @@ impl Error { pub fn upstream_response(&self) -> Option<(u16, &str)> { match &*self.0 { - Repr::Status { status, body } => Some((*status, &body.0)), + Repr::Status { status, body, .. } => Some((*status, &body.0)), Repr::Read { .. } | Repr::InvalidPem { .. } | Repr::Client(_) @@ -408,6 +474,7 @@ impl Error { | Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) | Repr::WebSocketSchemeRewrite(_) @@ -431,7 +498,7 @@ impl Error { pub fn status(&self) -> Option { match &*self.0 { Repr::Status { status, .. } | Repr::DownloadStatus(status) => Some(*status), - Repr::Timeout(_) | Repr::DownloadTimeout(_) => Some(408), + Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::DownloadTimeout(_) => Some(408), Repr::Read { .. } | Repr::InvalidPem { .. } | Repr::Client(_) @@ -482,6 +549,7 @@ impl Error { | Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index f948e8fffaf..cdcc3e911b9 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -108,7 +108,7 @@ pub(super) async fn execute( if !status.is_success() { return Err(litellm_http::Error::upstream(status.as_u16(), text).into()); } - let raw = RawResponse { body: text.clone() }; + let raw = RawResponse::new(text.clone()); if let Some(observers) = observers { observers.emit(litellm_host::lifecycle::CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index ac52cb65a0d..0c5be52550d 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -24,6 +24,7 @@ litellm-tracing.workspace = true reqwest.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } +tokio.workspace = true tracing.workspace = true [dev-dependencies] @@ -34,5 +35,4 @@ litellm-host-native.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-inference-testing.workspace = true rstest.workspace = true -tokio.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index a57eca1c6a9..e5ece2837d8 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -35,6 +35,14 @@ impl MessagesProvider { pub(crate) fn config(self) -> &'static dyn BaseMessagesConfig { self.config } + + /// Python's `MESSAGES_SERVICE_TIER_PROVIDERS`: everyone else has `service_tier` dropped. + pub(crate) fn accepts_service_tier(self) -> bool { + matches!( + self.provider, + LlmProviders::Anthropic | LlmProviders::AzureAi | LlmProviders::BedrockMantle + ) + } } /// Python's `get_provider_anthropic_messages_config`: Vertex AI serves only its Claude @@ -119,3 +127,69 @@ mod tests { ); } } + +/// Python's `_merge_messages_header_sources`: later sources win per header name, except +/// `anthropic-beta` (any casing), whose values from every source are unioned, deduplicated and +/// sorted into one lowercase `anthropic-beta` entry. +pub fn merge_messages_header_sources( + sources: impl IntoIterator>, +) -> Map { + let mut merged = Map::new(); + let mut betas = std::collections::BTreeSet::new(); + for (name, value) in sources.into_iter().flatten() { + if !name.eq_ignore_ascii_case("anthropic-beta") { + merged.insert(name, value); + continue; + } + let values: Vec<&str> = match &value { + Value::String(text) => vec![text.as_str()], + Value::Array(entries) => entries.iter().filter_map(Value::as_str).collect(), + _ => Vec::new(), + }; + betas.extend( + values + .iter() + .flat_map(|value| value.split(',')) + .map(str::trim) + .filter(|beta| !beta.is_empty()) + .map(str::to_owned), + ); + } + if !betas.is_empty() { + merged.insert( + "anthropic-beta".to_owned(), + Value::String(betas.into_iter().collect::>().join(",")), + ); + } + merged +} + +#[cfg(test)] +mod merge_tests { + use serde_json::{Map, Value, json}; + + use super::merge_messages_header_sources; + + fn map(value: Value) -> Map { + value.as_object().cloned().unwrap() + } + + #[test] + fn later_sources_win_except_anthropic_beta_which_is_unioned_and_sorted() { + let merged = merge_messages_header_sources([ + map(json!({"x-a": "1", "anthropic-beta": "b-2, a-1"})), + map(json!({"x-a": "2", "Anthropic-Beta": ["c-3", 7, "a-1"]})), + map(json!({"anthropic-beta": 5})), + ]); + assert_eq!( + Value::Object(merged), + json!({"x-a": "2", "anthropic-beta": "a-1,b-2,c-3"}) + ); + } + + #[test] + fn no_beta_values_leave_no_beta_header() { + let merged = merge_messages_header_sources([map(json!({"anthropic-beta": ""}))]); + assert!(merged.is_empty()); + } +} diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 2e647688a1a..e0a7aa5adb8 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -10,9 +10,9 @@ use litellm_core_utils::{ get_provider_specific_headers::get_provider_specific_headers, settings::Lookup, }; use litellm_host::interceptors::{ - Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, WireRequest, + Cost, ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext, + ResponseHead, ResultSource, WireRequest, }; -use litellm_http::request::with_default_headers; use litellm_llms::base_llm::{ auth::{Authenticated, ValidatedEnvironment, resolve_auth}, messages::{ @@ -21,14 +21,16 @@ use litellm_llms::base_llm::{ transformation::BaseMessagesConfig, }, }; -use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse}; +use litellm_llms_types::formats::messages::{ + MessagesOptionalParams, MessagesRequest, MessagesResponse, +}; use litellm_secrets::source::Secrets; use litellm_tracing::ByteChunk; use serde_json::Value; use super::{ Error, MessagesCall, MessagesCallResponse, MessagesRoute, - common_utils::{MessagesProvider, string_headers}, + common_utils::{MessagesProvider, merge_messages_header_sources, string_headers}, constants::MESSAGES_TIMEOUT_SECS, route::Messages, types::invalid_request, @@ -126,14 +128,18 @@ async fn provider_call( MessagesRequest { model, ..body }, shaping.settings.reasoning_auto_summary, )?; + let transform_context = + MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params); + let filtered = requested_optional_params(sanitized, provider, &transform_context)?; let trimmed = - without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?; + without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?; let scoped = get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str()); - let forwarded = string_headers(Some( - extra_headers.into_iter().flatten().chain(scoped).collect(), - ))?; + let forwarded = string_headers(Some(merge_messages_header_sources([ + extra_headers.unwrap_or_default(), + scoped, + ])))?; let validated = config.validate_environment( forwarded, api_key.as_deref(), @@ -144,16 +150,13 @@ async fn provider_call( )?; let body = config.request_body( - config.transform_anthropic_messages_request( - trimmed, - &MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params), - )?, + config.transform_anthropic_messages_request(trimmed, &transform_context)?, &validated.headers, shaping.settings.drop_params, )?; let environment = ValidatedEnvironment { headers: config.request_headers( - with_default_headers(validated.headers, config.default_headers()), + with_python_default_headers(validated.headers, config.default_headers()), &body, ), auth: validated.auth, @@ -184,7 +187,7 @@ async fn provider_call( .before_provider_request( WireRequest { url, - headers: authenticated.headers, + headers: joined_header_values(authenticated.headers), body: config.wire_body(serde_json::to_value(&body).map_err(serialize_failure)?), }, request_context, @@ -213,6 +216,64 @@ async fn provider_call( }) } +/// Python sends its header dict through `httpx.Headers`, which folds names that differ only in +/// case into one lowercase header whose values are joined with `", "`. +fn joined_header_values(headers: Vec<(String, String)>) -> Vec<(String, String)> { + let mut joined: Vec<(String, String)> = Vec::with_capacity(headers.len()); + let mut positions: std::collections::HashMap = + std::collections::HashMap::with_capacity(headers.len()); + for (name, value) in headers { + match positions.entry(name.to_ascii_lowercase()) { + std::collections::hash_map::Entry::Occupied(position) => { + let present = &mut joined[*position.get()].1; + present.push_str(", "); + present.push_str(&value); + } + std::collections::hash_map::Entry::Vacant(position) => { + position.insert(joined.len()); + joined.push((name, value)); + } + } + } + joined +} + +/// Python's `finalize_messages_headers` adds a default only when the exact lowercase name is +/// absent, so a caller's `Anthropic-Version` is sent alongside the default, not instead of it. +fn with_python_default_headers( + headers: Vec<(String, String)>, + defaults: &[(&str, &str)], +) -> Vec<(String, String)> { + let missing: Vec<(String, String)> = defaults + .iter() + .filter(|(name, _)| !headers.iter().any(|(present, _)| present == name)) + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect(); + headers.into_iter().chain(missing).collect() +} + +/// Python's `get_requested_anthropic_messages_optional_param`, which gates speed, service tier +/// and sampling params before `additional_drop_params` and the environment check run. +fn requested_optional_params( + request: MessagesRequest, + provider: MessagesProvider, + context: &MessagesTransformContext, +) -> Result { + let request = litellm_llms::anthropic::messages::transformation::drop_unsupported_params( + request, context, + )?; + if provider.accepts_service_tier() { + return Ok(request); + } + Ok(MessagesRequest { + params: MessagesOptionalParams { + service_tier: None, + ..request.params + }, + ..request + }) +} + fn without_additional_drop_params( request: MessagesRequest, paths: &[String], @@ -245,6 +306,7 @@ async fn call_provider( } = request; let provider_name = provider.as_str(); log_request_body(provider_name, stream, &wire.body); + let timeout = timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)); let response = send( http, Authenticated { @@ -254,10 +316,11 @@ async fn call_provider( &wire.url, &wire.body, timeout, + stream, ) .await?; if !response.status().is_success() { - return Err(provider_error(response).await); + return Err(provider_error(response, timeout).await); } let config = provider.config(); if stream { @@ -265,11 +328,24 @@ async fn call_provider( response, config.stream_decoder(), provider_name, + timeout, )); } + let head = ResponseHead { + status: response.status().as_u16(), + headers: litellm_http::request::header_pairs(response.headers()), + }; + let status = head.status; let text = response.text().await.map_err(litellm_http::Error::from)?; log_response_body(&text); - context.response_received(&text).await?; + context + .raw_response_received(RawResponse::new(text.clone()).with_head(head)) + .await?; + if serde_json::from_str::(&text).is_err() { + // Python raises `AnthropicError(status_code=, message=)` for a + // body it cannot parse, and the public mapping classifies it by that status. + return Err(litellm_http::Error::upstream(status, text).into()); + } decode_response(config, &identity.model, &text) .map(|message| MessagesCallResponse::Complete(Box::new(message))) } @@ -281,30 +357,53 @@ fn serialize_failure(err: serde_json::Error) -> Error { )) } +/// Python's httpx timeout bounds each phase and each read, never the whole exchange. A complete +/// response keeps the client's overall deadline; a stream is bounded until its head arrives and +/// then per read (see [`timed_chunks`]), so a long stream that keeps sending is not cut off. async fn send( http: &litellm_http::Client, authenticated: Authenticated, url: &str, body: &Value, - timeout: Option, + timeout: Duration, + stream: bool, ) -> Result { - let request = outbound_request( - authenticated, - url, - body, - Some(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), - )?; - Ok(litellm_inference::outbound::send(request, http).await?) + let request = outbound_request(authenticated, url, body, (!stream).then_some(timeout))?; + if !stream { + return Ok(litellm_inference::outbound::send(request, http).await?); + } + tokio::time::timeout(timeout, litellm_inference::outbound::send(request, http)) + .await + .map_err(litellm_http::Error::read_timeout)? + .map_err(Error::from) } -async fn provider_error(response: reqwest::Response) -> Error { +/// The stream's body, each read bounded by `timeout` as httpx bounds each read. +fn timed_chunks( + response: reqwest::Response, + timeout: Duration, +) -> impl futures_util::Stream> + Send + 'static { + futures_util::stream::try_unfold(response, move |mut response| async move { + let chunk = tokio::time::timeout(timeout, response.chunk()) + .await + .map_err(litellm_http::Error::body_read_timeout)? + .map_err(litellm_http::Error::from)?; + Ok(chunk.map(|chunk| (chunk, response))) + }) +} + +/// The provider's error, its body read under `timeout`: a stream has no overall deadline once +/// its head arrives, so an error body that stalls would otherwise hold the call forever. +async fn provider_error(response: reqwest::Response, timeout: Duration) -> Error { let status = response.status().as_u16(); - match response.text().await { - Ok(text) => { + let headers = litellm_http::request::header_pairs(response.headers()); + match tokio::time::timeout(timeout, response.text()).await { + Ok(Ok(text)) => { log_error_body(status, &text); - litellm_http::Error::upstream(status, text).into() + litellm_http::Error::upstream_with_headers(status, text, headers).into() } - Err(error) => litellm_http::Error::from(error).into(), + Ok(Err(error)) => litellm_http::Error::from(error).into(), + Err(elapsed) => litellm_http::Error::read_timeout(elapsed).into(), } } @@ -329,22 +428,22 @@ fn streaming_response( response: reqwest::Response, decoder: Option, provider: &'static str, + timeout: Duration, ) -> MessagesCallResponse { let headers = response .headers() .iter() .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) .collect(); + let chunks = timed_chunks(response, timeout) + .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)); let chunks = match decoder { - None => futures_util::stream::try_unfold(response, move |mut response| async move { - let chunk = response.chunk().await.map_err(litellm_http::Error::from)?; - Ok(chunk.map(|chunk| { - log_chunk(provider, "provider_response", &chunk); - (chunk, response) - })) - }) - .boxed(), - Some(decode) => decoded_chunks(response, decode, provider), + None => chunks.map_err(Error::from).boxed(), + Some(decode) => decoded_chunks( + chunks.map_err(std::io::Error::other).boxed(), + decode, + provider, + ), }; MessagesCallResponse::Stream { head: super::route::MessagesStreamHead { headers }, @@ -353,15 +452,10 @@ fn streaming_response( } fn decoded_chunks( - response: reqwest::Response, + bytes: ByteStream, decode: StreamDecoder, provider: &'static str, ) -> BoxStream<'static, Result> { - let bytes: ByteStream = response - .bytes_stream() - .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)) - .map_err(|error| std::io::Error::other(litellm_http::Error::from(error))) - .boxed(); futures_util::stream::try_unfold(decode(bytes), move |mut events| async move { let Some(event) = events.try_next().await? else { return Ok(None); @@ -964,9 +1058,12 @@ mod tests { .send() .await .unwrap(); - let MessagesCallResponse::Stream { mut chunks, .. } = - streaming_response(response, Some(streaming::event_stream), "test") - else { + let MessagesCallResponse::Stream { mut chunks, .. } = streaming_response( + response, + Some(streaming::event_stream), + "test", + Duration::from_secs(5), + ) else { panic!("a streaming response returns chunks"); }; diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index bfefb9c63b0..745442d0314 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -14,6 +14,7 @@ use std::sync::Arc; use crate::common_utils::messages_provider; +pub use common_utils::merge_messages_header_sources; pub use litellm_inference::Error; pub use types::{ MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index e53bdb8d252..216aa96fbe6 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -440,7 +440,12 @@ async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: M #[case::deepseek("deepseek")] #[case::minimax("minimax")] #[tokio::test] -async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) { +/// Python adds a default unless that exact lowercase name is present, then folds the casings +/// into one header, so a capitalised caller header is sent beside the default. +async fn capitalised_caller_protocol_headers_are_joined_with_the_defaults( + call: MessagesCall, + #[case] provider: &str, +) { let upstream = upstream([message_response()]).await; run_message(MessagesCall { @@ -457,10 +462,13 @@ async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[cas .await; let request = only_request(&upstream).await; - assert_eq!(request.header_values("anthropic-version"), ["2024-01-01"]); + assert_eq!( + request.header_values("anthropic-version"), + ["2024-01-01, 2023-06-01"] + ); assert_eq!( request.header_values("content-type"), - ["application/json; charset=utf-8"] + ["application/json; charset=utf-8, application/json"] ); } diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index 7aff1286858..76813fe0a34 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -225,14 +225,25 @@ async fn an_upstream_error_keeps_its_status_and_body(call: MessagesCall, #[case] } #[rstest] -#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))] -#[case::not_a_message(json_response(json!({"unexpected": true})))] #[tokio::test] -async fn an_unreadable_success_body_is_an_invalid_response( - call: MessagesCall, - #[case] response: ResponseTemplate, -) { - let upstream = upstream([response]).await; +async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: MessagesCall) { + let upstream = upstream([ResponseTemplate::new(200).set_body_string("not json")]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .expect_err("an unparseable body fails"); + + assert_eq!(error.upstream_response(), Some((200, "not json"))); +} + +#[rstest] +#[tokio::test] +async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) { + let upstream = upstream([json_response(json!({"unexpected": true}))]).await; let error = run(MessagesCall { api_key: Some("sk".into()), diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index 4a8a061f764..99ba08495aa 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -432,6 +432,49 @@ async fn dropping_the_sdk_stream_closes_the_unfinished_upstream( .unwrap(); } +/// Answers with an error head whose promised body never arrives. +async fn stalling_error_upstream() -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let connection = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all(b"HTTP/1.1 429 Too Many Requests\r\ncontent-type: application/json\r\ncontent-length: 100\r\n\r\n") + .await + .unwrap(); + let _ = socket.read_to_end(&mut Vec::new()).await; + }); + (base, connection) +} + +#[rstest] +#[tokio::test] +async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) { + let (base, connection) = stalling_error_upstream().await; + let host = RecordingStreamHost::new( + MessagesCall { + timeout: Some(Duration::from_millis(300)), + ..streaming(call, base) + }, + usize::MAX, + ); + + let error = tokio::time::timeout(Duration::from_secs(5), stream_through(&host)) + .await + .expect("the stalled error body gives up within the timeout") + .expect_err("a stalled error body fails the call"); + + assert!(error.dispatched(), "{error:?}"); + assert_eq!(error.upstream_response(), None); + assert!(host.seen.into_inner().unwrap().is_empty()); + tokio::time::timeout(Duration::from_secs(5), connection) + .await + .expect("timing out closes the upstream connection") + .unwrap(); +} + #[rstest] #[tokio::test] async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesCall) { diff --git a/litellm-rust/crates/inference-ocr/src/handler.rs b/litellm-rust/crates/inference-ocr/src/handler.rs index c126a192153..be23af60c69 100644 --- a/litellm-rust/crates/inference-ocr/src/handler.rs +++ b/litellm-rust/crates/inference-ocr/src/handler.rs @@ -91,9 +91,7 @@ impl> CallHooks for OcrCallHooks<'_, H> { } fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - let raw = RawResponse { - body: String::from_utf8_lossy(body).into_owned(), - }; + let raw = RawResponse::new(String::from_utf8_lossy(body).into_owned()); if let Some(observers) = self.observers { observers.emit(litellm_host::lifecycle::CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index 19092a63a14..f81a51ad192 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -114,7 +114,7 @@ pub(super) async fn execute( }); } let body = response.text().await.map_err(litellm_http::Error::from)?; - let raw = RawResponse { body: body.clone() }; + let raw = RawResponse::new(body.clone()); if let Some(observers) = observers { observers.emit(litellm_host::lifecycle::CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index 963aa7b90c5..ef2635d24a1 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -42,9 +42,11 @@ impl<'a, I: Interceptors> CallContext<'a, I> { } pub async fn response_received(&self, body: &str) -> Result<(), Error> { - let raw = RawResponse { - body: body.to_owned(), - }; + self.raw_response_received(RawResponse::new(body.to_owned())) + .await + } + + pub async fn raw_response_received(&self, raw: RawResponse) -> Result<(), Error> { if let Some(observers) = &self.observers { observers.emit(CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index a8e6fb4e1b3..841e0c53799 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -146,6 +146,34 @@ impl Error { } } + /// The provider's response headers on an upstream failure; empty for every other error. + pub fn upstream_headers(&self) -> &[(String, String)] { + match self { + Self::Llm(error) => error.upstream_headers(), + Self::Http(error) => error.upstream_headers(), + _ => &[], + } + } + + /// A read of a streamed body timed out after the response head arrived. + pub fn body_read_timeout(&self) -> bool { + match self { + Self::Llm(error) => error.is_body_read_timeout(), + Self::Http(error) => error.is_body_read_timeout(), + _ => false, + } + } + + /// `Some(connect)` when the call timed out, where `connect` says the connection was never + /// established; `None` for every other error. + pub fn timeout(&self) -> Option { + match self { + Self::Llm(error) => error.is_timeout(), + Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()), + _ => None, + } + } + pub fn upstream_response(&self) -> Option<(u16, &str)> { match self { Self::Llm(error) => error.upstream_response(), diff --git a/litellm-rust/crates/llms-types/src/formats/messages/request.rs b/litellm-rust/crates/llms-types/src/formats/messages/request.rs index 9dbfa068025..d30d21e74e7 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/request.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/request.rs @@ -264,7 +264,12 @@ pub enum ThinkingDisplay { #[macro_rules_attribute::apply(crate::wire_type)] #[derive(Default)] pub struct EnabledThinking { - #[serde(default, skip_serializing_if = "Option::is_none")] + /// An explicit `null` stays `Unrecognized(Null)`: Python rejects it, but defaults a missing key to 0. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "crate::recognized::present" + )] pub budget_tokens: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub display: Option>, diff --git a/litellm-rust/crates/llms-types/src/formats/messages/response.rs b/litellm-rust/crates/llms-types/src/formats/messages/response.rs index 9b0c5d31b3a..276a30a6cb4 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/response.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/response.rs @@ -10,14 +10,29 @@ pub struct MessagesResponse { pub content: Vec, pub stop_reason: Option, pub stop_sequence: Option, - #[serde(skip_serializing_if = "Option::is_none")] + /// `Some(Value::Null)` when the provider sent an explicit `null`, which Python keeps. + #[serde( + default, + deserialize_with = "present_value", + skip_serializing_if = "Option::is_none" + )] pub usage: Option, - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "present_value", + skip_serializing_if = "Option::is_none" + )] pub container: Option, #[serde(flatten)] pub extra: Map, } +fn present_value<'de, D: serde::Deserializer<'de>>( + deserializer: D, +) -> Result, D::Error> { + >::deserialize(deserializer).map(Some) +} + #[cfg(test)] mod tests { use super::*; @@ -73,4 +88,21 @@ mod tests { assert_eq!(body.get("usage").cloned(), usage); assert_eq!(body.get("container").cloned(), container); } + + #[test] + fn explicit_null_usage_and_container_survive_a_round_trip() { + let body = json!({"id": "msg_1", "type": "message", "role": "assistant", "model": "m", + "content": [], "stop_reason": null, "stop_sequence": null, + "usage": null, "container": null}); + let parsed: MessagesResponse = serde_json::from_value(body.clone()).expect("parses"); + assert_eq!(serde_json::to_value(parsed).expect("serializable"), body); + } + + #[test] + fn absent_usage_and_container_stay_absent() { + let body = json!({"id": "msg_1", "type": "message", "role": "assistant", "model": "m", + "content": [], "stop_reason": null, "stop_sequence": null}); + let parsed: MessagesResponse = serde_json::from_value(body.clone()).expect("parses"); + assert_eq!(serde_json::to_value(parsed).expect("serializable"), body); + } } diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index 01ed023999d..46c14a21c59 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -124,12 +124,86 @@ pub enum AnthropicBeta { ToolSearchTool20251019, #[strum(serialize = "inline-tools-2026-09-15")] InlineTools20260915, + #[strum(serialize = "afk-mode-2026-01-31")] + AfkMode20260131, + #[strum(serialize = "agent-memory-2026-07-22")] + AgentMemory20260722, + #[strum(serialize = "analytics-usage-cost-2026-08-24")] + AnalyticsUsageCost20260824, + #[strum(serialize = "cache-diagnosis-2026-04-07")] + CacheDiagnosis20260407, + #[strum(serialize = "ce-plugins-2026-09-01")] + CePlugins20260901, + #[strum(serialize = "ce-user-management-2026-07-13")] + CeUserManagement20260713, + #[strum(serialize = "code-execution-2025-05-22")] + CodeExecution20250522, + #[strum(serialize = "computer-use-2024-10-22")] + ComputerUse20241022, + #[strum(serialize = "dev-full-thinking-2025-05-14")] + DevFullThinking20250514, + #[strum(serialize = "dreaming-2026-04-21")] + Dreaming20260421, + #[strum(serialize = "experimental-claude-tag-2026-10-21")] + ExperimentalClaudeTag20261021, + #[strum(serialize = "extended-cache-ttl-2025-04-11")] + ExtendedCacheTtl20250411, + #[strum(serialize = "fallback-credit-2026-06-01")] + FallbackCredit20260601, + #[strum(serialize = "fallback-credit-2026-07-01")] + FallbackCredit20260701, + #[strum(serialize = "managed-agents-2026-04-01")] + ManagedAgents20260401, + #[strum(serialize = "mcp-client-2026-09-15")] + McpClient20260915, + #[strum(serialize = "mcp-tunnels-2026-06-22")] + McpTunnels20260622, + #[strum(serialize = "message-batches-2024-09-24")] + MessageBatches20240924, + #[strum(serialize = "message-threads-2026-08-12")] + MessageThreads20260812, + #[strum(serialize = "mid-conversation-system-clear-at-2026-08-21")] + MidConversationSystemClearAt20260821, + #[strum(serialize = "model-context-window-exceeded-2025-08-26")] + ModelContextWindowExceeded20250826, + #[strum(serialize = "output-300k-2026-03-24")] + Output300k20260324, + #[strum(serialize = "pdfs-2024-09-25")] + Pdfs20240925, + #[strum(serialize = "prompt-caching-2024-07-31")] + PromptCaching20240731, + #[strum(serialize = "redact-thinking-2026-02-12")] + RedactThinking20260212, + #[strum(serialize = "server-side-fallback-2026-06-01")] + ServerSideFallback20260601, + #[strum(serialize = "server-side-fallback-2026-07-01")] + ServerSideFallback20260701, + #[strum(serialize = "spend-limit-reads-2026-09-26")] + SpendLimitReads20260926, + #[strum(serialize = "structured-outputs-2025-12-15")] + StructuredOutputs20251215, + #[strum(serialize = "task-budgets-2026-03-13")] + TaskBudgets20260313, + #[strum(serialize = "telemetry-destinations-2026-08-11")] + TelemetryDestinations20260811, + #[strum(serialize = "thinking-token-count-2026-05-13")] + ThinkingTokenCount20260513, + #[strum(serialize = "token-counting-2024-11-01")] + TokenCounting20241101, + #[strum(serialize = "user-profiles-2026-03-24")] + UserProfiles20260324, + #[strum(serialize = "user-profiles-2026-08-18")] + UserProfiles20260818, + #[strum(serialize = "user-profiles-2026-09-04")] + UserProfiles20260904, + #[strum(serialize = "webhooks-2026-10-07")] + Webhooks20261007, #[strum(default, transparent)] Other(String), } impl AnthropicBeta { - pub const KNOWN: [Self; 40] = [ + pub const KNOWN: [Self; 77] = [ Self::Oauth20250420, Self::WebFetch20250910, Self::WebSearch20250305, @@ -170,6 +244,43 @@ impl AnthropicBeta { Self::ToolExamples20251029, Self::ToolSearchTool20251019, Self::InlineTools20260915, + Self::AfkMode20260131, + Self::AgentMemory20260722, + Self::AnalyticsUsageCost20260824, + Self::CacheDiagnosis20260407, + Self::CePlugins20260901, + Self::CeUserManagement20260713, + Self::CodeExecution20250522, + Self::ComputerUse20241022, + Self::DevFullThinking20250514, + Self::Dreaming20260421, + Self::ExperimentalClaudeTag20261021, + Self::ExtendedCacheTtl20250411, + Self::FallbackCredit20260601, + Self::FallbackCredit20260701, + Self::ManagedAgents20260401, + Self::McpClient20260915, + Self::McpTunnels20260622, + Self::MessageBatches20240924, + Self::MessageThreads20260812, + Self::MidConversationSystemClearAt20260821, + Self::ModelContextWindowExceeded20250826, + Self::Output300k20260324, + Self::Pdfs20240925, + Self::PromptCaching20240731, + Self::RedactThinking20260212, + Self::ServerSideFallback20260601, + Self::ServerSideFallback20260701, + Self::SpendLimitReads20260926, + Self::StructuredOutputs20251215, + Self::TaskBudgets20260313, + Self::TelemetryDestinations20260811, + Self::ThinkingTokenCount20260513, + Self::TokenCounting20241101, + Self::UserProfiles20260324, + Self::UserProfiles20260818, + Self::UserProfiles20260904, + Self::Webhooks20261007, ]; pub fn as_str(&self) -> &str { @@ -195,10 +306,9 @@ impl AnthropicBeta { fn rejected_by(&self) -> &'static [BetaProvider] { match self { - Self::AdvancedToolUse20251120 => { + Self::AdvancedToolUse20251120 | Self::ToolSearchTool20251019 => { &[BetaProvider::BedrockConverse, BetaProvider::Databricks] } - Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse], Self::AdvisorTool20260301 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, @@ -207,98 +317,99 @@ impl AnthropicBeta { BetaProvider::VertexAi, BetaProvider::Databricks, ], - Self::Compact20260904 => &[ + Self::AfkMode20260131 + | Self::CePlugins20260901 + | Self::CeUserManagement20260713 + | Self::Compact20260904 + | Self::Dreaming20260421 + | Self::ExperimentalClaudeTag20261021 + | Self::ExtendedCacheTtl20250411 + | Self::FallbackCredit20260601 + | Self::FallbackCredit20260701 + | Self::InlineTools20260915 + | Self::ManagedAgents20260401 + | Self::McpTunnels20260622 + | Self::MidConversationSystemClearAt20260821 + | Self::ModelContextWindowExceeded20250826 + | Self::Oauth20250420 + | Self::Pdfs20240925 + | Self::PromptCaching20240731 + | Self::PromptCachingScope20260105 + | Self::RedactThinking20260212 + | Self::ServerSideFallback20260601 + | Self::SpendLimitReads20260926 + | Self::StructuredOutputs20251215 + | Self::TaskBudgets20260313 + | Self::TelemetryDestinations20260811 + | Self::ThinkingTokenCount20260513 + | Self::TokenCounting20241101 + | Self::UserProfiles20260324 + | Self::UserProfiles20260818 + | Self::UserProfiles20260904 => &[ BetaProvider::BedrockConverse, BetaProvider::Bedrock, BetaProvider::BedrockMantle, BetaProvider::Databricks, ], - Self::Bash20241022 + Self::AgentMemory20260722 + | Self::Bash20241022 | Self::Bash20250124 + | Self::DevFullThinking20250514 | Self::McpServers20251204 | Self::StructuredOutput20240301 | Self::TextEditor20241022 | Self::TextEditor20250124 => BetaProvider::VARIANTS, - Self::ClaudeCode20250219 | Self::ToolExamples20251029 => &[ - BetaProvider::Anthropic, - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::CodeExecution20250825 + Self::AnalyticsUsageCost20260824 + | Self::CacheDiagnosis20260407 + | Self::CodeExecution20250522 + | Self::CodeExecution20250825 + | Self::ComputerUse20241022 + | Self::FastMode20260201 | Self::FilesApi20250414 | Self::McpClient20250404 | Self::McpClient20251120 - | Self::PromptCachingScope20260105 + | Self::McpClient20260915 + | Self::MessageBatches20240924 + | Self::MessageThreads20260812 + | Self::Output300k20260324 + | Self::ServerSideFallback20260701 | Self::Skills20251002 - | Self::WebFetch20250910 => &[ + | Self::WebFetch20250910 + | Self::Webhooks20261007 => &[ BetaProvider::BedrockConverse, BetaProvider::Bedrock, BetaProvider::BedrockMantle, BetaProvider::VertexAi, BetaProvider::Databricks, ], - Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse], + Self::ClaudeCode20250219 + | Self::PerTurnControl20260701 + | Self::ToolExamples20251029 => &[ + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::Databricks, + ], + Self::Compact20260112 + | Self::ContextManagement20250627 + | Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 + | Self::Effort20251124 + | Self::FineGrainedToolStreaming20250514 | Self::ThinkingBindingControls20260801 | Self::ThinkingDisplayUpdates20260818 => &[], - Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], - Self::Effort20251124 => &[BetaProvider::VertexAi], - Self::FastMode20260201 | Self::Oauth20250420 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::BedrockMantle, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::FineGrainedToolStreaming20250514 => { - &[BetaProvider::AzureAi, BetaProvider::VertexAi] - } - Self::InterleavedThinking20250514 | Self::WebSearch20250305 => { - &[BetaProvider::BedrockConverse, BetaProvider::Bedrock] - } - Self::MidConversationOutputConfig20260701 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::BedrockMantle, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::MidConversationToolChanges20260701 => &[ + Self::InterleavedThinking20250514 + | Self::Output128k20250219 + | Self::TokenEfficientTools20250219 + | Self::WebSearch20250305 => &[BetaProvider::BedrockConverse, BetaProvider::Bedrock], + Self::MidConversationOutputConfig20260701 + | Self::MidConversationToolChanges20260701 => &[ BetaProvider::BedrockConverse, BetaProvider::BedrockMantle, BetaProvider::Databricks, ], - Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::VertexAi, - ], - Self::PerTurnControl20260701 => &[ - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::Databricks, - ], - Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi], - Self::ToolSearchTool20251019 => &[ - BetaProvider::Anthropic, - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Databricks, - ], - Self::InlineTools20260915 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::BedrockMantle, - BetaProvider::Databricks, - ], + Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock], Self::Other(_) => BetaProvider::VARIANTS, } } @@ -451,7 +562,44 @@ mod tests { AnthropicBeta::TokenEfficientTools20250219, AnthropicBeta::ToolExamples20251029, AnthropicBeta::ToolSearchTool20251019, - AnthropicBeta::InlineTools20260915 + AnthropicBeta::InlineTools20260915, + AnthropicBeta::AfkMode20260131, + AnthropicBeta::AgentMemory20260722, + AnthropicBeta::AnalyticsUsageCost20260824, + AnthropicBeta::CacheDiagnosis20260407, + AnthropicBeta::CePlugins20260901, + AnthropicBeta::CeUserManagement20260713, + AnthropicBeta::CodeExecution20250522, + AnthropicBeta::ComputerUse20241022, + AnthropicBeta::DevFullThinking20250514, + AnthropicBeta::Dreaming20260421, + AnthropicBeta::ExperimentalClaudeTag20261021, + AnthropicBeta::ExtendedCacheTtl20250411, + AnthropicBeta::FallbackCredit20260601, + AnthropicBeta::FallbackCredit20260701, + AnthropicBeta::ManagedAgents20260401, + AnthropicBeta::McpClient20260915, + AnthropicBeta::McpTunnels20260622, + AnthropicBeta::MessageBatches20240924, + AnthropicBeta::MessageThreads20260812, + AnthropicBeta::MidConversationSystemClearAt20260821, + AnthropicBeta::ModelContextWindowExceeded20250826, + AnthropicBeta::Output300k20260324, + AnthropicBeta::Pdfs20240925, + AnthropicBeta::PromptCaching20240731, + AnthropicBeta::RedactThinking20260212, + AnthropicBeta::ServerSideFallback20260601, + AnthropicBeta::ServerSideFallback20260701, + AnthropicBeta::SpendLimitReads20260926, + AnthropicBeta::StructuredOutputs20251215, + AnthropicBeta::TaskBudgets20260313, + AnthropicBeta::TelemetryDestinations20260811, + AnthropicBeta::ThinkingTokenCount20260513, + AnthropicBeta::TokenCounting20241101, + AnthropicBeta::UserProfiles20260324, + AnthropicBeta::UserProfiles20260818, + AnthropicBeta::UserProfiles20260904, + AnthropicBeta::Webhooks20261007 )] beta: AnthropicBeta, ) { @@ -551,7 +699,7 @@ mod tests { BetaProvider::Anthropic, Some(AnthropicBeta::Oauth20250420) )] - #[case::rejected("effort-2025-11-24", BetaProvider::VertexAi, None)] + #[case::rejected("skills-2025-10-02", BetaProvider::VertexAi, None)] #[case::unknown("example-beta-2099-01-01", BetaProvider::Anthropic, None)] fn on_resolves_known_spellings_held_as_other( #[case] raw: &str, @@ -580,12 +728,11 @@ mod tests { #[test] fn unknown_values_are_kept_verbatim() { - let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap(); - assert_eq!( - parsed, - AnthropicBeta::Other("claude-code-20250219".to_string()) + let parsed: AnthropicBeta = "example-beta-2099-01-01".parse().unwrap(); + assert!( + matches!(&parsed, AnthropicBeta::Other(value) if value == "example-beta-2099-01-01") ); - assert_eq!(parsed.to_string(), "claude-code-20250219"); + assert_eq!(parsed.to_string(), "example-beta-2099-01-01"); } #[test] diff --git a/litellm-rust/crates/llms-types/src/recognized.rs b/litellm-rust/crates/llms-types/src/recognized.rs index 58e8dd7a2e7..0545389afbc 100644 --- a/litellm-rust/crates/llms-types/src/recognized.rs +++ b/litellm-rust/crates/llms-types/src/recognized.rs @@ -7,6 +7,16 @@ pub enum Recognized { Unrecognized(Value), } +/// `deserialize_with` for an optional `Recognized` field whose explicit `null` must stay +/// distinguishable from an absent key: `null` becomes `Some(Unrecognized(Null))`. +pub fn present<'de, D, T>(deserializer: D) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, + Recognized: serde::Deserialize<'de>, +{ + as serde::Deserialize<'de>>::deserialize(deserializer).map(Some) +} + impl Recognized { pub fn known(&self) -> Option<&T> { match self { diff --git a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs index a0843621043..7842b478dce 100644 --- a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs +++ b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs @@ -96,7 +96,10 @@ mod tests { } #[rstest] - #[case::azure_ai(BetaProvider::AzureAi, "fast-mode-2026-02-01,example-beta-2099-01-01")] + #[case::azure_ai( + BetaProvider::AzureAi, + "advisor-tool-2026-03-01,example-beta-2099-01-01" + )] #[case::anthropic(BetaProvider::Anthropic, "bash_20241022")] #[case::vertex_ai(BetaProvider::VertexAi, " , ")] fn filter_removes_header_when_nothing_survives( diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index aec5d9ebc27..13653f438f7 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -154,11 +154,42 @@ pub fn complete_anthropic_url( ) -> String { let api_base = resolve_anthropic_api_base(api_base, env_lookup); - let api_base = api_base.trim_end_matches('/'); - if api_base.ends_with(MESSAGES_PATH) { + ensure_api_route(&api_base, MESSAGES_PATH) +} + +/// Python's `urlsplit(url).path`, for the URLs an `api_base` holds. +fn url_path(url: &str) -> &str { + let before_fragment = url.split('#').next().unwrap_or_default(); + let before_query = before_fragment.split('?').next().unwrap_or_default(); + match before_query.split_once("://") { + Some((_, rest)) => rest.find('/').map_or("", |slash| &rest[slash..]), + None => before_query, + } +} + +/// Python's `_strip_trailing_slashes_keeping_the_scheme`. +fn strip_trailing_slashes(api_base: &str) -> String { + let stripped = api_base.trim_end_matches('/'); + if stripped.ends_with(':') && !stripped.contains('/') && !stripped.contains('?') { return api_base.to_string(); } - format!("{api_base}{MESSAGES_PATH}") + match stripped.split_once('?') { + Some((before, query)) if url_path(stripped).ends_with('/') => { + format!("{}?{query}", before.trim_end_matches('/')) + } + _ => stripped.to_string(), + } +} + +/// Python's `ensure_api_route`: an `api_base` whose path already ends in the route is kept as +/// it is, query string included; anything else gets the route appended to the whole string. +fn ensure_api_route(api_base: &str, route: &str) -> String { + let base = strip_trailing_slashes(api_base); + let route = route.trim_start_matches('/'); + if url_path(&base).ends_with(&format!("/{route}")) { + return base; + } + format!("{}/{route}", strip_trailing_slashes(&base)) } pub fn existing_betas(headers: &[(String, String)]) -> BetaSet { @@ -1922,4 +1953,24 @@ mod tests { expected ); } + + #[rstest] + #[case::bare_host("https://api.example.com", "https://api.example.com/v1/messages")] + #[case::trailing_slash("https://api.example.com/", "https://api.example.com/v1/messages")] + #[case::route_present( + "https://api.example.com/v1/messages/", + "https://api.example.com/v1/messages" + )] + #[case::route_and_query( + "https://api.example.com/v1/messages?beta=true", + "https://api.example.com/v1/messages?beta=true" + )] + #[case::query_without_route( + "https://api.example.com/?x=1", + "https://api.example.com?x=1/v1/messages" + )] + #[case::scheme_only("https://", "https:///v1/messages")] + fn the_messages_url_follows_python_ensure_api_route(#[case] api_base: &str, #[case] url: &str) { + assert_eq!(complete_anthropic_url(Some(api_base), &|_| None), url); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs index df1ae13fd28..ae3488a21f9 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs @@ -2,8 +2,8 @@ use litellm_llms_types::{ formats::{ chat_completions::ReasoningEffort, messages::{ - EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig, ThinkingConfig, - ThinkingDisplay, + AdaptiveThinking, EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig, + ThinkingConfig, ThinkingDisplay, }, }, recognized::Recognized, @@ -36,16 +36,17 @@ fn budget_for_effort(budgets: &ThinkingBudgets, effort: ReasoningEffort) -> Opti fn effort_for_budget( budgets: &ThinkingBudgets, - budget_tokens: u64, + budget_tokens: i64, capabilities: &MessagesModelCapabilities, ) -> EffortLevel { - if budget_tokens >= budgets.xhigh && capabilities.effort_tiers.xhigh { + let reaches = |tier: u64| i128::from(budget_tokens) >= i128::from(tier); + if reaches(budgets.xhigh) && capabilities.effort_tiers.xhigh { return EffortLevel::Xhigh; } - if budget_tokens >= budgets.high { + if reaches(budgets.high) { return EffortLevel::High; } - if budget_tokens >= budgets.medium { + if reaches(budgets.medium) { return EffortLevel::Medium; } EffortLevel::Low @@ -231,32 +232,129 @@ fn drop_disabled_thinking(request: MessagesRequest, context: &ThinkingContext) - } } +/// `_numeric_budget_tokens`: Python reads any int, float (truncated) or int-parseable string +/// as a budget, and nothing else. +fn numeric_budget_tokens(budget: &Recognized) -> Option { + let value = match budget { + Recognized::Known(budget) => return Some(i64::try_from(*budget).unwrap_or(i64::MAX)), + Recognized::Unrecognized(value) => value, + }; + match value { + Value::Number(number) => number + .as_i64() + .or_else(|| number.as_f64().map(|float| float.trunc() as i64)), + Value::String(text) => python_int(text), + Value::Null | Value::Bool(_) | Value::Array(_) | Value::Object(_) => None, + } +} + +/// Python's `int(str)`: surrounding whitespace, a sign, and single underscores between digits. +fn python_int(text: &str) -> Option { + let trimmed = text.trim(); + let digits = trimmed.trim_start_matches(['+', '-']); + let well_formed = trimmed.len() - digits.len() <= 1 + && !digits.is_empty() + && !digits.starts_with('_') + && !digits.ends_with('_') + && !digits.contains("__") + && digits.chars().all(|ch| ch.is_ascii_digit() || ch == '_'); + if !well_formed { + return None; + } + let cleaned: String = trimmed.chars().filter(|ch| *ch != '_').collect(); + cleaned.parse::().ok().or_else(|| { + Some(if cleaned.starts_with('-') { + i64::MIN + } else { + i64::MAX + }) + }) +} + +fn python_type_name(value: &Value) -> &'static str { + match value { + Value::Null => "NoneType", + Value::Bool(_) => "bool", + Value::Number(number) if number.is_f64() => "float", + Value::Number(_) => "int", + Value::String(_) => "str", + Value::Array(_) => "list", + Value::Object(_) => "dict", + } +} + +fn kept_display( + display: Option<&Recognized>, + context: &ThinkingContext, +) -> Option> { + match display? { + Recognized::Known(ThinkingDisplay::Summarized | ThinkingDisplay::Omitted) => { + display.cloned() + } + Recognized::Known(ThinkingDisplay::Updates) if context.keeps_display_updates => { + display.cloned() + } + Recognized::Known(ThinkingDisplay::Updates) | Recognized::Unrecognized(_) => None, + } +} + +/// Keys Python's `carried_thinking_keys` keeps on the rebuilt adaptive config: the adaptive +/// keys always, and any other key unless `drop_params` is set. +fn carried_thinking_keys( + extra: &serde_json::Map, + context: &ThinkingContext, +) -> serde_json::Map { + extra + .iter() + .filter(|(key, _)| key.as_str() == "block_binding" || !context.drop_params) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} + +/// `AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently`. fn translate_legacy_thinking_for_adaptive_model( request: MessagesRequest, context: &ThinkingContext, -) -> MessagesRequest { +) -> Result { let capabilities = &context.capabilities; if !capabilities.supports_adaptive_thinking || capabilities.supports_legacy_thinking { - return request; + return Ok(request); } let Some(ThinkingConfig::Enabled(enabled)) = known_thinking(&request) else { - return request; + return Ok(request); }; - let budget = enabled - .budget_tokens - .as_ref() - .and_then(Recognized::known) - .copied() - .unwrap_or(0); - let level = effort_for_budget(&context.budgets, budget, capabilities); - MessagesRequest { + let budget = match &enabled.budget_tokens { + None => Some(0), + Some(budget) => numeric_budget_tokens(budget), + }; + if budget.is_none() && !context.drop_params { + let actual = match &enabled.budget_tokens { + Some(Recognized::Unrecognized(value)) => python_type_name(value), + Some(Recognized::Known(_)) | None => "int", + }; + return Err(Error::InvalidRequest( + crate::ErrorDetail::NonNumericThinkingBudget { actual }, + )); + } + let thinking = ThinkingConfig::Adaptive(AdaptiveThinking { + display: kept_display(enabled.display.as_ref(), context), + extra: carried_thinking_keys(&enabled.extra, context), + }); + let output_config = match budget { + Some(budget) => with_default_effort( + request.params.output_config, + effort_for_budget(&context.budgets, budget, capabilities), + ), + None => request.params.output_config, + }; + Ok(MessagesRequest { params: MessagesOptionalParams { - thinking: Some(Recognized::Known(ThinkingConfig::adaptive(None))), - output_config: with_default_effort(request.params.output_config, level), + thinking: Some(Recognized::Known(thinking)), + output_config, ..request.params }, ..request - } + }) } fn translate_adaptive_effort_for_non_adaptive_model( @@ -337,7 +435,7 @@ pub fn translate_thinking( ) -> Result { let request = translate_reasoning_effort(request, context)?; let request = drop_disabled_thinking(request, context); - let request = translate_legacy_thinking_for_adaptive_model(request, context); + let request = translate_legacy_thinking_for_adaptive_model(request, context)?; let request = translate_adaptive_effort_for_non_adaptive_model(request, context)?; Ok(drop_incompatible_temperature_for_thinking(request, context)) } @@ -363,6 +461,7 @@ mod tests { ThinkingContext { capabilities, budgets: ThinkingBudgets::default(), + ..ThinkingContext::default() } } @@ -1239,10 +1338,102 @@ mod tests { let context = ThinkingContext { capabilities, budgets: overridden_budgets(overrides), + ..ThinkingContext::default() }; assert_eq!( translate_thinking(request(input), &context), Ok(request(expected)) ); } + + use serde_json::json; + + fn legacy_context(drop_params: bool, keeps_display_updates: bool) -> ThinkingContext { + ThinkingContext { + drop_params, + keeps_display_updates, + ..context(opus_4_7()) + } + } + + #[rstest] + #[case::float_budget(json!(8192.0), "xhigh")] + #[case::float_below_a_tier(json!(1023.9), "low")] + #[case::string_budget(json!(" 16_000 "), "xhigh")] + #[case::negative_budget(json!(-5), "low")] + fn legacy_budgets_read_like_python_numbers(#[case] budget: Value, #[case] effort: &str) { + let translated = translate_thinking( + request( + json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": budget}}), + ), + &legacy_context(false, false), + ); + assert_eq!( + translated, + Ok(request(json!({ + "max_tokens": 64000, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": effort}, + }))) + ); + } + + #[rstest] + #[case::text(json!("lots"), "str")] + #[case::list(json!([1]), "list")] + #[case::null(Value::Null, "NoneType")] + #[case::boolean(json!(true), "bool")] + fn a_non_numeric_legacy_budget_is_rejected_unless_dropped( + #[case] budget: Value, + #[case] actual: &'static str, + ) { + let fields = + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": budget}}); + assert_eq!( + translate_thinking(request(fields.clone()), &legacy_context(false, false)), + Err(Error::InvalidRequest( + crate::ErrorDetail::NonNumericThinkingBudget { actual } + )) + ); + assert_eq!( + translate_thinking(request(fields), &legacy_context(true, false)), + Ok(request( + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + )) + ); + } + + #[rstest] + #[case::summarized("summarized", false, Some("summarized"))] + #[case::updates_kept("updates", true, Some("updates"))] + #[case::updates_dropped("updates", false, None)] + #[case::unknown("verbose", true, None)] + fn legacy_display_survives_only_where_python_keeps_it( + #[case] display: &str, + #[case] keeps_display_updates: bool, + #[case] kept: Option<&str>, + ) { + let translated = translate_thinking( + request(json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 2048, "display": display}})), + &legacy_context(false, keeps_display_updates), + ) + .unwrap(); + let thinking = serde_json::to_value(translated.params.thinking).unwrap(); + assert_eq!(thinking.get("display").and_then(Value::as_str), kept); + } + + #[rstest] + #[case::kept(false, json!({"type": "adaptive", "block_binding": "strict", "note": "x"}))] + #[case::dropped(true, json!({"type": "adaptive", "block_binding": "strict"}))] + fn legacy_extra_keys_follow_drop_params(#[case] drop_params: bool, #[case] expected: Value) { + let translated = translate_thinking( + request(json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 2048, "block_binding": "strict", "note": "x"}})), + &legacy_context(drop_params, false), + ) + .unwrap(); + assert_eq!( + serde_json::to_value(translated.params.thinking).unwrap(), + expected + ); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index fb6d7259bf9..18789e5b2be 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -56,7 +56,7 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - transform_messages_request(request, context) + transform_messages_request(request, &context.keeping_display_updates()) } fn secret_names(&self) -> &'static [&'static str] { @@ -206,6 +206,8 @@ fn feature_betas(request: &MessagesRequest) -> BetaSet { .then_some(AnthropicBeta::PerTurnControl20260701), uses_thinking_display_updates(params) .then_some(AnthropicBeta::ThinkingDisplayUpdates20260818), + uses_thinking_block_binding(params) + .then_some(AnthropicBeta::ThinkingBindingControls20260801), messages_carry_tool_changes(&request.messages) .then_some(AnthropicBeta::MidConversationToolChanges20260701), has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301), @@ -267,6 +269,25 @@ fn uses_thinking_display_updates(params: &MessagesOptionalParams) -> bool { display.and_then(Recognized::known) == Some(&ThinkingDisplay::Updates) } +/// `AnthropicModelInfo.thinking_betas`: any thinking object carrying `block_binding`. +fn uses_thinking_block_binding(params: &MessagesOptionalParams) -> bool { + match params.thinking.as_ref() { + Some(Recognized::Known(ThinkingConfig::Enabled(thinking))) => { + thinking.extra.contains_key("block_binding") + } + Some(Recognized::Known(ThinkingConfig::Adaptive(thinking))) => { + thinking.extra.contains_key("block_binding") + } + Some(Recognized::Known(ThinkingConfig::Disabled(thinking))) => { + thinking.extra.contains_key("block_binding") + } + Some(Recognized::Unrecognized(Value::Object(thinking))) => { + thinking.contains_key("block_binding") + } + Some(Recognized::Unrecognized(_)) | None => false, + } +} + fn is_tool_reference_change(block: &ContentBlock) -> bool { matches!( &block.block_type, @@ -297,7 +318,8 @@ fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error }) } -fn drop_unsupported_params( +/// The speed and sampling gates of Python's `get_requested_anthropic_messages_optional_param`. +pub fn drop_unsupported_params( request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { @@ -838,6 +860,8 @@ mod tests { low: 2000, ..ThinkingBudgets::default() }, + drop_params: true, + keeps_display_updates: false, }, drop_params: true, } diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index a2ba2cee0d3..d13cfa81a94 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -68,7 +68,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { .capabilities .supports_mid_conversation_system, )), - context, + &context.keeping_display_updates(), ) } diff --git a/litellm-rust/crates/llms/src/base_llm/messages/context.rs b/litellm-rust/crates/llms/src/base_llm/messages/context.rs index e59a63813b0..b514e975aa4 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/context.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/context.rs @@ -110,6 +110,11 @@ impl ThinkingBudgets { pub struct ThinkingContext { pub capabilities: MessagesModelCapabilities, pub budgets: ThinkingBudgets, + /// `drop_params` as Python's thinking translation reads it. + pub drop_params: bool, + /// The provider accepts `display: "updates"` on thinking translated from the legacy shape + /// (Python's `LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS`). + pub keeps_display_updates: bool, } #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub struct MessagesTransformContext { @@ -130,8 +135,22 @@ impl MessagesTransformContext { thinking: ThinkingContext { capabilities, budgets: ThinkingBudgets::from_lookup(env), + drop_params, + keeps_display_updates: false, }, drop_params, } } + + /// The same context for a provider in Python's `LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS`. + #[must_use] + pub fn keeping_display_updates(&self) -> Self { + Self { + thinking: ThinkingContext { + keeps_display_updates: true, + ..self.thinking + }, + ..*self + } + } } diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index b3d1421083a..2844fbf3e01 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -92,6 +92,24 @@ impl Error { } } + pub fn upstream_headers(&self) -> &[(String, String)] { + match self { + Self::Http(error) => error.upstream_headers(), + _ => &[], + } + } + + pub fn is_body_read_timeout(&self) -> bool { + matches!(self, Self::Http(error) if error.is_body_read_timeout()) + } + + pub fn is_timeout(&self) -> Option { + match self { + Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()), + _ => None, + } + } + pub fn upstream_response(&self) -> Option<(u16, &str)> { match self { Self::Http(error) => error.upstream_response(), @@ -158,6 +176,10 @@ pub enum ErrorDetail { value: String, hint: String, }, + #[error( + "thinking.budget_tokens must be a number, got {actual}. To drop it, set `litellm.drop_params=True`, or `drop_params: true` under `litellm_settings` on the proxy" + )] + NonNumericThinkingBudget { actual: &'static str }, #[error( "Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({{LLM_PROVIDER}}_API_KEY) or via params" )] diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index e5e3a056fab..038617c2e10 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -77,7 +77,7 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - let request = transform_messages_request(request, context)?; + let request = transform_messages_request(request, &context.keeping_display_updates())?; let request = strip_cache_control_scope(normalize_system_role_messages( request, context diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index 8f7549c1c6d..141702bf857 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -186,6 +186,9 @@ pub struct ModelInfoFields { pub supports_image_size: Option, pub supported_audio_formats: Option, pub token_count_framing: Option, + pub token_count_tool_framing: Option, + pub token_count_responses_framing: Option, + pub token_count_responses_tool_framing: Option, pub token_count_model: Option, pub token_count_tokenizer: Option, pub vertex_ai_audio_api: Option, diff --git a/litellm-rust/crates/model-catalog/src/vocabulary.rs b/litellm-rust/crates/model-catalog/src/vocabulary.rs index 2b7ac8bd1ab..6bf2cc55224 100644 --- a/litellm-rust/crates/model-catalog/src/vocabulary.rs +++ b/litellm-rust/crates/model-catalog/src/vocabulary.rs @@ -93,6 +93,12 @@ pub enum Endpoint { Realtime, #[serde(rename = "/v1/realtime/transcription_sessions")] RealtimeTranscriptionSessions, + #[serde(rename = "/v1/realtime/translations")] + RealtimeTranslations, + #[serde(rename = "/v1/realtime/translations/calls")] + RealtimeTranslationCalls, + #[serde(rename = "/v1/realtime/translations/client_secrets")] + RealtimeTranslationClientSecrets, #[serde(rename = "/v1/videos")] Videos, #[serde(rename = "/v1/ocr")] diff --git a/litellm-rust/crates/model-catalog/src/wire/attributes.rs b/litellm-rust/crates/model-catalog/src/wire/attributes.rs index 93c10a81524..307f25f0bda 100644 --- a/litellm-rust/crates/model-catalog/src/wire/attributes.rs +++ b/litellm-rust/crates/model-catalog/src/wire/attributes.rs @@ -11,7 +11,7 @@ use crate::error::{DateCause, Error, Repr}; use crate::vocabulary::{ AudioFormat, EffortLevel, Endpoint, InputModality, Mode, OutputModality, VertexAiAudioApi, }; -use crate::wire::lax::Count; +use crate::wire::lax::{Count, NumericCount}; macro_rules! attributes { ($($(#[$doc:meta])* $name:ident: $ty:ty $(as $via:ty)?),* $(,)?) => { @@ -104,9 +104,14 @@ attributes! { /// Cloud regions the model is available in ('global' or region ids). supported_regions: Vec, thinking_always_on: bool, - token_count_framing: u64 as Count, + /// Whether the model thinks when the request carries no thinking param. + thinking_on_by_default: bool, + token_count_framing: u64 as NumericCount, token_count_model: TokenCountModel, + token_count_responses_framing: u64 as NumericCount, + token_count_responses_tool_framing: u64 as NumericCount, token_count_tokenizer: TokenCountTokenizer, + token_count_tool_framing: u64 as NumericCount, /// Provider default tokens-per-minute limit. tpm: u64 as Count, use_openai_responses_path: bool, @@ -298,10 +303,24 @@ mod tests { #[case::known_mode("mode", json!("chat"), true)] #[case::unknown_mode("mode", json!("chatt"), false)] #[case::negative_dbu("input_dbu_cost_per_token", json!(-1), false)] + #[case::thinking_on_by_default("thinking_on_by_default", json!(true), true)] + #[case::thinking_on_by_default_not_a_flag("thinking_on_by_default", json!("yes"), false)] #[case::reasoning_contexts("reasoning_contexts", json!(["reasoning", "non_reasoning"]), true)] #[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)] #[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)] #[case::reasoning_mode_not_a_string("reasoning_modes", json!([1]), false)] + #[case::framing("token_count_framing", json!(12), true)] + #[case::framing_written_as_text("token_count_framing", json!("12"), false)] + #[case::framing_written_as_whole_float("token_count_framing", json!(12.0), true)] + #[case::fractional_tool_framing("token_count_tool_framing", json!(1.5), false)] + #[case::tool_framing("token_count_tool_framing", json!(214), true)] + #[case::tool_framing_written_as_text("token_count_tool_framing", json!("214"), false)] + #[case::tool_framing_as_a_flag("token_count_tool_framing", json!(true), false)] + #[case::negative_tool_framing("token_count_tool_framing", json!(-1), false)] + #[case::responses_framing("token_count_responses_framing", json!(38), true)] + #[case::responses_framing_as_a_flag("token_count_responses_framing", json!(false), false)] + #[case::responses_tool_framing("token_count_responses_tool_framing", json!(267), true)] + #[case::responses_tool_framing_written_as_text("token_count_responses_tool_framing", json!("267"), false)] fn each_attribute_reads_its_own_type( #[case] key: &str, #[case] value: Value, diff --git a/litellm-rust/crates/model-catalog/src/wire/lax.rs b/litellm-rust/crates/model-catalog/src/wire/lax.rs index d514a6a82c1..d98a601e8c1 100644 --- a/litellm-rust/crates/model-catalog/src/wire/lax.rs +++ b/litellm-rust/crates/model-catalog/src/wire/lax.rs @@ -8,6 +8,8 @@ pub(crate) struct LaxI64; pub(crate) struct Count; +pub(crate) struct NumericCount; + impl<'de> DeserializeAs<'de, i64> for LaxI64 { fn deserialize_as>(deserializer: D) -> Result { deserializer.deserialize_any(Self) @@ -52,6 +54,34 @@ impl<'de> DeserializeAs<'de, u64> for Count { } } +impl<'de> DeserializeAs<'de, u64> for NumericCount { + fn deserialize_as>(deserializer: D) -> Result { + deserializer.deserialize_any(Self) + } +} + +impl Visitor<'_> for NumericCount { + type Value = u64; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a whole number of zero or more") + } + + fn visit_u64(self, value: u64) -> Result { + Ok(value) + } + + fn visit_i64(self, value: i64) -> Result { + u64::try_from(value).map_err(|_| E::invalid_value(de::Unexpected::Signed(value), &self)) + } + + fn visit_f64(self, value: f64) -> Result { + integral_float(value) + .and_then(|whole| u64::try_from(whole).ok()) + .ok_or_else(|| E::invalid_value(de::Unexpected::Float(value), &self)) + } +} + fn integer_string(value: &str) -> Option { let integer = match value.split_once('.') { Some((integer, fraction)) => { @@ -99,6 +129,12 @@ mod tests { .map(DeserializeAsWrap::into_inner) } + fn numeric_count(value: Value) -> Option { + DeserializeAsWrap::::deserialize(value) + .ok() + .map(DeserializeAsWrap::into_inner) + } + use serde::Deserialize; #[rstest] @@ -125,4 +161,23 @@ mod tests { ) { assert_eq!(count(written), expected); } + + #[rstest] + #[case::integer(json!(214), Some(214))] + #[case::zero(json!(0), Some(0))] + #[case::whole_float(json!(214.0), Some(214))] + #[case::negative_zero(json!(-0.0), Some(0))] + #[case::above_i64(json!(u64::MAX), Some(u64::MAX))] + #[case::negative(json!(-1), None)] + #[case::negative_whole_float(json!(-1.0), None)] + #[case::fraction(json!(1.5), None)] + #[case::string(json!("214"), None)] + #[case::boolean(json!(true), None)] + #[case::list(json!([214]), None)] + fn numeric_counts_accept_only_whole_json_numbers( + #[case] written: Value, + #[case] expected: Option, + ) { + assert_eq!(numeric_count(written), expected); + } } diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index b42c08b7e78..426cd861729 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -23,6 +23,17 @@ pyo3::create_exception!( ); const MESSAGES_REQUEST_ERROR: &str = "messages_request_error"; +/// The rejection is a parameter the model does not support, which Python raises as +/// `UnsupportedParamsError`. +const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param"; +/// The phase label (`Connection` or `Read`) of a timed-out call, for Python's timeout text, or +/// `Body` for a streamed body read that timed out after the head arrived. +const MESSAGES_TIMEOUT: &str = "messages_timeout"; +/// No credential reached the provider, which Python raises as `AuthenticationError`. +const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential"; +/// Python refuses the request only once its logger knows the resolved model and provider (the +/// provider config's request transform), so the failure logs them. +const MESSAGES_LATE_REJECTION: &str = "messages_late_rejection"; pub(crate) fn upstream_error( py: Python<'_>, @@ -78,17 +89,57 @@ pub(crate) fn route_failure( return Ok(raised); } if let Some((status, body)) = error.upstream_response() { - return public(upstream_error(py, status, body, &HeaderMap::new())?); + let headers = error + .upstream_headers() + .iter() + .filter_map(|(name, value)| { + Some(( + reqwest::header::HeaderName::from_bytes(name.as_bytes()).ok()?, + reqwest::header::HeaderValue::from_str(value).ok()?, + )) + }) + .collect::(); + return public(upstream_error(py, status, body, &headers)?); } let native = match (marker, error.rejected_request()) { (RequestMarker::Messages, Some(rejection)) => { let native = PyValueError::new_err(Report::new(rejection).to_string()); - native.value(py).setattr(MESSAGES_REQUEST_ERROR, true)?; + let value = native.value(py); + value.setattr(MESSAGES_REQUEST_ERROR, true)?; + if matches!( + rejection.downcast_ref::(), + Some( + litellm_llms::ErrorDetail::UnsupportedParameter { .. } + | litellm_llms::ErrorDetail::NonNumericThinkingBudget { .. } + ) + ) { + value.setattr(MESSAGES_UNSUPPORTED_PARAM, true)?; + } + if matches!( + rejection.downcast_ref::(), + Some(litellm_llms::ErrorDetail::NonNumericThinkingBudget { .. }) + ) { + value.setattr(MESSAGES_LATE_REJECTION, true)?; + } native } - (RequestMarker::Messages, None) | (RequestMarker::Unmarked, _) => { - by_fault(error.is_request(), &error) + (RequestMarker::Messages, None) => { + let native = by_fault(error.is_request(), &error); + let value = native.value(py); + if let Some(connect) = error.timeout() { + let phase = match (connect, error.body_read_timeout()) { + (true, _) => "Connection", + (false, true) => "Body", + (false, false) => "Read", + }; + value.setattr(MESSAGES_TIMEOUT, phase)?; + } + if error.is_missing_credential() { + value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?; + } + native } + (RequestMarker::Unmarked, _) => by_fault(error.is_request(), &error), }; public(native) } 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 27ce8c943f0..af614f17360 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -34,11 +34,9 @@ fn merge_headers( forwarded: Option>, extra_headers: Option>, ) -> Option> { - let merged: Map = forwarded - .into_iter() - .flatten() - .chain(extra_headers.into_iter().flatten()) - .collect(); + let merged = litellm_inference_messages::merge_messages_header_sources( + forwarded.into_iter().chain(extra_headers), + ); (!merged.is_empty()).then_some(merged) } @@ -107,10 +105,19 @@ impl MessagesPythonHost { .map(|value| python_timeout_seconds(py, value.unbind())) .transpose()? .flatten(); - let custom_llm_provider = string("custom_llm_provider")?; - let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; let api_key = string("api_key")?; let api_base = string("api_base")?; + // Python's handler resolves a bare model name (`claude-sonnet-4-5`) through + // `get_llm_provider`'s model lists, which the native prefix lookup does not carry. + let custom_llm_provider = match string("custom_llm_provider")? { + Some(provider) => Some(provider), + None => py + .import(ROUTE_HOST_MODULE)? + .getattr("inferred_provider")? + .call1((&model, api_base.as_deref()))? + .extract()?, + }; + let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; let extra_headers = self.merged_headers(py, arguments)?; let provider_specific_header = self.provider_specific_header(py, arguments)?; let litellm_params = project_litellm_params(argument, |name| module_global(py, name))?; 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 c788760b268..a9ae01bd7ee 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -50,9 +50,18 @@ fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyR ) } +/// Python's sync Messages handler raises for every native provider before it sends anything +/// or opens a logging object (`BaseLLMHTTPHandler.anthropic_messages_handler`), and Rust copies +/// that until both sides serve sync calls. +const SYNC_NOT_IMPLEMENTED: &str = "anthropic_messages_handler is not implemented for sync calls"; + #[pyfunction] -pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { - run_messages(py, call, false) +#[pyo3(signature = (call))] +pub(crate) fn messages(call: NativeCall<'_>) -> PyResult> { + drop(call); + Err(pyo3::exceptions::PyValueError::new_err( + SYNC_NOT_IMPLEMENTED, + )) } #[pyfunction] diff --git a/litellm/__init__.py b/litellm/__init__.py index 0e5101b49c3..102c4d1440f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -8,7 +8,7 @@ _apply_process_defaults() import threading import os -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import ( Any, Callable, @@ -1562,7 +1562,9 @@ from .realtime_api.main import ( _arealtime, acreate_realtime_client_secret, acreate_realtime_transcription_session, + acreate_realtime_translation_client_secret as acreate_realtime_translation_client_secret, arealtime_calls, + arealtime_translation_calls as arealtime_translation_calls, ) from .responses.main import _aresponses_websocket @@ -1754,6 +1756,9 @@ if TYPE_CHECKING: from .llms.typesafe.decisions.transformation import ( TypeSafeDecisionsConfig as TypeSafeDecisionsConfig, ) + from .llms.parallel_ai.decisions.transformation import ( + ParallelAIDecisionsConfig as ParallelAIDecisionsConfig, + ) from .llms.openrouter.decisions.transformation import ( OpenRouterDecisionsConfig as OpenRouterDecisionsConfig, ) @@ -1808,6 +1813,9 @@ if TYPE_CHECKING: from .llms.vertex_ai.rerank.transformation import ( VertexAIRerankConfig as VertexAIRerankConfig, ) + from .llms.together_ai.chat.transformation import ( + TogetherAIChatConfig as TogetherAIChatConfig, + ) from .llms.fireworks_ai.rerank.transformation import ( FireworksAIRerankConfig as FireworksAIRerankConfig, ) @@ -1853,9 +1861,6 @@ if TYPE_CHECKING: BedrockMantleAnthropicMessagesConfig as BedrockMantleAnthropicMessagesConfig, ) from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig - from .llms.together_ai.chat.transformation import ( - TogetherAIChatConfig as TogetherAIChatConfig, - ) from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig as VertexGeminiConfig, @@ -2014,6 +2019,9 @@ if TYPE_CHECKING: from .llms.xai.responses.transformation import ( XAIResponsesAPIConfig as XAIResponsesAPIConfig, ) + from .llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig as VertexAIInteractionsConfig, + ) from .llms.litellm_proxy.responses.transformation import ( LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig, ) @@ -2041,9 +2049,6 @@ if TYPE_CHECKING: from .llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig, ) - from .llms.vertex_ai.interactions.transformation import ( - VertexAIInteractionsConfig as VertexAIInteractionsConfig, - ) from .llms.openai.chat.o_series_transformation import ( OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config, @@ -2374,6 +2379,7 @@ if TYPE_CHECKING: supports_web_search: Callable[..., bool] supports_tool_search: Callable[..., bool] supports_url_context: Callable[..., bool] + supports_continuation: Callable[..., bool] supports_response_schema: Callable[..., bool] supports_parallel_function_calling: Callable[..., bool] supports_vision: Callable[..., bool] diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index 40d3377bf66..f2f7869876d 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -52,6 +52,18 @@ def internal_sub_call(parent_request_id: str | None = None) -> Generator[None]: internal_parent_request_id.reset(parent_token) +@contextmanager +def billed_separately() -> Generator[None]: + """A call that writes its own spend row even inside a sub-call whose parent bills the model response.""" + token: Final = is_internal_call.set(False) + parent_token: Final = internal_parent_request_id.set(None) + try: + yield + finally: + is_internal_call.reset(token) + internal_parent_request_id.reset(parent_token) + + @contextmanager def post_response_phase() -> Generator[None]: """Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns.""" diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index de290070ed3..65074093dbb 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -33,6 +33,7 @@ UTILS_NAMES: Final = ( "supports_web_search", "supports_tool_search", "supports_url_context", + "supports_continuation", "supports_response_schema", "supports_parallel_function_calling", "supports_vision", @@ -162,6 +163,7 @@ LLM_CONFIG_NAMES: Final = ( "HostedVLLMRerankConfig", "PerplexityDecisionsConfig", "TypeSafeDecisionsConfig", + "ParallelAIDecisionsConfig", "NaceDecisionsConfig", "OpenRouterDecisionsConfig", "FireworksAIDecisionsConfig", @@ -488,6 +490,7 @@ _UTILS_IMPORT_MAP: Final = { "supports_web_search": (".utils", "supports_web_search"), "supports_tool_search": (".utils", "supports_tool_search"), "supports_url_context": (".utils", "supports_url_context"), + "supports_continuation": (".utils", "supports_continuation"), "supports_response_schema": (".utils", "supports_response_schema"), "supports_parallel_function_calling": ( ".utils", @@ -738,6 +741,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "PerplexityDecisionsConfig": (".llms.perplexity.decisions.transformation", "PerplexityDecisionsConfig"), "NaceDecisionsConfig": (".llms.nace.decisions.transformation", "NaceDecisionsConfig"), "TypeSafeDecisionsConfig": (".llms.typesafe.decisions.transformation", "TypeSafeDecisionsConfig"), + "ParallelAIDecisionsConfig": (".llms.parallel_ai.decisions.transformation", "ParallelAIDecisionsConfig"), "OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"), "FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"), "CloudflareDecisionsConfig": (".llms.cloudflare.decisions.transformation", "CloudflareDecisionsConfig"), diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 253074b0e1f..bdfddb48522 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -37,7 +37,47 @@ "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01", "thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18", "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", - "inline-tools-2026-09-15": "inline-tools-2026-09-15" + "inline-tools-2026-09-15": "inline-tools-2026-09-15", + "agent-memory-2026-07-22": null, + "analytics-usage-cost-2026-08-24": "analytics-usage-cost-2026-08-24", + "cache-diagnosis-2026-04-07": "cache-diagnosis-2026-04-07", + "ce-plugins-2026-09-01": "ce-plugins-2026-09-01", + "ce-user-management-2026-07-13": "ce-user-management-2026-07-13", + "claude-code-20250219": "claude-code-20250219", + "code-execution-2025-05-22": "code-execution-2025-05-22", + "computer-use-2024-10-22": "computer-use-2024-10-22", + "dev-full-thinking-2025-05-14": null, + "dreaming-2026-04-21": "dreaming-2026-04-21", + "experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21", + "extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11", + "fallback-credit-2026-06-01": "fallback-credit-2026-06-01", + "fallback-credit-2026-07-01": "fallback-credit-2026-07-01", + "managed-agents-2026-04-01": "managed-agents-2026-04-01", + "mcp-client-2026-09-15": "mcp-client-2026-09-15", + "mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22", + "message-batches-2024-09-24": "message-batches-2024-09-24", + "message-threads-2026-08-12": "message-threads-2026-08-12", + "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", + "model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26", + "output-300k-2026-03-24": "output-300k-2026-03-24", + "pdfs-2024-09-25": "pdfs-2024-09-25", + "prompt-caching-2024-07-31": "prompt-caching-2024-07-31", + "server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01", + "server-side-fallback-2026-07-01": "server-side-fallback-2026-07-01", + "spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26", + "task-budgets-2026-03-13": "task-budgets-2026-03-13", + "telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11", + "thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13", + "token-counting-2024-11-01": "token-counting-2024-11-01", + "tool-examples-2025-10-29": "tool-examples-2025-10-29", + "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", + "user-profiles-2026-03-24": "user-profiles-2026-03-24", + "user-profiles-2026-08-18": "user-profiles-2026-08-18", + "user-profiles-2026-09-04": "user-profiles-2026-09-04", + "webhooks-2026-10-07": "webhooks-2026-10-07", + "afk-mode-2026-01-31": "afk-mode-2026-01-31", + "redact-thinking-2026-02-12": "redact-thinking-2026-02-12", + "structured-outputs-2025-12-15": "structured-outputs-2025-12-15" }, "azure_ai": { "advisor-tool-2026-03-01": null, @@ -45,23 +85,21 @@ "bash_20241022": null, "bash_20250124": null, "code-execution-2025-08-25": "code-execution-2025-08-25", - "compact-2026-01-12": null, - "compact-2026-09-04": "compact-2026-09-04", + "compact-2026-01-12": "compact-2026-01-12", "computer-use-2025-01-24": "computer-use-2025-01-24", "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", - "fast-mode-2026-02-01": null, + "fast-mode-2026-02-01": "fast-mode-2026-02-01", "files-api-2025-04-14": "files-api-2025-04-14", - "fine-grained-tool-streaming-2025-05-14": null, + "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-11-20": "mcp-client-2025-11-20", "mcp-client-2025-04-04": "mcp-client-2025-04-04", "mcp-servers-2025-12-04": null, - "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", - "output-128k-2025-02-19": null, + "output-128k-2025-02-19": "output-128k-2025-02-19", "structured-output-2024-03-01": null, "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05", @@ -71,9 +109,54 @@ "text_editor_20250124": null, "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", "thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18", - "token-efficient-tools-2025-02-19": null, + "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19", "web-fetch-2025-09-10": "web-fetch-2025-09-10", - "web-search-2025-03-05": "web-search-2025-03-05" + "web-search-2025-03-05": "web-search-2025-03-05", + "agent-memory-2026-07-22": null, + "analytics-usage-cost-2026-08-24": "analytics-usage-cost-2026-08-24", + "cache-diagnosis-2026-04-07": "cache-diagnosis-2026-04-07", + "ce-plugins-2026-09-01": "ce-plugins-2026-09-01", + "ce-user-management-2026-07-13": "ce-user-management-2026-07-13", + "claude-code-20250219": "claude-code-20250219", + "code-execution-2025-05-22": "code-execution-2025-05-22", + "compact-2026-09-04": "compact-2026-09-04", + "computer-use-2024-10-22": "computer-use-2024-10-22", + "dev-full-thinking-2025-05-14": null, + "dreaming-2026-04-21": "dreaming-2026-04-21", + "experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21", + "extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11", + "fallback-credit-2026-06-01": "fallback-credit-2026-06-01", + "fallback-credit-2026-07-01": "fallback-credit-2026-07-01", + "inline-tools-2026-09-15": "inline-tools-2026-09-15", + "managed-agents-2026-04-01": "managed-agents-2026-04-01", + "mcp-client-2026-09-15": "mcp-client-2026-09-15", + "mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22", + "message-batches-2024-09-24": "message-batches-2024-09-24", + "message-threads-2026-08-12": "message-threads-2026-08-12", + "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01", + "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", + "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", + "model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26", + "oauth-2025-04-20": "oauth-2025-04-20", + "output-300k-2026-03-24": "output-300k-2026-03-24", + "pdfs-2024-09-25": "pdfs-2024-09-25", + "prompt-caching-2024-07-31": "prompt-caching-2024-07-31", + "server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01", + "server-side-fallback-2026-07-01": "server-side-fallback-2026-07-01", + "spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26", + "task-budgets-2026-03-13": "task-budgets-2026-03-13", + "telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11", + "thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13", + "token-counting-2024-11-01": "token-counting-2024-11-01", + "tool-examples-2025-10-29": "tool-examples-2025-10-29", + "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", + "user-profiles-2026-03-24": "user-profiles-2026-03-24", + "user-profiles-2026-08-18": "user-profiles-2026-08-18", + "user-profiles-2026-09-04": "user-profiles-2026-09-04", + "webhooks-2026-10-07": "webhooks-2026-10-07", + "afk-mode-2026-01-31": "afk-mode-2026-01-31", + "redact-thinking-2026-02-12": "redact-thinking-2026-02-12", + "structured-outputs-2025-12-15": "structured-outputs-2025-12-15" }, "bedrock_converse": { "advisor-tool-2026-03-01": null, @@ -196,34 +279,92 @@ "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", - "effort-2025-11-24": null, + "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, - "fine-grained-tool-streaming-2025-05-14": null, + "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, "mcp-servers-2025-12-04": null, - "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", - "output-128k-2025-02-19": null, + "output-128k-2025-02-19": "output-128k-2025-02-19", "structured-output-2024-03-01": null, "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", "prompt-caching-scope-2026-01-05": null, "skills-2025-10-02": null, - "structured-outputs-2025-11-13": null, + "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", "text_editor_20241022": null, "text_editor_20250124": null, "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", "thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18", - "token-efficient-tools-2025-02-19": null, + "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19", "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", "web-fetch-2025-09-10": null, "web-search-2025-03-05": "web-search-2025-03-05", - "inline-tools-2026-09-15": "inline-tools-2026-09-15" + "inline-tools-2026-09-15": "inline-tools-2026-09-15", + "agent-memory-2026-07-22": null, + "analytics-usage-cost-2026-08-24": null, + "cache-diagnosis-2026-04-07": null, + "ce-plugins-2026-09-01": null, + "ce-user-management-2026-07-13": null, + "claude-code-20250219": "claude-code-20250219", + "code-execution-2025-05-22": null, + "compact-2026-09-04": null, + "computer-use-2024-10-22": null, + "dev-full-thinking-2025-05-14": null, + "dreaming-2026-04-21": null, + "experimental-claude-tag-2026-10-21": null, + "extended-cache-ttl-2025-04-11": null, + "fallback-credit-2026-06-01": "fallback-credit-2026-06-01", + "fallback-credit-2026-07-01": "fallback-credit-2026-07-01", + "managed-agents-2026-04-01": null, + "mcp-client-2026-09-15": null, + "mcp-tunnels-2026-06-22": null, + "message-batches-2024-09-24": null, + "message-threads-2026-08-12": null, + "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01", + "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", + "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", + "model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26", + "oauth-2025-04-20": null, + "output-300k-2026-03-24": null, + "pdfs-2024-09-25": "pdfs-2024-09-25", + "prompt-caching-2024-07-31": null, + "server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01", + "server-side-fallback-2026-07-01": null, + "spend-limit-reads-2026-09-26": null, + "task-budgets-2026-03-13": "task-budgets-2026-03-13", + "telemetry-destinations-2026-08-11": null, + "thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13", + "token-counting-2024-11-01": "token-counting-2024-11-01", + "tool-examples-2025-10-29": "tool-examples-2025-10-29", + "user-profiles-2026-03-24": null, + "user-profiles-2026-08-18": null, + "user-profiles-2026-09-04": null, + "webhooks-2026-10-07": null, + "afk-mode-2026-01-31": "afk-mode-2026-01-31", + "redact-thinking-2026-02-12": null, + "structured-outputs-2025-12-15": "structured-outputs-2025-12-15" }, "vertex_ai_header_only": { + "ce-plugins-2026-09-01": "ce-plugins-2026-09-01", + "ce-user-management-2026-07-13": "ce-user-management-2026-07-13", "compact-2026-09-04": "compact-2026-09-04", - "inline-tools-2026-09-15": "inline-tools-2026-09-15" + "dreaming-2026-04-21": "dreaming-2026-04-21", + "experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21", + "extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11", + "inline-tools-2026-09-15": "inline-tools-2026-09-15", + "managed-agents-2026-04-01": "managed-agents-2026-04-01", + "mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22", + "oauth-2025-04-20": "oauth-2025-04-20", + "prompt-caching-2024-07-31": "prompt-caching-2024-07-31", + "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05", + "redact-thinking-2026-02-12": "redact-thinking-2026-02-12", + "spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26", + "telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11", + "user-profiles-2026-03-24": "user-profiles-2026-03-24", + "user-profiles-2026-08-18": "user-profiles-2026-08-18", + "user-profiles-2026-09-04": "user-profiles-2026-09-04" }, "databricks": { "advisor-tool-2026-03-01": null, diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index be2ff07e088..e284845b2ff 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -31,6 +31,7 @@ from concurrent.futures import ThreadPoolExecutor from functools import cache from importlib.resources import files from typing import Final, Literal, TypeAlias +from urllib.parse import urlparse import httpx from pydantic import TypeAdapter @@ -48,6 +49,7 @@ _BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) _DECODED_JSON: Final = TypeAdapter(object) BetaPlacement: TypeAlias = Literal["header", "body"] +UnlistedBetas: TypeAlias = Literal["drop", "forward"] _HEADER_ONLY_COLUMN_SUFFIX: Final = "_header_only" @@ -280,12 +282,13 @@ def filter_and_transform_beta_headers( beta_headers: list[str], provider: str, placement: BetaPlacement = "header", + unlisted: UnlistedBetas = "drop", ) -> list[str]: """ Filter and transform beta headers based on provider's mapping configuration. This function: - 1. Only allows headers that are present in the provider's mapping keys + 1. Only allows headers that are present in the provider's mapping keys, unless ``unlisted`` is "forward" 2. Filters out headers with null values (unsupported) 3. Maps headers to provider-specific names (e.g., advanced-tool-use -> tool-search-tool) @@ -293,6 +296,7 @@ def filter_and_transform_beta_headers( beta_headers: List of Anthropic beta header values provider: Provider name (e.g., "anthropic", "bedrock", "vertex_ai") placement: Whether the betas travel as the anthropic-beta header or the anthropic_beta body field + unlisted: Whether a beta missing from the mapping is dropped or forwarded as sent Returns: List of filtered and transformed beta headers for the provider @@ -307,6 +311,10 @@ def filter_and_transform_beta_headers( for header in beta_headers: header = header.strip() + if header not in provider_mapping and unlisted == "forward": + filtered_headers.add(header) + continue + # Check if header is in the mapping if header not in provider_mapping: verbose_logger.debug( @@ -388,9 +396,38 @@ def get_provider_beta_header( return provider_mapping[anthropic_beta_header] +FOUNDRY_API_BASE_HOST_SUFFIXES: Final[tuple[str, ...]] = ( + "services.ai.azure.com", + "openai.azure.com", + "cognitiveservices.azure.com", +) + + +def is_foundry_api_base(api_base: str | None) -> bool: + """Whether ``api_base`` points at an Azure AI Foundry host, where Claude is served by Microsoft.""" + if not api_base: + return False + host: Final = urlparse(api_base if "://" in api_base else f"https://{api_base}").hostname or "" + return any(host == suffix or host.endswith(f".{suffix}") for suffix in FOUNDRY_API_BASE_HOST_SUFFIXES) + + +def resolve_beta_provider(provider: str, api_base: str | None = None) -> str: + """ + Pick the beta allowlist row for one deployment, not one provider string. + + An ``anthropic`` deployment whose ``api_base`` is a Foundry host is served by Microsoft and + gets the ``azure_ai`` row, since Foundry accepts a different set of betas than api.anthropic.com. + """ + if provider == "anthropic" and is_foundry_api_base(api_base): + return "azure_ai" + return provider + + def update_headers_with_filtered_beta( headers: dict, provider: str, + api_base: str | None = None, + unlisted: UnlistedBetas = "drop", ) -> dict: """ Update headers dict by filtering and transforming anthropic-beta header values. @@ -399,6 +436,8 @@ def update_headers_with_filtered_beta( Args: headers: Request headers dict (will be modified in place) provider: Provider name + api_base: Deployment api_base, used to pick the allowlist row (Foundry hosts use ``azure_ai``) + unlisted: Whether a beta missing from the row is dropped or forwarded as sent Returns: Updated headers dict @@ -413,7 +452,8 @@ def update_headers_with_filtered_beta( # Filter and transform based on provider filtered_beta_values: Final = filter_and_transform_beta_headers( beta_headers=beta_values, - provider=provider, + provider=resolve_beta_provider(provider, api_base), + unlisted=unlisted, ) # Update or remove the header @@ -430,6 +470,7 @@ def update_request_with_filtered_beta( headers: dict, request_data: dict, provider: str, + api_base: str | None = None, ) -> tuple[dict[str, str], dict[str, object]]: """ Update both headers and request body beta fields based on provider support. @@ -439,11 +480,12 @@ def update_request_with_filtered_beta( headers: Request headers dict (will be modified in place) request_data: Request body dict (will be modified in place) provider: Provider name + api_base: Deployment api_base, used to pick the allowlist row (Foundry hosts use ``azure_ai``) Returns: Tuple of (updated headers, updated request_data) """ - headers = update_headers_with_filtered_beta(headers=headers, provider=provider) + headers = update_headers_with_filtered_beta(headers=headers, provider=provider, api_base=api_base) existing_body_betas: Final = request_data.get("anthropic_beta") if not existing_body_betas: @@ -451,7 +493,7 @@ def update_request_with_filtered_beta( filtered_body_betas: Final = filter_and_transform_beta_headers( beta_headers=existing_body_betas, - provider=provider, + provider=resolve_beta_provider(provider, api_base), placement="body", ) diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index d8198e41239..b4c95fafff3 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -16,7 +16,7 @@ from .exceptions import AnthropicErrorDetail, AnthropicErrorResponse, AnthropicE _PARSED_ERROR_BODY: Final = TypeAdapter(object) # HTTP status code -> Anthropic error type -# Source: https://docs.anthropic.com/en/api/errors +# Source: https://platform.claude.com/docs/en/api/errors ANTHROPIC_ERROR_TYPE_MAP: Final[dict[int, AnthropicErrorType]] = { 400: "invalid_request_error", 401: "authentication_error", @@ -46,6 +46,7 @@ class AnthropicExceptionMapping: status_code: int, message: str, request_id: str | None = None, + error_type: AnthropicErrorType | None = None, ) -> AnthropicErrorResponse: """ Create an Anthropic-formatted error response dict. @@ -57,12 +58,10 @@ class AnthropicExceptionMapping: "request_id": "req_..." } """ - error_type: Final = AnthropicExceptionMapping.get_error_type(status_code) - response: Final[AnthropicErrorResponse] = { "type": "error", "error": { - "type": error_type, + "type": error_type or AnthropicExceptionMapping.get_error_type(status_code), "message": message, }, } @@ -131,6 +130,7 @@ class AnthropicExceptionMapping: status_code: int, raw_message: str, request_id: str | None = None, + error_type: AnthropicErrorType | None = None, ) -> AnthropicErrorResponse: """ Transform an error message to Anthropic format. @@ -144,6 +144,7 @@ class AnthropicExceptionMapping: status_code: HTTP status code raw_message: Raw error message (may be JSON string or plain text) request_id: Optional request ID to include + error_type: Anthropic error type to use instead of the one the status maps to Returns: AnthropicErrorResponse dict @@ -169,6 +170,7 @@ class AnthropicExceptionMapping: status_code=status_code, message=message, request_id=request_id, + error_type=error_type, ) diff --git a/litellm/anthropic_interface/exceptions/exceptions.py b/litellm/anthropic_interface/exceptions/exceptions.py index b48cd2fee6f..18b3ca37a3b 100644 --- a/litellm/anthropic_interface/exceptions/exceptions.py +++ b/litellm/anthropic_interface/exceptions/exceptions.py @@ -6,15 +6,18 @@ from typing import Literal from typing_extensions import NotRequired, ReadOnly, Required, TypedDict # Known Anthropic error types -# Source: https://docs.anthropic.com/en/api/errors +# Source: https://platform.claude.com/docs/en/api/errors AnthropicErrorType = Literal[ "invalid_request_error", "authentication_error", + "billing_error", "permission_error", "not_found_error", + "conflict_error", "request_too_large", "rate_limit_error", "api_error", + "timeout_error", "overloaded_error", ] @@ -25,6 +28,7 @@ class AnthropicErrorDetail(TypedDict): type: AnthropicErrorType message: str provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]] + litellm_error_type: NotRequired[ReadOnly[str]] litellm_call_id: NotRequired[ReadOnly[str]] diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 91cf9bb2cfe..3b7ea66ad6b 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -45,13 +45,27 @@ from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache RESPONSE_CACHE_TARGET: Final = "llm_response" -_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"}) +_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token", "return_token_ids"}) _DECODED_LITERAL: Final = TypeAdapter(object) _REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter( list[dict[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True) ) +def _keyed_provider_params_in_extra_body(kwargs: Mapping[str, object]) -> str: + raw_extra_body: Final = kwargs.get("extra_body") + if not isinstance(raw_extra_body, dict): + return "" + extra_body: Final = cast( # cast-ok: isinstance leaves the key and value types unknown + Mapping[str, object], raw_extra_body + ) + return "".join( + f"extra_body.{param}: {extra_body[param]}" + for param in sorted(_ALWAYS_KEYED_PROVIDER_PARAMS) + if extra_body.get(param) + ) + + @contextmanager def response_cache_phase(operation: Literal["get", "set"]) -> Generator[None]: """The ``cache.get llm_response`` / ``cache.set llm_response`` span a response-cache read or write runs @@ -429,7 +443,10 @@ class Cache: "" if embeds_thought_signature_in_tool_call_id(kwargs) else "thought_signature_in_tool_call_id: False" ) hashed_cache_key = Cache._get_hashed_cache_key( - cache_key + reasoning_summary_material + tool_call_id_shape_material + cache_key + + _keyed_provider_params_in_extra_body(kwargs) + + reasoning_summary_material + + tool_call_id_shape_material ) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) verbose_logger.debug( diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index e6214a3650c..c82d25c9425 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -673,7 +673,7 @@ class DualCache(BaseCache): if self.redis_cache is not None: self.redis_cache.flush_cache() - def delete_cache(self, key): + def delete_cache(self, key: str) -> None: """ Delete a key from the cache """ diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 76d17ea9693..f6879c313de 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -11,7 +11,7 @@ from typing import ( cast, # noqa: TID251 # optional_params reaches the bridge as an untyped dict ) -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.openai import ResponsesAPIResponse @@ -20,15 +20,15 @@ if TYPE_CHECKING: class ResponsesToCompletionBridgeHandlerInputKwargs(TypedDict): - model: str - messages: list - optional_params: dict - litellm_params: dict - headers: dict - model_response: "ModelResponse" - logging_obj: "LiteLLMLoggingObj" - custom_llm_provider: str - encoding: object + model: ReadOnly[str] + messages: ReadOnly[list] + optional_params: ReadOnly[dict[str, object]] + litellm_params: ReadOnly[dict[str, object]] + headers: ReadOnly[dict] + model_response: ReadOnly["ModelResponse"] + logging_obj: ReadOnly["LiteLLMLoggingObj"] + custom_llm_provider: ReadOnly[str] + encoding: ReadOnly[object] def _restore_routing_prefix(model: str, custom_llm_provider: str) -> str: @@ -184,6 +184,7 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} + logging_obj.record_bridged_chat_request(optional_params, litellm_params) request_data: Final = self.transformation_handler.transform_request( model=model, @@ -285,6 +286,7 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} + logging_obj.record_bridged_chat_request(optional_params, litellm_params) try: request_data: Final = self.transformation_handler.transform_request( diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index bb59b4617af..8dd4d87ae13 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -16,6 +16,7 @@ from openai.types.responses.response_input_param import ( ResponseCustomToolCallOutputParam, ResponseCustomToolCallParam, ) +from openai.types.responses.tool_choice_allowed_param import ToolChoiceAllowedParam from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam @@ -33,11 +34,17 @@ from litellm.completion_extras.litellm_responses_transformation.history import ( replay_items, requested_phase, ) -from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint +from litellm.integrations.anthropic_cache_control_hook import routed_model_info, supports_openai_prompt_cache_breakpoint +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + parse_chat_allowed_tools_tool_choice, + responses_allowed_tools_tool_choice, +) +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( - prompt_cache_breakpoint_for_wire, + prompt_cache_breakpoint_with_readable_mode, responses_reasoning_items_from_thinking_blocks, + whole_prompt_cache_breakpoint_or_none, with_prompt_cache_breakpoint, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -97,6 +104,12 @@ _RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *Response ) +def _bridged_prompt_cache_breakpoint(marker: object, drop_params: bool, target_trims_breakpoints: bool) -> object: + if target_trims_breakpoints: + return prompt_cache_breakpoint_with_readable_mode(marker, drop_params) + return whole_prompt_cache_breakpoint_or_none(marker) if drop_params else marker + + def _strip_prompt_cache_breakpoint_from_content_block(value: object) -> object: if not isinstance(value, dict): return value @@ -775,13 +788,28 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _normalize_tool_choice_for_responses_api( self, tool_choice: _ToolChoiceT - ) -> _ToolChoiceT | ToolChoiceFunctionParam | ToolChoiceCustomParam | Literal["auto", "none", "required"]: + ) -> ( + _ToolChoiceT + | ToolChoiceFunctionParam + | ToolChoiceCustomParam + | ToolChoiceAllowedParam + | Literal["auto", "none", "required"] + ): """Chat tool_choice nests the name under function/custom; Responses API expects top-level name.""" if not isinstance(tool_choice, dict): return tool_choice choice_type: Final = tool_choice.get("type") if isinstance(choice_type, str) and choice_type in ("auto", "none", "required"): return choice_type + if choice_type == "allowed_tools": + parsed: Final = parse_chat_allowed_tools_tool_choice( + cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic + object, tool_choice + ) + ) + flat: Final = responses_allowed_tools_tool_choice(parsed) if parsed is not None else None + if flat is not None: + return flat if choice_type not in ("function", "custom"): return tool_choice if isinstance(tool_choice.get("name"), str) and tool_choice.get("name"): @@ -844,6 +872,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): messages: list["AllMessageValues"], *, drop_params: bool = False, + target_trims_breakpoints: bool = True, keep_prompt_cache_breakpoints: bool = False, replay_tool_search: bool = False, include_namespace: bool = False, @@ -852,6 +881,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input( messages, drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, replay_tool_search=replay_tool_search, include_namespace=include_namespace, replay_openai_history=replay_openai_history, @@ -868,6 +898,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): messages: list["AllMessageValues"], *, drop_params: bool = False, + target_trims_breakpoints: bool = True, replay_tool_search: bool = False, include_namespace: bool = False, replay_openai_history: bool = False, @@ -921,6 +952,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): content, role, drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, ), } ) @@ -940,6 +972,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): content, "user", # Use "user" role to get input_* types drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, ) else: # Fallback: convert unexpected types to input_text @@ -968,7 +1001,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "type": "message", "role": "assistant", "content": self._convert_content_to_responses_format( - content, "assistant", drop_params=drop_params + content, + "assistant", + drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, ), **_assistant_phase_field(msg, replay_openai_history, drop_params), } @@ -1009,7 +1045,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "type": "message", "role": role, "content": self._convert_content_to_responses_format( - content, cast(str, role), drop_params=drop_params + content, + cast(str, role), + drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, ), **_assistant_phase_field(msg, replay_openai_history and role == "assistant", drop_params), } @@ -1127,14 +1166,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): optional_params = self._extract_extra_body_params(optional_params) provider: Final = _served_provider(litellm_params) base_model: Final = litellm_params.get("base_model") - supports_prompt_cache_breakpoint: Final = supports_openai_prompt_cache_breakpoint(model, provider) or ( - isinstance(base_model, str) - and bool(base_model) - and supports_openai_prompt_cache_breakpoint(base_model, provider) + model_info: Final = routed_model_info(litellm_params.get("model_info")) + named_models: Final = (model, base_model) if isinstance(base_model, str) and base_model else (model,) + supports_prompt_cache_breakpoint: Final = any( + supports_openai_prompt_cache_breakpoint(name, provider, model_info) for name in named_models ) converted_input_items, converted_instructions = self.convert_chat_completion_messages_to_responses_api( messages, drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params), + target_trims_breakpoints=drops_unsupported_params(litellm_params.get("drop_params")), keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint, replay_tool_search=_requests_tool_search(cast(object, optional_params)), # cast-ok: upcast include_namespace=True, @@ -1195,10 +1235,6 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): api_model: Final = model - from litellm.types.utils import CallTypes - - setattr(litellm_logging_obj, "call_type", CallTypes.responses.value) - sanitized_litellm_params: Final = self._build_sanitized_litellm_params(litellm_params) request_data: Final = { @@ -1640,6 +1676,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): | None, role: str, drop_params: bool = False, + target_trims_breakpoints: bool = True, ) -> list[dict[str, object]]: """Convert chat completion content to responses API format""" from litellm.types.llms.openai import ChatCompletionImageObject @@ -1666,7 +1703,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if original_type == "text": converted = with_prompt_cache_breakpoint( self._convert_content_str_to_input_text(item.get("text", ""), role), - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: text -> %s", converted) @@ -1679,14 +1718,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): cast(ChatCompletionImageObject, item), role ), ), - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: image_url -> %s", converted) elif original_type == "input_audio": converted = with_prompt_cache_breakpoint( {"type": "input_audio", "input_audio": item.get("input_audio")}, - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: input_audio -> %s", converted) @@ -1702,7 +1745,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): _input_file_from_file_value( cast("ChatCompletionFileObject", item).get("file"), # cast-ok: type tag checked ), - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: file -> %s", converted) @@ -1726,7 +1771,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Default to input_text for unknown types converted = with_prompt_cache_breakpoint( self._convert_content_str_to_input_text(str(item.get("text", item)), role), - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: unknown(%s) -> %s", original_type, converted) diff --git a/litellm/constants.py b/litellm/constants.py index cd7c0176428..97631b9997a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -451,6 +451,7 @@ OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: Final = frozenset( {"response.completed", "response.incomplete", "response.failed"} ) REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" +REALTIME_SESSION_ENDED_ABNORMALLY_KEY: Final = "realtime_session_ended_abnormally" # SSL/TLS cipher configuration for faster handshakes # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones @@ -699,6 +700,9 @@ STREAM_SSE_DONE_STRING: Final[str] = "[DONE]" STREAM_SSE_DATA_PREFIX: Final[str] = "data: " STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n' STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8") +STREAM_FINISH_CHOICE_FIELDS: Final = ("stop_reason",) +STREAM_CHUNK_CHOICE_FIELDS: Final = ("token_ids",) +VLLM_CHOICE_FIELDS: Final = STREAM_FINISH_CHOICE_FIELDS + STREAM_CHUNK_CHOICE_FIELDS ### SPEND TRACKING ### DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = get_env_float( "DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400 @@ -968,7 +972,10 @@ OPENAI_CHAT_COMPLETION_PARAMS: Final = [ OPENAI_TRANSCRIPTION_PARAMS: Final = [ "language", + "languages", + "keywords", "response_format", + "stream", "timestamp_granularities", ] diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index e66be233be4..527981a1dc6 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -6,6 +6,7 @@ import time from collections.abc import Mapping, Sequence from datetime import datetime from functools import lru_cache +from itertools import accumulate from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast @@ -22,9 +23,14 @@ from litellm.constants import ( OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES, REALTIME_METERED_USAGE_KEY, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping, billed_usage_cost +from litellm.litellm_core_utils.audio_utils.utils import normalized_audio_duration_seconds +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, billed_usage_cost from litellm.litellm_core_utils.cost_discount import resolve_cost_discount -from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type +from litellm.litellm_core_utils.decisions_pricing import ( + decisions_pricing_fields, + decisions_request_price, + is_decisions_call_type, +) from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, @@ -319,12 +325,41 @@ def _get_additional_costs( return None -def _transcription_usage_has_token_details( +_TranscriptionRateKey: TypeAlias = Literal[ + "input_cost_per_second", + "output_cost_per_second", + "input_cost_per_audio_token", + "output_cost_per_audio_token", + "input_cost_per_token", + "output_cost_per_token", +] +_TRANSCRIPTION_TOKEN_RATE_KEYS: Final[tuple[_TranscriptionRateKey, ...]] = ( + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_audio_token", + "output_cost_per_audio_token", +) +_TRANSCRIPTION_SECOND_RATE_KEYS: Final[tuple[_TranscriptionRateKey, ...]] = ( + "input_cost_per_second", + "output_cost_per_second", +) + + +def _transcription_uses_token_pricing( + model: str, + custom_llm_provider: str | None, usage_block: Usage | None, ) -> bool: if usage_block is None: return False + model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + if ( + _declared_transcription_rate(model_info, _TRANSCRIPTION_TOKEN_RATE_KEYS) is None + and _declared_transcription_rate(model_info, _TRANSCRIPTION_SECOND_RATE_KEYS) is not None + ): + return False + prompt_tokens_val: Final = getattr(usage_block, "prompt_tokens", 0) or 0 completion_tokens_val: Final = getattr(usage_block, "completion_tokens", 0) or 0 prompt_details: Final[PromptTokensDetailsWrapper | None] = getattr(usage_block, "prompt_tokens_details", None) @@ -338,13 +373,38 @@ def _transcription_usage_has_token_details( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) -def _bills_provider_transcription_seconds(model: str | None, custom_llm_provider: str | None) -> bool: - if model is None or custom_llm_provider != LlmProviders.OPENAI.value: +def _bills_transcription_usage_duration(pricing_model: str | None, custom_llm_provider: str | None) -> bool: + if pricing_model is None or custom_llm_provider is None: return False + transcription_info: Final = _lookup_model_info_or_none(pricing_model, custom_llm_provider) + transcription_metadata: Final = ( + as_str_mapping(transcription_info.get("provider_specific_entry")) if transcription_info is not None else None + ) + return transcription_metadata is not None and transcription_metadata.get("transcription_usage_duration") == 1 + + +def _transcription_duration_for_cost( + completion_response: object, pricing_model: str | None, model: str | None, custom_llm_provider: str | None +) -> float: + hidden_params: Final = as_str_mapping(attribute_or_none(completion_response, "_hidden_params")) or {} + measured_duration: Final = normalized_audio_duration_seconds( + hidden_params.get("audio_transcription_duration", attribute_or_none(completion_response, "duration")) + ) + if measured_duration is not None and measured_duration > 0: + return measured_duration + provider_duration: Final = _get_transcription_usage_duration(completion_response) + if ( + provider_duration is not None + and provider_duration > 0 + and any(_bills_transcription_usage_duration(name, custom_llm_provider) for name in (pricing_model, model)) + ): + return provider_duration + if model is None or custom_llm_provider != LlmProviders.OPENAI.value: + return 0.0 transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI ) - return isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig) + return (provider_duration or 0.0) if isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig) else 0.0 OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"] @@ -604,6 +664,9 @@ def cost_per_token( decisions_overrides: Final = decisions_pricing_fields( litellm.model_cost.get(decisions_model_info["key"], {}) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # model cost map is dynamically loaded ) + request_price: Final = decisions_request_price(decisions_model_info, decisions_overrides) + if request_price is not None: + return request_price, 0.0 if decisions_overrides: return generic_cost_per_token( model=model, @@ -693,7 +756,7 @@ def cost_per_token( vertex_location=vertex_location, ) elif call_type == "atranscription" or call_type == "transcription": - if _transcription_usage_has_token_details(usage_block): + if _transcription_uses_token_pricing(model_without_prefix, custom_llm_provider, usage_block): return generic_cost_per_token( model=model, usage=usage_block, @@ -1275,6 +1338,23 @@ def _decisions_usage(usage: DecisionsUsage | OpenAIDecisionUsage) -> Usage: ) +def _get_transcription_usage_duration(completion_response: object) -> float | None: + response_fields: Final = as_str_mapping(completion_response) + usage_object: Final = ( + response_fields.get("usage") if response_fields is not None else attribute_or_none(completion_response, "usage") + ) + usage_fields: Final = as_str_mapping(usage_object) + usage_type: Final = ( + usage_fields.get("type") if usage_fields is not None else attribute_or_none(usage_object, "type") + ) + if usage_type != "duration": + return None + seconds: Final = ( + usage_fields.get("seconds") if usage_fields is not None else attribute_or_none(usage_object, "seconds") + ) + return normalized_audio_duration_seconds(seconds) + + def _extraction_usage(usage: ExtractionUsage) -> Usage: return Usage( prompt_tokens=usage.input_tokens, @@ -2000,20 +2080,8 @@ def completion_cost( prompt, ) elif call_type in _TRANSCRIPTION_CALL_TYPES: - # Check _hidden_params first (duration stored there to - # avoid polluting the response body), then fall back to - # the response attribute (for verbose_json responses that - # naturally include duration from the provider). - _hidden = getattr(completion_response, "_hidden_params", {}) or {} - audio_transcription_file_duration = _hidden.get( - "audio_transcription_duration", - getattr(completion_response, "duration", 0.0), - ) or ( - completion_response.usage.seconds - if isinstance(completion_response, TranscriptionResponse) - and isinstance(completion_response.usage, TranscriptionUsageDurationObject) - and _bills_provider_transcription_seconds(requested_model, custom_llm_provider) - else 0.0 + audio_transcription_file_duration = _transcription_duration_for_cost( + completion_response, model, requested_model, custom_llm_provider ) elif call_type in _RERANK_CALL_TYPES: if completion_response is not None and isinstance(completion_response, RerankResponse): @@ -2123,6 +2191,7 @@ def completion_cost( litellm_logging_obj=litellm_logging_obj, custom_pricing_model=selected_model if custom_pricing else None, base_pricing_model=(selected_model if base_model is not None and not custom_pricing else None), + deployment_models=(base_model, requested_model), ) elif call_type == _MCP_CALL_TYPE: from litellm.proxy.mcp_server.cost_calculator import ( @@ -3401,8 +3470,11 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): def create_logging_realtime_object( usage: Usage, results: OpenAIRealtimeStreamList ) -> LiteLLMRealtimeStreamLoggingObject: + translated_input_seconds, _ = _translation_session_seconds(results) return LiteLLMRealtimeStreamLoggingObject( - usage=usage, + usage=usage.model_copy(update={"seconds": translated_input_seconds}) + if translated_input_seconds > 0 + else usage, results=results, ) @@ -3487,6 +3559,7 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" +_TRANSLATION_CLOSED_EVENT_TYPE: Final = "session.closed" class _LiveResponsePayload(BaseModel): @@ -3709,6 +3782,7 @@ def handle_realtime_stream_cost_calculation( litellm_logging_obj: LitellmLoggingObject | None = None, custom_pricing_model: str | None = None, base_pricing_model: str | None = None, + deployment_models: Sequence[str | None] = (), ) -> float: """ Handles the cost calculation for realtime stream responses. @@ -3753,6 +3827,13 @@ def handle_realtime_stream_cost_calculation( if any(r.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE for r in results) else 0.0 ) + translation_cost: Final = handle_realtime_translation_cost_calculation( + results=results, + custom_llm_provider=custom_llm_provider, + litellm_model_name=litellm_model_name, + deployment_pricing_models=(custom_pricing_model, base_pricing_model, *deployment_models), + session_model=received_model, + ) live_lifecycle_event_types: Final = frozenset({"session.started", "session.usage.updated", "session.closed"}) live_session_started: Final = any(result["type"] == "session.started" for result in results) live_session_closed: Final = any(result["type"] == "session.closed" for result in results) @@ -3809,12 +3890,19 @@ def handle_realtime_stream_cost_calculation( input_cost: Final = input_cost_per_token + live_response_input_cost output_cost: Final = output_cost_per_token + live_response_output_cost total_cost: Final = ( - input_cost + output_cost + transcription_cost + live_voice_cost + audio_duration_cost + text_message_cost + input_cost + + output_cost + + transcription_cost + + translation_cost + + live_voice_cost + + audio_duration_cost + + text_message_cost ) additional_costs: Final = tuple( (cost_name, cost) for cost_name, cost in ( ("transcription_cost", transcription_cost), + ("translation_cost", translation_cost), ("live_voice_cost", live_voice_cost), ("audio_duration_cost", audio_duration_cost), ("text_message_cost", text_message_cost), @@ -3835,6 +3923,55 @@ def handle_realtime_stream_cost_calculation( return total_cost +def _translation_session_seconds(results: OpenAIRealtimeStreamList) -> tuple[float, float]: + usage_events: Final = tuple( + usage + for result in results + if (event := as_str_mapping(result)) is not None + and event.get("type") == _TRANSLATION_CLOSED_EVENT_TYPE + and (usage := as_str_mapping(event.get("usage"))) is not None + ) + input_seconds: Final = sum( + seconds + for usage in usage_events + if (seconds := normalized_audio_duration_seconds(usage.get("input_seconds"))) is not None + ) + output_seconds: Final = sum( + output_duration + for usage in usage_events + if (output_duration := normalized_audio_duration_seconds(usage.get("output_seconds"))) is not None + ) + return float(input_seconds), float(output_seconds) + + +def _translation_per_second_rates( + models: Sequence[str | None], custom_llm_provider: str +) -> tuple[float | None, float | None]: + model_infos: Final = tuple( + _lookup_model_info_or_none(model, custom_llm_provider) for model in models if model is not None + ) + return ( + _first_declared_transcription_rate(("input_cost_per_second",), model_infos), + _first_declared_transcription_rate(("output_cost_per_second",), model_infos), + ) + + +def handle_realtime_translation_cost_calculation( + results: OpenAIRealtimeStreamList, + custom_llm_provider: str, + litellm_model_name: str, + deployment_pricing_models: Sequence[str | None] = (), + session_model: str | None = None, +) -> float: + input_seconds, output_seconds = _translation_session_seconds(results) + if input_seconds <= 0 and output_seconds <= 0: + return 0.0 + input_rate, output_rate = _translation_per_second_rates( + (*deployment_pricing_models, litellm_model_name, session_model), custom_llm_provider + ) + return input_seconds * (input_rate or 0.0) + output_seconds * (output_rate or 0.0) + + def handle_realtime_transcription_cost_calculation( results: OpenAIRealtimeStreamList, custom_llm_provider: str, @@ -3850,57 +3987,77 @@ def handle_realtime_transcription_cost_calculation( - {"type": "duration", "seconds": } → priced via input_cost_per_second - {"type": "tokens", "input_tokens": ...} → priced via input/audio token cost """ - completed_events: Final = [ - cast(dict, result) for result in results if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE - ] - if not completed_events: + configured_models: Final = tuple(accumulate(map(_session_transcription_model, results), _newer_model)) + clips: Final = tuple( + (as_str_mapping(result.get("usage")) or {}, configured) + for result, configured in zip(results, configured_models, strict=True) + if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE + ) + if not clips: return 0.0 - model_name: Final = _get_transcription_model_name_from_results(results) or litellm_model_name - model_info: Final = _cost_map_model_info(model_name, custom_llm_provider) + fallback_model: Final = _get_transcription_model_name_from_results(results) or litellm_model_name + clip_models: Final = tuple((usage, configured or fallback_model) for usage, configured in clips) + model_infos: Final = MappingProxyType( + {model: _cost_map_model_info(model, custom_llm_provider) for model in frozenset(m for _, m in clip_models)} + ) override_info: Final = ( _cost_map_model_info(custom_pricing_model, custom_llm_provider) if custom_pricing_model is not None else None ) + return sum( + (_transcription_usage_cost(dict(usage), model_infos[model], override_info) for usage, model in clip_models), + 0.0, + ) - total_cost = 0.0 - for event in completed_events: - usage = event.get("usage") or {} - total_cost += _transcription_usage_cost(usage, model_info, override_info) - return total_cost + +_TRANSCRIPTION_SESSION_EVENT_TYPES: Final = frozenset( + {"transcription_session.created", "transcription_session.updated", "session.created", "session.updated"} +) + + +def _str_field(mapping: Mapping[str, object] | None, key: str) -> str | None: + value: Final = None if mapping is None else mapping.get(key) + return value if isinstance(value, str) and value else None + + +def _configured_transcription_model(session: Mapping[str, object]) -> str | None: + audio_input: Final = as_str_mapping((as_str_mapping(session.get("audio")) or {}).get("input")) + nested: Final = _str_field(as_str_mapping((audio_input or {}).get("transcription")), "model") + return nested or _str_field(as_str_mapping(session.get("input_audio_transcription")), "model") + + +def _session_of(result: Mapping[str, object]) -> Mapping[str, object] | None: + return as_str_mapping(result.get("session")) if result.get("type") in _TRANSCRIPTION_SESSION_EVENT_TYPES else None + + +def _session_transcription_model(result: Mapping[str, object]) -> str | None: + session: Final = _session_of(result) + return None if session is None else _configured_transcription_model(session) + + +def _newer_model(current: str | None, newer: str | None) -> str | None: + return newer or current def _get_transcription_model_name_from_results( results: OpenAIRealtimeStreamList, ) -> str | None: - """Resolve the ASR model from a transcription_session.* / session.* event.""" - for result in results: - if result.get("type") in ( - "transcription_session.created", - "transcription_session.updated", - "session.created", - "session.updated", - ): - session = cast(dict, result).get("session", {}) or {} - transcription = ((session.get("audio", {}) or {}).get("input", {}) or {}).get( - "transcription", {} - ) or session.get("input_audio_transcription", {}) - model = (transcription or {}).get("model") or session.get("model") - if model: - return model - return None + return next( + (model for result in results if (model := _str_field(_session_of(result), "model"))), + None, + ) -_TranscriptionRateKey: TypeAlias = Literal[ - "input_cost_per_second", "input_cost_per_audio_token", "input_cost_per_token", "output_cost_per_token" -] - - -def _model_info_rate(info: ModelInfo, key: _TranscriptionRateKey) -> float | None: +def _model_info_rate(info: ModelInfoBase, key: _TranscriptionRateKey) -> float | None: match key: case "input_cost_per_second": return info.get("input_cost_per_second") + case "output_cost_per_second": + return info.get("output_cost_per_second") case "input_cost_per_audio_token": return info.get("input_cost_per_audio_token") + case "output_cost_per_audio_token": + return info.get("output_cost_per_audio_token") case "input_cost_per_token": return info.get("input_cost_per_token") case "output_cost_per_token": @@ -3909,7 +4066,7 @@ def _model_info_rate(info: ModelInfo, key: _TranscriptionRateKey) -> float | Non assert_never(key) -def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[_TranscriptionRateKey, ...]) -> float | None: +def _declared_transcription_rate(info: ModelInfoBase | None, keys: tuple[_TranscriptionRateKey, ...]) -> float | None: """First of ``keys`` this entry prices, read off the raw ``litellm.model_cost`` entry because ``get_model_info`` synthesizes zero token rates for entries that omit them.""" if info is None: @@ -3927,15 +4084,19 @@ def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[_Transcript ) -def _transcription_rate( - keys: tuple[_TranscriptionRateKey, ...], override: ModelInfo | None, base: ModelInfo | None -) -> float: - rates: Final = (_declared_transcription_rate(info, keys) for info in (override, base)) - return next((rate for rate in rates if rate is not None), 0.0) +def _first_declared_transcription_rate( + keys: tuple[_TranscriptionRateKey, ...], infos: Sequence[ModelInfoBase | None] +) -> float | None: + rates: Final = (_declared_transcription_rate(info, keys) for info in infos) + return next((rate for rate in rates if rate is not None), None) + + +def _transcription_rate(keys: tuple[_TranscriptionRateKey, ...], *infos: ModelInfoBase | None) -> float: + return _first_declared_transcription_rate(keys, infos) or 0.0 def _transcription_usage_cost( - usage: dict, + usage: Mapping[str, object], model_info: ModelInfo | None, override_info: ModelInfo | None = None, ) -> float: @@ -3944,13 +4105,13 @@ def _transcription_usage_cost( usage_type: Final = usage.get("type") if usage_type == "duration": - seconds: Final = usage.get("seconds") or 0.0 + seconds: Final = normalized_audio_duration_seconds(usage.get("seconds")) or 0.0 return float(seconds) * _transcription_rate(("input_cost_per_second",), override_info, model_info) if usage_type == "tokens": - input_token_details: Final = usage.get("input_token_details") or {} - audio_tokens: Final = input_token_details.get("audio_tokens") or 0 - text_tokens: Final = input_token_details.get("text_tokens") or 0 - output_tokens: Final = usage.get("output_tokens") or 0 + input_token_details: Final = as_str_mapping(usage.get("input_token_details")) or {} + audio_tokens: Final = normalized_audio_duration_seconds(input_token_details.get("audio_tokens")) or 0.0 + text_tokens: Final = normalized_audio_duration_seconds(input_token_details.get("text_tokens")) or 0.0 + output_tokens: Final = normalized_audio_duration_seconds(usage.get("output_tokens")) or 0.0 audio_cost: Final = float(audio_tokens) * _transcription_rate( ("input_cost_per_audio_token", "input_cost_per_token"), override_info, model_info ) diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 871fa8227ea..9c0e51e14ac 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -13,6 +13,7 @@ from typing_extensions import assert_never import litellm from litellm.cost_calculator import cost_map_key_prices_anything from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params +from litellm.litellm_core_utils.decisions_pricing import QUERY_PRICED_DECISIONS_PROVIDERS from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.image_handling import ( @@ -555,7 +556,7 @@ def prepare_call( resolved_api_base, static_api_key = provider_config.resolve_connection( dynamic_api_base or call_api_base, dynamic_api_key or api_key ) - except ValueError as error: + except (ValueError, httpx.InvalidURL) as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error if resolved_api_base is None: raise litellm.BadRequestError( @@ -613,6 +614,9 @@ def prepare_call( litellm_params={ "litellm_call_id": kwargs.get("litellm_call_id"), "api_base": provider_config.get_complete_url(resolved_api_base, canonical_model), + "input_cost_per_query": ( + kwargs.get("input_cost_per_query") if provider in QUERY_PRICED_DECISIONS_PROVIDERS else None + ), }, custom_llm_provider=provider, ) @@ -787,7 +791,9 @@ def _formatted_response( case DecisionsRequestBody(): return ir_to_systemone_response(caller_response, caller_request) case OpenAIDecisionRequestBody(): - return ir_to_openai_response(caller_response, caller_request, call.requested_model) + return ir_to_openai_response( + caller_response, caller_request, call.requested_model, call.provider_config.openai_evidence_fields + ) case _: assert_never(call.request) diff --git a/litellm/images/main.py b/litellm/images/main.py index 165543aadba..27917cf9289 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -28,8 +28,7 @@ from litellm.litellm_core_utils.mock_functions import mock_image_generation from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.llms.custom_llm import CustomLLM -from litellm.utils import exception_type, filter_out_litellm_params, get_litellm_params +from litellm.utils import exception_type, filter_out_litellm_params, get_litellm_params, header_auth_params #################### Initialize provider clients #################### llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() @@ -608,10 +607,7 @@ def image_generation( ) elif custom_llm_provider in litellm._custom_providers: # Assume custom LLM provider # Get the Custom Handler - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) @@ -915,10 +911,7 @@ def image_edit( # Check for custom provider if custom_llm_provider in litellm._custom_providers: - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) @@ -995,10 +988,15 @@ def image_edit( or custom_llm_provider == "azure" or custom_llm_provider in litellm.openai_compatible_providers ): + dropped_form_params: Final = header_auth_params(custom_llm_provider) image_edit_request_params.update( flatten_form_field_values( - non_default_params, - extra_body if isinstance(extra_body, dict) else None, + {key: value for key, value in non_default_params.items() if key not in dropped_form_params}, + ( + {key: value for key, value in extra_body.items() if key not in dropped_form_params} + if isinstance(extra_body, dict) + else None + ), ) ) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index c05407b7972..4e473c2f3f0 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -9,24 +9,29 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and """ +import bisect import copy import itertools import os import re from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Sequence +from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar, cast from urllib.parse import urlparse -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm._logging import verbose_logger +from litellm.exceptions import UnsupportedParamsError from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.integrations.prompt_management_base import PromptManagementClient -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.prompt_templates.common_utils import ( + PROMPT_CACHE_BREAKPOINT_DROP_PARAMS_HINT, cache_control_as_prompt_cache_breakpoint, + prompt_cache_breakpoint_as_cache_control, with_prompt_cache_breakpoint, + without_prompt_cache_breakpoints, ) from litellm.litellm_core_utils.prompt_templates.factory import ( tool_call_ids_answered_by_client, @@ -73,7 +78,7 @@ MAX_CACHE_CONTROL_BLOCKS: Final = 4 CACHE_BREAKPOINT_KEYS: Final = ("cache_control", "prompt_cache_breakpoint") OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION: Final = (5, 6) -_GPT_VERSION_PATTERN: Final = re.compile(r"^gpt-(\d+)(?:\.(\d+))?") +_GPT_VERSION_PATTERN: Final = re.compile(r"^gpt-(\d+)(?:\.(\d+)|-(0|[1-9]\d?)(?![\da-z]))?") _FINE_TUNED_BASE_PATTERN: Final = re.compile(r"^ft:([^:]+)") _O_SERIES_PATTERN: Final = re.compile(r"^o\d+(?:-|$)") OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset( @@ -85,8 +90,40 @@ _OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _CONTENT_PART_ADAPTER: Final = TypeAdapter(dict[str, object]) _ATTRIBUTE_VALUE: Final = TypeAdapter(object) +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues +CostMap: TypeAlias = Mapping[str, Mapping[str, object]] + + +def routed_model_info(value: object) -> Mapping[str, object] | None: + try: + return _MODEL_INFO_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _live_cost_map_unless(model_cost: CostMap | None) -> CostMap: + import litellm + + if model_cost is not None: + return model_cost + return cast(CostMap, litellm.model_cost) # cast-ok: the live map holds catalog entries + + +def declared_prompt_cache_breakpoint(model_info: Mapping[str, object] | None) -> bool | None: + declared: Final = None if model_info is None else model_info.get("supports_prompt_cache_breakpoint") + return declared if isinstance(declared, bool) else None + + +def routed_cost_map() -> CostMap: + from litellm.utils import catalog_with_runtime_registrations + + return catalog_with_runtime_registrations() + + +def _catalog_unless_unrouted(model_info: Mapping[str, object] | None) -> CostMap | None: + return None if model_info is None else routed_cost_map() def _validated_object_mapping(value: object) -> dict[object, object] | None: @@ -295,6 +332,12 @@ def _part_with_cache_control_renamed(part: object) -> object: return cache_control_as_prompt_cache_breakpoint(_CONTENT_PART_ADAPTER.validate_python(part)) +def _part_with_prompt_cache_breakpoint_renamed(part: object, drop_params: bool) -> object: + if not _carries_prompt_cache_breakpoint(part): + return part + return prompt_cache_breakpoint_as_cache_control(_CONTENT_PART_ADAPTER.validate_python(part), drop_params) + + def configured_injection_points(value: object) -> Sequence[CacheControlInjectionPoint]: if not isinstance(value, (list, tuple)): return () @@ -314,19 +357,27 @@ def _served_provider_for_model(model: str, api_base: object = None) -> str | Non return provider -def supports_openai_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None = None) -> bool: - return openai_prompt_cache_breakpoint_support(model, custom_llm_provider) is True +def supports_openai_prompt_cache_breakpoint( + model: str, custom_llm_provider: str | None = None, model_info: Mapping[str, object] | None = None +) -> bool: + return openai_prompt_cache_breakpoint_support(model, custom_llm_provider, model_info) is True -def openai_prompt_cache_breakpoint_support(model: str, custom_llm_provider: str | None = None) -> bool | None: +def openai_prompt_cache_breakpoint_support( + model: str, custom_llm_provider: str | None = None, model_info: Mapping[str, object] | None = None +) -> bool | None: + declared: Final = declared_prompt_cache_breakpoint(model_info) + if declared is not None: + return declared + model_cost: Final = _catalog_unless_unrouted(model_info) hosted_flag: Final = ( None if custom_llm_provider is None - else _hosted_openai_dialect_flag(model, custom_llm_provider, _served_provider_for_model) + else _hosted_openai_dialect_flag(model, custom_llm_provider, _served_provider_for_model, model_cost) ) if hosted_flag is not None: return hosted_flag - model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model) + model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model, model_cost) if model_map_flag is not None: return model_map_flag fine_tuned_base: Final = _FINE_TUNED_BASE_PATTERN.match(model.rsplit("/", 1)[-1]) @@ -335,6 +386,35 @@ def openai_prompt_cache_breakpoint_support(model: str, custom_llm_provider: str return _named_model_breakpoint_support(model) +class _DeclaredBaseModels(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + + base_model: object = None + model_info: object = None + + def names(self) -> tuple[str, ...]: + model_info: Final = _validated_object_mapping(self.model_info) or {} + candidates: Final = (self.base_model, model_info.get("base_model")) + return tuple(name for name in candidates if isinstance(name, str) and name) + + +def deployment_model_info(litellm_params: object) -> Mapping[str, object] | None: + return routed_model_info(_DeclaredBaseModels.model_validate(litellm_params).model_info) + + +def deployment_prompt_cache_breakpoint_support( + model: str, custom_llm_provider: str, litellm_params: object +) -> bool | None: + """``openai_prompt_cache_breakpoint_support`` for a deployment, read from its declared base model first.""" + declared: Final = _DeclaredBaseModels.model_validate(litellm_params) + model_info: Final = routed_model_info(declared.model_info) + supports: Final = ( + openai_prompt_cache_breakpoint_support(name, custom_llm_provider, model_info) + for name in dict.fromkeys((*declared.names(), model)) + ) + return next((support for support in supports if support is not None), None) + + def _named_model_breakpoint_support(model: str) -> bool | None: bare_model: Final = model.rsplit("/", 1)[-1].lower() if _O_SERIES_PATTERN.match(bare_model) is not None: @@ -342,15 +422,14 @@ def _named_model_breakpoint_support(model: str) -> bool | None: version_match: Final = _GPT_VERSION_PATTERN.match(bare_model) if version_match is None: return None - version: Final = (int(version_match.group(1)), int(version_match.group(2) or 0)) + version: Final = (int(version_match.group(1)), int(version_match.group(2) or version_match.group(3) or 0)) return version >= OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION -def _model_map_prompt_cache_breakpoint_flag(model: str) -> bool | None: - import litellm - - entries: Final = (litellm.model_cost.get(key) for key in (model, model.rsplit("/", 1)[-1])) - flags: Final = (entry.get("supports_prompt_cache_breakpoint") for entry in entries if isinstance(entry, dict)) +def _model_map_prompt_cache_breakpoint_flag(model: str, model_cost: CostMap | None = None) -> bool | None: + cost_map: Final = _live_cost_map_unless(model_cost) + entries: Final = (cost_map.get(key) for key in (model, model.rsplit("/", 1)[-1])) + flags: Final = (entry.get("supports_prompt_cache_breakpoint") for entry in entries if isinstance(entry, Mapping)) return next((flag is True for flag in flags if flag is not None), None) @@ -366,7 +445,7 @@ def _hosted_openai_dialect_flag( model: str, custom_llm_provider: str | None, resolve_provider: Callable[[str], str | None], - model_cost: Mapping[str, Mapping[str, object]] | None = None, + model_cost: CostMap | None = None, ) -> bool | None: """ Explicit opt-in for an OpenAI-shaped deployment served by another provider. @@ -383,13 +462,7 @@ def _hosted_openai_dialect_flag( row of the same name, so an exact entry only speaks for a deployment when it is keyed for that deployment's provider. """ - import litellm - - cost_map: Final = ( - cast(Mapping[str, Mapping[str, object]], litellm.model_cost) # cast-ok: the live map holds catalog entries - if model_cost is None - else model_cost - ) + cost_map: Final = _live_cost_map_unless(model_cost) exact_entry: Final = cost_map.get(model) exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, Mapping) else None if custom_llm_provider is None and (exact_provider == "openai" or ("/" not in model and exact_provider is None)): @@ -399,10 +472,16 @@ def _hosted_openai_dialect_flag( return None if isinstance(exact_entry, Mapping) and exact_provider == provider: return _hosted_entry_flag(exact_entry, lambda: provider) + entries: Final = (cost_map.get(candidate) for candidate in _provider_cost_map_keys(model, provider)) + flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, Mapping)) + return next((flag for flag in flags if flag is not None), None) + + +def _provider_cost_map_keys(model: str, provider: str) -> tuple[str, ...]: from litellm.utils import get_potential_model_names names: Final = get_potential_model_names(model, provider) - candidates: Final = ( + return ( names["combined_model_name"], names["region_free_combined_model_name"], names["split_model"], @@ -410,9 +489,60 @@ def _hosted_openai_dialect_flag( names["stripped_model_name"], names["provider_prefixed_model_name"], ) - entries: Final = (cost_map.get(candidate) for candidate in candidates) - flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, Mapping)) - return next((flag for flag in flags if flag is not None), None) + + +def _cost_map_lists_model_for_provider(model: str, provider: str, model_cost: CostMap) -> bool: + entries: Final = (model_cost.get(key) for key in (model, *_provider_cost_map_keys(model, provider))) + return any(isinstance(entry, Mapping) and entry.get("litellm_provider") == provider for entry in entries) + + +_PROVIDERS_READING_THE_FLAG_AS_CACHE_POINT: Final = frozenset(("bedrock", "bedrock_converse")) +_PROVIDERS_HOSTING_OPENAIS_API: Final = frozenset(("azure", "azure_ai", "bedrock_mantle")) + + +def _declared_hosted_flag( + declared: bool, model: str, provider: str | None, model_cost: CostMap, deployment_provider: object +) -> bool | None: + """ + A deployment's own ``supports_prompt_cache_breakpoint`` picks OpenAI's breakpoint only where the + request goes out OpenAI-shaped. Bedrock reads the same flag on every route as whether Converse takes + a ``cachePoint``, and a provider building Anthropic messages takes ``cache_control``, so both keep + the Anthropic mark. A gateway or self-hosted server (OpenRouter, a LiteLLM proxy, vLLM) may pass + the request on to Anthropic, so there the flag only picks OpenAI's breakpoint for a model the cost + map, or the deployment's own ``litellm_provider``, lists under that provider. + """ + if provider is None or provider == "openai": + return None + return ( + declared + and provider not in _PROVIDERS_READING_THE_FLAG_AS_CACHE_POINT + and ( + provider in _PROVIDERS_HOSTING_OPENAIS_API + or deployment_provider == provider + or _cost_map_lists_model_for_provider(model, provider, model_cost) + ) + and _sends_openai_shaped_chat(model, provider) + ) + + +@lru_cache(maxsize=512) +def _sends_openai_shaped_chat(model: str, custom_llm_provider: str) -> bool: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata + from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig + from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + from litellm.llms.openai.openai import OpenAIConfig + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + try: + served_model, provider = resolve_model_and_provider_for_metadata(model, custom_llm_provider) + chat_config: Final = ProviderConfigManager.get_provider_chat_config( + model=served_model, provider=LlmProviders(provider) + ) + except (BadRequestError, ValueError): + return False + return isinstance(chat_config, (OpenAIGPTConfig, OpenAIConfig, AzureOpenAIConfig)) def targets_openai_api(api_base: object) -> bool: @@ -432,6 +562,10 @@ def _carries_cache_breakpoint(block: object) -> bool: return any(_attribute_or_key(block, key) is not None for key in CACHE_BREAKPOINT_KEYS) +def _carries_cache_control(block: object) -> bool: + return _attribute_or_key(block, "cache_control") is not None + + def _attribute_or_key(value: object, key: str) -> object | None: if hasattr(value, key): return _ATTRIBUTE_VALUE.validate_python(getattr(value, key)) @@ -511,6 +645,12 @@ def _builds_anthropic_messages(model: str, custom_llm_provider: str | None) -> b return isinstance(chat_config, AnthropicConfig) and not isinstance(chat_config, OpenAIGPTConfig) +def _translates_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None) -> bool: + return (custom_llm_provider or _served_provider_for_model(model)) == "bedrock" or _builds_anthropic_messages( + model, custom_llm_provider + ) + + def _custom_llm_provider(non_default_params: Mapping[str, object], point: Mapping[str, object]) -> str | None: provider: Final = non_default_params.get("custom_llm_provider") or point.get(CUSTOM_LLM_PROVIDER_STAMP) return provider if isinstance(provider, str) else None @@ -554,10 +694,79 @@ def _accepts_prompt_cache_breakpoint(block: object) -> bool: return isinstance(block, dict) and block.get("type") in OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES +def _carries_prompt_cache_breakpoint(block: object) -> bool: + return _accepts_prompt_cache_breakpoint(block) and _attribute_or_key(block, "prompt_cache_breakpoint") is not None + + def _may_carry_prompt_cache_breakpoint(message: Mapping[str, object]) -> bool: return message.get("role") != "assistant" +def _tool_result_carries_prompt_cache_breakpoint(message: Mapping[str, object]) -> bool: + return message.get("role") == "tool" and any( + _carries_prompt_cache_breakpoint(part) for part in _as_object_list(message.get("content")) or () + ) + + +def _part_without_prompt_cache_breakpoint(part: object) -> object: + if not _carries_prompt_cache_breakpoint(part): + return part + return { + key: value + for key, value in _CONTENT_PART_ADAPTER.validate_python(part).items() + if key != "prompt_cache_breakpoint" + } + + +def _tool_result_without_prompt_cache_breakpoints(message: AllMessageValues) -> AllMessageValues: + if not _tool_result_carries_prompt_cache_breakpoint(message): + return message + parts: Final = _as_object_list(_attribute_or_key(message, "content")) or [] + return cast( # cast-ok: the same tool message with its part marks removed + AllMessageValues, {**message, "content": [_part_without_prompt_cache_breakpoint(part) for part in parts]} + ) + + +def _created_cache_control(original: object, translated: object) -> bool: + return ( + _attribute_or_key(original, "cache_control") is None + and _attribute_or_key(translated, "cache_control") is not None + ) + + +def _created_cache_controls(original: object, translated: object) -> tuple[bool, ...]: + if translated is original: + return () + original_parts: Final = _as_object_list(_attribute_or_key(original, "content")) or [] + translated_parts: Final = _as_object_list(_attribute_or_key(translated, "content")) or [] + return tuple(_created_cache_control(before, after) for before, after in zip(original_parts, translated_parts)) + + +def _part_without_cache_control(part: object) -> object: + return {key: value for key, value in _CONTENT_PART_ADAPTER.validate_python(part).items() if key != "cache_control"} + + +def _message_dropping_created_cache_controls_between( + message: AllMessageValues, created: Sequence[bool], marks_before: int, first_dropped: int, first_kept_after: int +) -> AllMessageValues: + if not any(created): + return message + parts: Final = _as_object_list(_attribute_or_key(message, "content")) or [] + ordinals: Final = itertools.accumulate((int(flag) for flag in created), initial=marks_before) + return cast( # cast-ok: the same message with its middle translated marks removed + AllMessageValues, + { + **message, + "content": [ + _part_without_cache_control(part) + if is_created and first_dropped <= ordinal < first_kept_after + else part + for part, is_created, ordinal in zip(parts, created, ordinals) + ], + }, + ) + + # Set by a caller whose message list is not the one that goes upstream -- today the # Responses API layer, whose `instructions` only becomes a system message further down. # Tells this hook to hand role-targeted points to the pass holding the final messages @@ -643,9 +852,6 @@ class AnthropicCacheControlHook(CustomPromptManagement): if not injection_points: return model, messages, non_default_params - # Create a deep copy of messages to avoid modifying the original list - processed_messages = copy.deepcopy(messages) - message_points: Final = tuple( cast(CacheControlMessageInjectionPoint, point) for point in injection_points @@ -653,18 +859,26 @@ class AnthropicCacheControlHook(CustomPromptManagement): ) remaining_points: Final = tuple(point for point in injection_points if point.get("location") != "message") - custom_llm_provider: Final = _custom_llm_provider(non_default_params, injection_points[0]) + request_params: Final = cast(Mapping[str, object], non_default_params) # cast-ok: base hook's bare dict + custom_llm_provider: Final = _custom_llm_provider(request_params, injection_points[0]) stamped_dialect: Final = injection_points[0].get("_litellm_openai_dialect") + model_info: Final = routed_model_info(request_params.get("model_info")) openai_dialect: Final = ( stamped_dialect if isinstance(stamped_dialect, bool) else AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( model, custom_llm_provider, - non_default_params.get("api_base") or non_default_params.get("base_url"), - non_default_params.get("prompt_cache_options"), + request_params.get("api_base") or request_params.get("base_url"), + request_params.get("prompt_cache_options"), + model_info, ) ) + processed_messages = copy.deepcopy( + [without_prompt_cache_breakpoints(message) for message in messages] + if not openai_dialect and message_points and _translates_prompt_cache_breakpoint(model, custom_llm_provider) + else messages + ) # A provisional message list defers every role-targeted point to the pass holding # the final one: a role with no message here may have one there, and settling all # of them in one pass is what lets config order decide the shared breakpoint @@ -733,32 +947,35 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider: str | None, api_base: object = None, prompt_cache_options: object = None, + model_info: Mapping[str, object] | None = None, ) -> bool: if model is None: return False - hosted_flag: Final = _hosted_openai_dialect_flag( - model, custom_llm_provider, lambda name: AnthropicCacheControlHook._resolve_provider(name, api_base) + declared: Final = declared_prompt_cache_breakpoint(model_info) + hosted_flag: Final = ( + _hosted_openai_dialect_flag( + model, + custom_llm_provider, + lambda name: AnthropicCacheControlHook._resolve_provider(name, api_base), + _catalog_unless_unrouted(model_info), + ) + if declared is None + else _declared_hosted_flag( + declared, + model, + custom_llm_provider or AnthropicCacheControlHook._resolve_provider(model, api_base), + routed_cost_map(), + None if model_info is None else model_info.get("litellm_provider"), + ) ) if hosted_flag is not None: - return hosted_flag - if not supports_openai_prompt_cache_breakpoint(model): + return hosted_flag and not _translates_prompt_cache_breakpoint(model, custom_llm_provider) + if not supports_openai_prompt_cache_breakpoint(model, model_info=model_info): return False if (custom_llm_provider or AnthropicCacheControlHook._resolve_provider(model)) != "openai": return False return prompt_cache_options is not None or targets_openai_api(api_base) - @staticmethod - def catalog_targets_openai_prompt_cache_breakpoint(model: str, custom_llm_provider: str) -> bool: - return ( - _hosted_openai_dialect_flag( - model, - custom_llm_provider, - AnthropicCacheControlHook._resolve_provider, - GetModelCostMap.loaded_model_cost_map(), - ) - is True - ) - @staticmethod def _resolve_provider(model: str, api_base: object = None) -> str | None: return _served_provider_for_model(model, api_base) @@ -805,6 +1022,19 @@ class AnthropicCacheControlHook(CustomPromptManagement): ) return int(wire_cache_control is not None) + tool_blocks + envelope_blocks + @staticmethod + def tools_keeping_cache_control_on_chat(tools: object) -> tuple[object, ...]: + return tuple( + tool for tool in _as_object_list(tools) or () if not _chat_transform_drops_tool_cache_control(tool) + ) + + @staticmethod + def count_chat_external_cache_breakpoints(optional_params: object) -> int: + params: Final = _validated_object_mapping(optional_params) or {} + return AnthropicCacheControlHook.count_external_cache_breakpoints( + _as_object_list(params.get("tools")), params.get("cache_control") + ) + @staticmethod def count_external_cache_breakpoints_on_messages_route( tools: Iterable[object] | None, cache_control: object, request_kwargs: object @@ -960,6 +1190,16 @@ class AnthropicCacheControlHook(CustomPromptManagement): ) return message_count + content_count + tool_call_count + @staticmethod + def count_converse_cache_points(message: AllMessageValues) -> int: + content: Final = _as_object_list(_attribute_or_key(message, "content")) + tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls")) or [] + marked_parts: Final = sum(1 for block in content or () if _carries_cache_control(block)) + marked_tool_calls: Final = sum(1 for tool_call in tool_calls if _carries_cache_control(tool_call)) + if message.get("role") == "tool": + return int(_carries_cache_control(message) or marked_parts > 0) + return (int(_carries_cache_control(message)) if content is None else marked_parts) + marked_tool_calls + @staticmethod def _message_has_cache_control(message: AllMessageValues, answered_tool_ids: Container[str]) -> bool: """Return True if the message already carries any cache_control.""" @@ -1018,6 +1258,92 @@ class AnthropicCacheControlHook(CustomPromptManagement): return renamed return AnthropicCacheControlHook._insert_prompt_cache_breakpoint_in_message(renamed) + @staticmethod + def prompt_cache_breakpoints_as_cache_control(message: AllMessageValues, drop_params: bool) -> AllMessageValues: + """Every OpenAI ``prompt_cache_breakpoint`` on a part OpenAI accepts one on moved to Anthropic's + ``cache_control``. Marks with no OpenAI slot are left for the caller to strip.""" + if not _may_carry_prompt_cache_breakpoint(message): + return message + parts: Final = _as_object_list(_attribute_or_key(message, "content")) + if parts is None or not any(_carries_prompt_cache_breakpoint(part) for part in parts): + return message + return cast( # cast-ok: the same message with its part marks renamed + AllMessageValues, + {**message, "content": [_part_with_prompt_cache_breakpoint_renamed(part, drop_params) for part in parts]}, + ) + + @staticmethod + def prompt_cache_breakpoints_dropped_from_tool_results( + messages: Sequence[AllMessageValues], drop_params: bool, model: str + ) -> tuple[AllMessageValues, ...]: + if not any(_tool_result_carries_prompt_cache_breakpoint(message) for message in messages): + return tuple(messages) + if not drop_params: + raise UnsupportedParamsError( + message=( + f"{model} takes no cache breakpoint on a tool result, so prompt_cache_breakpoint on a tool " + f"message has no equivalent. Remove it, or {PROMPT_CACHE_BREAKPOINT_DROP_PARAMS_HINT} to drop it." + ), + status_code=400, + ) + verbose_logger.debug("Dropping prompt_cache_breakpoint from tool results for %s under drop_params", model) + return tuple(_tool_result_without_prompt_cache_breakpoints(message) for message in messages) + + @staticmethod + def prompt_cache_breakpoints_as_capped_cache_control( + messages: Sequence[AllMessageValues], + drop_params: bool, + sent_breakpoints: Callable[[AllMessageValues], int], + external_breakpoints: Callable[[], int], + ) -> tuple[AllMessageValues, ...]: + """``prompt_cache_breakpoints_as_cache_control`` over the whole request, held to the provider's + breakpoint cap, counting each message by the breakpoints ``sent_breakpoints`` says its provider sends. + Under ``drop_params`` the first and last translated marks fill the slots the client's own + ``cache_control`` marks leave, the first half taking the odd slot, so the stable prefix and the newest + turns stay cached, and without it a request over the cap is rejected.""" + translated: Final = tuple( + AnthropicCacheControlHook.prompt_cache_breakpoints_as_cache_control(message, drop_params) + for message in messages + ) + if all(after is before for before, after in zip(messages, translated)): + return translated + created: Final = tuple(_created_cache_controls(before, after) for before, after in zip(messages, translated)) + created_marks: Final = sum(sum(flags) for flags in created) + if created_marks == 0: + return translated + external: Final = external_breakpoints() + breakpoints: Final = sum(map(sent_breakpoints, translated)) + external + excess: Final = breakpoints - MAX_CACHE_CONTROL_BLOCKS + if excess <= 0: + return translated + if not drop_params: + raise UnsupportedParamsError( + message=( + f"This model takes at most {MAX_CACHE_CONTROL_BLOCKS} cache breakpoints per request, and this " + f"request has {breakpoints} once its prompt_cache_breakpoint marks are translated. Remove {excess} " + f"of them, or {PROMPT_CACHE_BREAKPOINT_DROP_PARAMS_HINT} to keep the first and last of them." + ), + status_code=400, + ) + verbose_logger.debug("Dropping the middle prompt_cache_breakpoint marks past the cap under drop_params") + marks_before: Final = tuple(itertools.accumulate((sum(flags) for flags in created), initial=0)) + + def capped_request(kept_marks: int) -> tuple[AllMessageValues, ...]: + first_dropped: Final = (kept_marks + 1) // 2 + first_kept_after: Final = created_marks - kept_marks // 2 + return tuple( + _message_dropping_created_cache_controls_between( + message, flags, before, first_dropped, first_kept_after + ) + for message, flags, before in zip(translated, created, marks_before) + ) + + def over_cap(kept_marks: int) -> bool: + return sum(map(sent_breakpoints, capped_request(kept_marks))) + external > MAX_CACHE_CONTROL_BLOCKS + + first_over_cap: Final = bisect.bisect_left(range(max(0, created_marks - excess) + 1), True, key=over_cap) + return capped_request(max(0, first_over_cap - 1)) + @staticmethod def _insert_prompt_cache_breakpoint_in_message(message: AllMessageValues) -> AllMessageValues: if not _may_carry_prompt_cache_breakpoint(message): @@ -1157,6 +1483,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider: str | None, api_base: object, prompt_cache_options: object, + model_info: Mapping[str, object] | None = None, ) -> Sequence[Mapping[str, object]]: """Carry onto the points what the prompt-management hook never receives. @@ -1168,7 +1495,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): across requests. """ with_dialect: Final = AnthropicCacheControlHook._stamped_with_dialect( - points, model, custom_llm_provider, api_base, prompt_cache_options + points, model, custom_llm_provider, api_base, prompt_cache_options, model_info ) stamped_earlier: Final = any(EXTERNAL_BREAKPOINTS_STAMP in point for point in with_dialect) if external_breakpoints == 0 and not stamped_earlier: @@ -1182,14 +1509,17 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider: str | None, api_base: object, prompt_cache_options: object, + model_info: Mapping[str, object] | None = None, ) -> Sequence[Mapping[str, object]]: - if not AnthropicCacheControlHook._may_target_openai_prompt_cache_breakpoint(model, custom_llm_provider): + if model_info is None and not AnthropicCacheControlHook._may_target_openai_prompt_cache_breakpoint( + model, custom_llm_provider + ): return points return AnthropicCacheControlHook._stamped( points, "_litellm_openai_dialect", AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( - model, custom_llm_provider, api_base, prompt_cache_options + model, custom_llm_provider, api_base, prompt_cache_options, model_info ), ) @@ -1248,6 +1578,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): cache_control: object = None, request_kwargs: object = None, on_messages_route: bool = False, + model_info: Mapping[str, object] | None = None, ) -> list[CacheControlInjectionPoint]: """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on. @@ -1265,7 +1596,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): if litellm.enable_anthropic_prompt_caching is not True and enable_prompt_caching is not True: return [] - if not supports_anthropic_cache_control(model, custom_llm_provider): + if not supports_anthropic_cache_control(model, custom_llm_provider, model_info): return [] if AnthropicCacheControlHook._request_has_cache_control( @@ -1343,6 +1674,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): tools: list | None = None, enable_prompt_caching: bool | None = None, api_base: object = None, + model_info: Mapping[str, object] | None = None, ) -> None: """For /chat/completions: resolve the injection points the request should carry. @@ -1357,9 +1689,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): configured: Final = configured_injection_points(non_default_params.get("cache_control_injection_points")) if configured: - tools_keeping_marks: Final = tuple( - tool for tool in tools or () if not _chat_transform_drops_tool_cache_control(tool) - ) + tools_keeping_marks: Final = AnthropicCacheControlHook.tools_keeping_cache_control_on_chat(tools) stamped: Final = AnthropicCacheControlHook._stamped_for_prompt_hook( configured, AnthropicCacheControlHook.count_external_cache_breakpoints( @@ -1369,6 +1699,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider, api_base, non_default_params.get("prompt_cache_options"), + model_info, ) non_default_params["cache_control_injection_points"] = ( AnthropicCacheControlHook._stamped(stamped, CUSTOM_LLM_PROVIDER_STAMP, custom_llm_provider) @@ -1385,6 +1716,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): enable_prompt_caching=enable_prompt_caching, cache_control=non_default_params.get("cache_control"), request_kwargs=non_default_params, + model_info=model_info, ) if points: non_default_params["cache_control_injection_points"] = points @@ -1503,6 +1835,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): cache_control=cache_control, request_kwargs=kwargs, on_messages_route=True, + model_info=routed_model_info(kwargs.get("model_info")), ) if model is not None else () @@ -1511,7 +1844,11 @@ class AnthropicCacheControlHook(CustomPromptManagement): return messages, system openai_dialect: Final = AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( - model, custom_llm_provider, api_base, kwargs.get("prompt_cache_options") + model, + custom_llm_provider, + api_base, + kwargs.get("prompt_cache_options"), + routed_model_info(kwargs.get("model_info")), ) breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) messages, system, remaining = AnthropicCacheControlHook.apply_to_anthropic_messages_request( @@ -1645,9 +1982,17 @@ def messages_within_cache_control_cap( Tool hints and a top-level ``cache_control`` stay on the wire, so they are reserved first. """ - reserved_marks: Final = AnthropicCacheControlHook.count_external_cache_breakpoints( - _validated_object_list(tools), cache_control, request_kwargs + return messages_within_reserved_cache_control_cap( + messages, + AnthropicCacheControlHook.count_external_cache_breakpoints( + _validated_object_list(tools), cache_control, request_kwargs + ), ) + + +def messages_within_reserved_cache_control_cap( + messages: Sequence[_MessageT], reserved_marks: int +) -> Sequence[_MessageT]: kept: Final = newest_marks_within_cache_control_limit(messages, reserved_marks) dropped: Final = sum(map(cache_control_count, messages)) - sum(map(cache_control_count, kept)) if dropped > 0: diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 37b4d2d3937..2ff523571f2 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -1,7 +1,8 @@ import asyncio import os import time -from collections.abc import Callable +from collections.abc import Callable, Mapping +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from functools import cache from typing import Final @@ -15,6 +16,7 @@ from litellm.constants import ( AZURE_STORAGE_MSFT_VERSION, ) from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id from litellm.llms.custom_httpx.http_handler import ( @@ -47,6 +49,20 @@ def adls_safe_file_name(payload_id: str | None) -> str: return f"{name.translate(_ADLS_SAFE_NAME)}.json" +@dataclass(frozen=True, slots=True) +class AzureStorageLogQueueItem: + payload: StandardLoggingPayload + file_name: str + + +def _queued_log(kwargs: Mapping[str, object], payload: StandardLoggingPayload) -> AzureStorageLogQueueItem: + payload_id: Final = payload.get("id") + return AzureStorageLogQueueItem( + payload=payload, + file_name=adls_safe_file_name(payload_id and logged_event_key(kwargs, payload_id)), + ) + + @cache def _cached_credential_chain_token_provider() -> Callable[[], str]: return get_azure_ad_token_provider( @@ -99,7 +115,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - self.log_queue: list[StandardLoggingPayload] = [] + self.log_queue: list[AzureStorageLogQueueItem] = [] super().__init__(**kwargs, flush_lock=self.flush_lock) except Exception as e: verbose_logger.exception( @@ -115,7 +131,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): def azure_storage_blob_endpoint(self) -> str: return f"https://{self.azure_storage_account_name}.blob.{self.azure_storage_endpoint_suffix}" - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs, response_obj: object, start_time: datetime, end_time: datetime): """ Async Log success events to Azure Blob Storage @@ -133,12 +149,12 @@ class AzureBlobStorageLogger(CustomBatchLogger): if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") - self.log_queue.append(standard_logging_payload) + self.log_queue.append(_queued_log(kwargs, standard_logging_payload)) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event(self, kwargs, response_obj: object, start_time: datetime, end_time: datetime): """ Async Log failure events to Azure Blob Storage @@ -156,7 +172,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") - self.log_queue.append(standard_logging_payload) + self.log_queue.append(_queued_log(kwargs, standard_logging_payload)) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e) @@ -195,13 +211,15 @@ class AzureBlobStorageLogger(CustomBatchLogger): len(self.log_queue), ) - for payload in self.log_queue: - await self.async_upload_payload_to_azure_blob_storage(payload=payload) + for item in self.log_queue: + await self.async_upload_payload_to_azure_blob_storage(payload=item.payload, file_name=item.file_name) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Error sending batch API - %s", e) - async def async_upload_payload_to_azure_blob_storage(self, payload: StandardLoggingPayload): + async def async_upload_payload_to_azure_blob_storage( + self, payload: StandardLoggingPayload, file_name: str | None = None + ): """ Uploads the payload to Azure Blob Storage using a 3-step process: 1. Create file resource @@ -210,14 +228,14 @@ class AzureBlobStorageLogger(CustomBatchLogger): """ try: if self.azure_storage_account_key: - await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload) + await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload, file_name=file_name) else: # Get a valid token instead of always requesting a new one await self.set_valid_azure_ad_token() async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) json_payload: Final = safe_dumps(payload) + "\n" # Add newline for each log entry payload_bytes: Final = json_payload.encode("utf-8") - filename: Final = adls_safe_file_name(payload.get("id")) + filename: Final = file_name or adls_safe_file_name(payload.get("id")) base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{filename}" # Execute the 3-step upload process @@ -377,7 +395,9 @@ class AzureBlobStorageLogger(CustomBatchLogger): self._service_client_timeout = self._clock() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS return self._service_client - async def upload_to_azure_data_lake_with_azure_account_key(self, payload: StandardLoggingPayload): + async def upload_to_azure_data_lake_with_azure_account_key( + self, payload: StandardLoggingPayload, file_name: str | None = None + ): """ Uploads the payload to Azure Data Lake using the Azure SDK @@ -386,11 +406,11 @@ class AzureBlobStorageLogger(CustomBatchLogger): try: today: Final = datetime.now().strftime("%Y-%m-%d") - file_name: Final = adls_safe_file_name(payload.get("id")) + resolved_file_name: Final = file_name or adls_safe_file_name(payload.get("id")) await self._upload_bytes_with_account_key( - file_path=f"{today}/{file_name}", content=safe_dumps(payload).encode("utf-8") + file_path=f"{today}/{resolved_file_name}", content=safe_dumps(payload).encode("utf-8") ) - verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, file_name) + verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, resolved_file_name) except Exception as e: verbose_logger.exception("Error occurred: %s", e) diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index 1d01965cd2c..c8453456052 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -19,6 +19,7 @@ from litellm._logging import verbose_logger from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger from litellm.integrations.clickhouse.context import is_lens_analysis from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.tracing.types import SpendLogPayload, SpendLogRecord @@ -134,7 +135,7 @@ def spend_log_row_from_payload( usage: Mapping[str, Any] = metadata.get("usage_object") or hidden_params.get("usage_object") or MappingProxyType({}) cache_read_tokens, cache_write_tokens = _cache_tokens(usage) trace_id, span_id = _find_traceparent(metadata, kwargs) - request_id = str(payload.get("id") or "") + response_id: Final = str(payload.get("id") or "") redact = litellm.turn_off_message_logging is True track_end_user: Final = not litellm.disable_end_user_cost_tracking completion_start_ms = _to_ms(payload.get("completionStartTime")) @@ -157,8 +158,8 @@ def spend_log_row_from_payload( ) ) return SpendLogRecord( - request_id=request_id, - response_id=strip_cache_hit_suffix(request_id), + request_id=logged_event_key(kwargs, response_id), + response_id=strip_cache_hit_suffix(response_id), provider_request_id=( get_provider_request_id(kwargs.get("response_headers")) or get_provider_request_id(hidden_params.get("additional_headers")) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index dfe508f3630..68876812c3e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1089,6 +1089,12 @@ class CustomGuardrail(CustomLogger): enrich_http_exception_with_guardrail_context(result, self) raise result + def runs_post_call_deployment_hook(self, request_data: Mapping[str, object]) -> bool: + return ( + isinstance(request_data.get("guardrails"), list) + and self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is True + ) + async def async_post_call_success_deployment_hook( self, request_data: dict, @@ -1098,12 +1104,7 @@ class CustomGuardrail(CustomLogger): """ Allow modifying / reviewing the response just after it's received from the deployment. """ - # should run guardrail - litellm_guardrails: Final = request_data.get("guardrails") - if litellm_guardrails is None or not isinstance(litellm_guardrails, list): - return None - - if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: + if not self.runs_post_call_deployment_hook(request_data): return None target: Final = self._deployment_hook_target() diff --git a/litellm/integrations/deepeval/deepeval.py b/litellm/integrations/deepeval/deepeval.py index 71d149e4014..e38b987d46f 100644 --- a/litellm/integrations/deepeval/deepeval.py +++ b/litellm/integrations/deepeval/deepeval.py @@ -15,6 +15,7 @@ from litellm.integrations.deepeval.utils import ( to_zod_compatible_iso, validate_environment, ) +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key # This file includes the custom callbacks for LiteLLM Proxy @@ -109,8 +110,10 @@ class DeepEvalLogger(CustomLogger): ) else: output = str(standard_logging_object.get("error_string", "")) + logged_id: Final = standard_logging_object.get("id", uuid.uuid4()) + span_id: Final = logged_event_key(kwargs, logged_id) if isinstance(logged_id, str) else logged_id return BaseApiSpan( - uuid=standard_logging_object.get("id", uuid.uuid4()), + uuid=span_id, name=("litellm_success_callback" if is_success else "litellm_failure_callback"), status=(TraceSpanApiStatus.SUCCESS if is_success else TraceSpanApiStatus.ERRORED), type=SpanApiType.LLM, diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py index 6569e70329b..8672eb00ce5 100644 --- a/litellm/integrations/dynamodb.py +++ b/litellm/integrations/dynamodb.py @@ -9,6 +9,8 @@ from typing import Final, Protocol import litellm from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key class _DynamoTable(Protocol): @@ -22,11 +24,13 @@ class _DynamoResource(Protocol): class DyanmoDBLogger: # Class variables or attributes - def __init__(self): + def __init__(self, dynamodb: _DynamoResource | None = None): # Instance variables import boto3 - self.dynamodb: Final[_DynamoResource] = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"]) + self.dynamodb: Final[_DynamoResource] = dynamodb or boto3.resource( + "dynamodb", region_name=os.environ["AWS_REGION_NAME"] + ) if litellm.dynamodb_table_name is None: raise ValueError( "LiteLLM Error, trying to use DynamoDB but not table name passed. Create a table and set `litellm.dynamodb_table_name=`" @@ -45,23 +49,23 @@ class DyanmoDBLogger: _async_log_event = async_log_event - def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): + def log_event(self, kwargs: Mapping[str, object], response_obj, start_time, end_time, print_verbose): try: print_verbose(f"DynamoDB Logging - Enters logging function for model {kwargs}") # construct payload to send to DynamoDB # follows the same params as langfuse.py - litellm_params: Final = kwargs.get("litellm_params", {}) - metadata: Final = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} messages: Final = kwargs.get("messages") optional_params: Final = kwargs.get("optional_params", {}) call_type: Final = kwargs.get("call_type", "litellm.completion") usage: Final = response_obj["usage"] - id: Final = response_obj.get("id", str(uuid.uuid4())) + response_id: Final = response_obj.get("id", str(uuid.uuid4())) # Build the initial payload payload: Final[dict[str, object]] = { - "id": id, + "id": logged_event_key(kwargs, response_id) if isinstance(response_id, str) else response_id, "call_type": call_type, "startTime": start_time, "endTime": end_time, diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 99cd8b31b63..7057e719764 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( sanitize_cloud_object_component, ) from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload @@ -296,30 +297,28 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): else: await self._send_individual_logs(items_to_process) - def _get_object_name(self, kwargs: dict, logging_payload: StandardLoggingPayload, response_obj: Any) -> str: + def _get_object_name( + self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: Any + ) -> str: """ Get the object name to use for the current payload """ current_date: Final = self._get_object_date_from_datetime(datetime.now(timezone.utc)) + metadata: Final = as_str_mapping((as_str_mapping(kwargs.get("litellm_params")) or {}).get("metadata")) or {} + gcs_log_id: Final = metadata.get("gcs_log_id") + safe_log_id: Final = sanitize_cloud_object_component( + gcs_log_id if isinstance(gcs_log_id, str) else None, fallback="" + ) + if safe_log_id: + return f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" if logging_payload.get("error_str", None) is not None: - object_name = self._generate_failure_object_name( - request_date_str=current_date, - ) - else: - object_name = self._generate_success_object_name( - request_date_str=current_date, - response_id=response_obj.get("id", ""), - ) - - # used for testing - _litellm_params: Final = kwargs.get("litellm_params", None) or {} - _metadata: Final = _litellm_params.get("metadata", None) or {} - if "gcs_log_id" in _metadata: - safe_log_id: Final = sanitize_cloud_object_component(_metadata.get("gcs_log_id"), fallback="") - if safe_log_id: - object_name = f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" - - return object_name + return self._generate_failure_object_name(request_date_str=current_date) + raw_response_id: Final[object] = response_obj.get("id", "") + response_id: Final = raw_response_id if isinstance(raw_response_id, str) else "" + return self._generate_success_object_name( + request_date_str=current_date, + response_id=logged_event_key(kwargs, response_id), + ) async def get_request_response_payload( self, diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index db2fe386dec..918e849a086 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -3,12 +3,15 @@ import json import os +from collections.abc import Mapping from typing import Final import httpx import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -48,7 +51,7 @@ class OpenMeterLogger(CustomLogger): if len(missing_keys) > 0: raise Exception(f"Missing keys={missing_keys} in environment.") - def _common_logic(self, kwargs: dict, response_obj): + def _common_logic(self, kwargs: Mapping[str, object], response_obj): call_id: Final = response_obj.get("id", kwargs.get("litellm_call_id")) dt: Final = get_utc_datetime().isoformat() cost: Final = kwargs.get("response_cost", None) @@ -74,8 +77,8 @@ class OpenMeterLogger(CustomLogger): # If no user provided directly, try to get it from token user_id if user_param is None: # Check if user_id is available from the API key metadata - litellm_params: Final = kwargs.get("litellm_params", {}) - metadata: Final = litellm_params.get("metadata", {}) + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} user_api_key_user_id: Final = metadata.get("user_api_key_user_id", None) if user_api_key_user_id is not None: @@ -89,7 +92,7 @@ class OpenMeterLogger(CustomLogger): return { "specversion": "1.0", "type": os.getenv("OPENMETER_EVENT_TYPE", "litellm_tokens"), - "id": call_id, + "id": logged_event_key(kwargs, call_id) if isinstance(call_id, str) else call_id, "time": dt, "subject": subject, "source": "litellm-proxy", diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 42fe0837c29..1c9a443b546 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3036,7 +3036,7 @@ class PrometheusLogger(CustomLogger): async def async_post_call_failure_hook( self, - request_data: dict, + request_data: dict[str, object], original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, traceback_str: str | None = None, @@ -3059,10 +3059,14 @@ class PrometheusLogger(CustomLogger): try: _tags: Final = StandardLoggingPayloadSetup.get_request_tags( litellm_params=request_data, - proxy_server_request=request_data.get("proxy_server_request", {}), + proxy_server_request=_logged_fields(request_data.get("proxy_server_request")), + ) + _metadata: Final = _logged_fields(request_data.get("metadata")) + model_id: Final = ( + _logged_fields(_metadata.get("model_info")).get("id") + or _logged_fields(request_data.get("model_info")).get("id") + or _logged_fields(request_data.get("standard_logging_object")).get("model_id") ) - _metadata: Final = request_data.get("metadata", {}) or {} - model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id") rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception) api_provider: Final = self._extract_api_provider_from_request_data( request_data diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 2bb89248be4..d41647fe5f1 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast from typing_extensions import Never, ReadOnly import litellm +from litellm._internal_context import billed_separately from litellm._logging import verbose_logger from litellm.anthropic_interface import messages as anthropic_messages from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME @@ -1628,19 +1629,20 @@ class WebSearchInterceptionLogger(CustomLogger): search_kwargs: Final = MappingProxyType( {**configured_search_kwargs, **parent_correlation.as_search_kwargs()} ) - result: Final = ( - await litellm.asearch( - query=query_arg, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs + with billed_separately(): + result: Final = ( + await litellm.asearch( + query=query_arg, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs + ) + if search_metadata is None + else await litellm.asearch( + query=query_arg, + search_provider=search_provider, + litellm_metadata=search_metadata, + **_NO_ASEARCH_NAMED, + **search_kwargs, + ) ) - if search_metadata is None - else await litellm.asearch( - query=query_arg, - search_provider=search_provider, - litellm_metadata=search_metadata, - **_NO_ASEARCH_NAMED, - **search_kwargs, - ) - ) # Format using transformation function search_result_text: Final = WebSearchTransformation.format_search_response(result) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py new file mode 100644 index 00000000000..5c8f5d83ece --- /dev/null +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -0,0 +1,111 @@ +from collections.abc import Mapping +from typing import Final, Literal + +from openai.types.responses.tool_choice_allowed_param import ToolChoiceAllowedParam +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +from litellm.exceptions import BadRequestError + +ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: Final = frozenset({"openai", "azure", "litellm_proxy"}) + +CHAT_ALLOWED_TOOLS_LAYOUT: Final = ( + '{"type": "allowed_tools", "allowed_tools": {"mode": "auto" | "required", "tools": []}}' +) + + +class ChatAllowedTools(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + mode: Literal["auto", "required"] + tools: tuple[Mapping[str, object], ...] = Field(min_length=1) + + +class ChatAllowedToolsToolChoice(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["allowed_tools"] + allowed_tools: ChatAllowedTools + + +_CHAT_ALLOWED_TOOLS_TOOL_CHOICE: Final = TypeAdapter(ChatAllowedToolsToolChoice) + + +_SELECTOR: Final = TypeAdapter(dict[str, object]) + + +class _ToolChoiceType(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str + + +_TOOL_CHOICE_TYPE: Final = TypeAdapter(_ToolChoiceType) + + +def is_allowed_tools_tool_choice(tool_choice: object) -> bool: + try: + return _TOOL_CHOICE_TYPE.validate_python(tool_choice).type == "allowed_tools" + except ValidationError: + return False + + +def raise_if_allowed_tools_tool_choice_unsupported( + tool_choice: object, + custom_llm_provider: str, + model: str, +) -> None: + """Providers outside the set would silently send the request with no tool restriction at all, so the call stops + here with a 400. ``drop_params`` is deliberately not consulted: dropping a restriction widens what the model may call.""" + if not is_allowed_tools_tool_choice(tool_choice) or custom_llm_provider in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: + return + raise BadRequestError( + message=( + "tool_choice type 'allowed_tools' is only supported for " + f"{', '.join(sorted(ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS))}, got custom_llm_provider={custom_llm_provider}. " + "drop_params does not drop it, since that would let the model call tools the request excluded. " + "Send only the allowed tools in `tools` with tool_choice='auto' or 'required' instead." + ), + model=model, + llm_provider=custom_llm_provider, + ) + + +def parse_chat_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsToolChoice | None: + try: + return _CHAT_ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) + except ValidationError: + return None + + +def _has_name(selector: Mapping[str, object]) -> bool: + name: Final = selector.get("name") + return isinstance(name, str) and name != "" + + +def _responses_allowed_tool(tool: Mapping[str, object]) -> Mapping[str, object] | None: + tool_type: Final = tool.get("type") + if tool_type not in ("function", "custom"): + return tool + if tool_type not in tool: + return tool if _has_name(tool) else None + try: + nested: Final = _SELECTOR.validate_python(tool[tool_type]) + except ValidationError: + return None + outer: Final = {key: value for key, value in tool.items() if key != tool_type} + if not _has_name(nested) or outer.keys() & nested.keys(): + return None + return {**outer, **nested} + + +def responses_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ToolChoiceAllowedParam | None: + """The flat Responses API layout of a chat ``allowed_tools`` choice, or ``None`` when an entry is malformed or + carries conflicting keys, so the caller sends the original rather than a rewrite that could mean something looser.""" + flattened: Final = tuple(_responses_allowed_tool(tool) for tool in choice.allowed_tools.tools) + if any(tool is None for tool in flattened): + return None + return ToolChoiceAllowedParam( + type="allowed_tools", + mode=choice.allowed_tools.mode, + tools=[dict(tool) for tool in flattened if tool is not None], + ) diff --git a/litellm/litellm_core_utils/audio_utils/transcription_streaming.py b/litellm/litellm_core_utils/audio_utils/transcription_streaming.py new file mode 100644 index 00000000000..4650b7db2f2 --- /dev/null +++ b/litellm/litellm_core_utils/audio_utils/transcription_streaming.py @@ -0,0 +1,341 @@ +from __future__ import annotations + +import datetime +import json +import traceback +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from typing import Final, Protocol + +import anyio +from openai import AsyncStream, Stream +from openai.types.audio import ( + TranscriptionStreamEvent, + TranscriptionTextDeltaEvent, + TranscriptionTextDoneEvent, +) +from pydantic import TypeAdapter, ValidationError + +from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( + TranscriptionUsageObjectTransformation, +) +from litellm.types.utils import TranscriptionResponse, TranscriptionUsageTokensObject, Usage + +_STREAM_FLAG: Final = TypeAdapter(bool) + + +def transcription_stream_flag(value: object) -> bool | None: + try: + return _STREAM_FLAG.validate_python(value) + except ValidationError: + return None + + +def transcription_stream_unset(value: object) -> bool: + return value is None or value == "" + + +def transcription_stream_requested(value: object) -> bool: + return transcription_stream_flag(value) is True + + +TRANSCRIPTION_STREAM_DONE_FRAME: Final = "data: [DONE]\n\n" + +TranscriptionStreamExceptionMapper = Callable[[Exception], Exception] + + +def litellm_exception_mapper(call_details: Mapping[str, object]) -> TranscriptionStreamExceptionMapper: + from litellm.litellm_core_utils.exception_mapping_utils import exception_type, with_deployment_api_base + + model: Final = call_details.get("model") + custom_llm_provider: Final = call_details.get("custom_llm_provider") + extra_kwargs: Final = with_deployment_api_base(call_details.get("litellm_params")) + + def map_exception(exception: Exception) -> Exception: + try: + mapped: Final[object] = exception_type( + model=model if isinstance(model, str) else None, + original_exception=exception, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + extra_kwargs=extra_kwargs, + ) + except Exception as mapped_exception: # noqa: BLE001 # exception_type raises the mapped exception + return mapped_exception + return mapped if isinstance(mapped, Exception) else exception + + return map_exception + + +def _guarded_transcription_event_payloads( + response: TranscriptionResponse, +) -> tuple[dict[str, object], dict[str, object]]: + text: Final = response.text or "" + usage: Final = response.usage + languages: Final = TranscriptionResponse.detected_languages(getattr(response, "languages", None)) + return ( + {"type": "transcript.text.delta", "delta": text}, + { + "type": "transcript.text.done", + "text": text, + **({} if usage is None else {"usage": usage.model_dump(exclude_none=True)}), + **({} if languages is None else {"languages": [language.model_dump() for language in languages]}), + }, + ) + + +def guarded_transcription_events(response: TranscriptionResponse) -> tuple[str, ...]: + return ( + *( + f"data: {json.dumps(event, separators=(',', ':'))}\n\n" + for event in _guarded_transcription_event_payloads(response) + ), + TRANSCRIPTION_STREAM_DONE_FRAME, + ) + + +class GuardedTranscriptionReplay: + def __init__(self, response: TranscriptionResponse) -> None: + self.response: Final = response + self._hidden_params: dict[str, object] = response.hidden_params + + def __aiter__(self) -> AsyncIterator[TranscriptionStreamEvent]: + return self._events() + + async def _events(self) -> AsyncIterator[TranscriptionStreamEvent]: + delta, done = _guarded_transcription_event_payloads(self.response) + yield TranscriptionTextDeltaEvent.model_construct(**delta) + yield TranscriptionTextDoneEvent.model_construct(**done) + + +class TranscriptionStreamLogging(Protocol): + def success_handler( + self, + result: TranscriptionResponse, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: ... + + async def async_success_handler( + self, + result: TranscriptionResponse, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: ... + + def handle_sync_success_callbacks_for_async_calls( + self, + result: TranscriptionResponse, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: ... + + def failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime, + end_time: datetime.datetime, + *, + for_async_request: bool = False, + ) -> None: ... + + async def async_failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: ... + + def response_cost_calculator(self, result: TranscriptionResponse) -> float | None: ... + + def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None: ... + + +class _TranscriptionEventCollector: + def __init__(self, duration: float | None) -> None: + self.duration = duration + self.text_deltas: list[str] = [] # mutable-ok: streaming deltas accumulate until the terminal event + self.done_event: TranscriptionTextDoneEvent | None = None + + def add(self, event: TranscriptionStreamEvent) -> None: + if isinstance(event, TranscriptionTextDeltaEvent): + self.text_deltas.append(event.delta) + elif isinstance(event, TranscriptionTextDoneEvent): + self.done_event = event + + def response(self) -> TranscriptionResponse: + done_event: Final = self.done_event + response: Final = TranscriptionResponse( + text=done_event.text if done_event is not None else "".join(self.text_deltas), + usage=done_event.usage.model_dump() if done_event is not None and done_event.usage is not None else None, + languages=getattr(done_event, "languages", None) if done_event is not None else None, + ) + if self.duration is not None: + response.set_audio_transcription_duration(self.duration) + return response + + +def _billed_usage(response: TranscriptionResponse) -> Usage: + usage: Final[object] = getattr(response, "usage", None) + token_usage: Final = ( + TranscriptionUsageObjectTransformation.transform_transcription_usage_object(usage) + if isinstance(usage, TranscriptionUsageTokensObject) + else None + ) + return token_usage or Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + + +def _record_billable_partial_spend( + logging_obj: TranscriptionStreamLogging, collector: _TranscriptionEventCollector +) -> None: + partial_response: Final = collector.response() + partial_cost: Final = logging_obj.response_cost_calculator(result=partial_response) or 0.0 + if partial_cost <= 0: + return + logging_obj.record_partial_usage_for_failure(_billed_usage(partial_response), partial_cost) + + +class LoggingTranscriptionStream(Stream[TranscriptionStreamEvent]): + def __init__( + self, + stream: Stream[TranscriptionStreamEvent], + logging_obj: TranscriptionStreamLogging, + start_time: datetime.datetime, + map_exception: TranscriptionStreamExceptionMapper, + ) -> None: + self.__dict__.update(stream.__dict__) + self._hidden_params: dict[str, object] = { # mutable-ok: ResponseMetadata fills in the call's metadata + "additional_headers": process_response_headers(stream.response.headers) + } + self._logging_obj = logging_obj + self._start_time = start_time + self._map_exception = map_exception + self._collector = _TranscriptionEventCollector(getattr(stream, "_litellm_audio_duration", None)) + self._finalized = False + self._failed = False + source_iterator: Final = self._iterator + self._iterator = self._logging_iterator(source_iterator) + + def _logging_iterator( + self, source_iterator: Iterator[TranscriptionStreamEvent] + ) -> Iterator[TranscriptionStreamEvent]: + try: + for event in source_iterator: + self._collector.add(event) + yield event + except Exception as exception: + self._failed = True + _record_billable_partial_spend(self._logging_obj, self._collector) + mapped_exception: Final = self._map_exception(exception) + end_time: Final = datetime.datetime.now() # noqa: DTZ005 # callback timestamps use the legacy naive contract + self._logging_obj.failure_handler(mapped_exception, traceback.format_exc(), self._start_time, end_time) + raise mapped_exception from exception + finally: + self._finalize() + + def _finalize(self) -> None: + if self._finalized or self._failed: + return + self._finalized = True + self._logging_obj.success_handler( + self._collector.response(), + self._start_time, + datetime.datetime.now(), # noqa: DTZ005 # callback timestamps use the legacy naive contract + ) + + def close(self) -> None: + try: + super().close() + finally: + self._finalize() + + @property + def logging_obj(self) -> TranscriptionStreamLogging: + return self._logging_obj + + +class LoggingAsyncTranscriptionStream(AsyncStream[TranscriptionStreamEvent]): + def __init__( + self, + stream: AsyncStream[TranscriptionStreamEvent], + logging_obj: TranscriptionStreamLogging, + start_time: datetime.datetime, + map_exception: TranscriptionStreamExceptionMapper, + ) -> None: + self.__dict__.update(stream.__dict__) + self._hidden_params: dict[str, object] = { # mutable-ok: ResponseMetadata fills in the call's metadata + "additional_headers": process_response_headers(stream.response.headers) + } + self._logging_obj = logging_obj + self._start_time = start_time + self._map_exception = map_exception + self._collector = _TranscriptionEventCollector(getattr(stream, "_litellm_audio_duration", None)) + self._finalized = False + self._failed = False + source_iterator: Final = self._iterator + self._iterator = self._logging_iterator(source_iterator) + + async def _logging_iterator( + self, source_iterator: AsyncIterator[TranscriptionStreamEvent] + ) -> AsyncIterator[TranscriptionStreamEvent]: + try: + async for event in source_iterator: + self._collector.add(event) + yield event + except Exception as exception: + self._failed = True + _record_billable_partial_spend(self._logging_obj, self._collector) + mapped_exception: Final = self._map_exception(exception) + end_time: Final = datetime.datetime.now() # noqa: DTZ005 # callback timestamps use the legacy naive contract + self._logging_obj.failure_handler( + mapped_exception, traceback.format_exc(), self._start_time, end_time, for_async_request=True + ) + await self._logging_obj.async_failure_handler( + mapped_exception, + traceback.format_exc(), + self._start_time, + end_time, + ) + raise mapped_exception from exception + finally: + await self._finalize() + + async def _finalize(self) -> None: + if self._finalized or self._failed: + return + self._finalized = True + response: Final = self._collector.response() + end_time: Final = datetime.datetime.now() # noqa: DTZ005 # callback timestamps use the legacy naive contract + self._logging_obj.handle_sync_success_callbacks_for_async_calls( + result=response, + start_time=self._start_time, + end_time=end_time, + ) + with anyio.CancelScope(shield=True): + await self._logging_obj.async_success_handler( + result=response, + start_time=self._start_time, + end_time=end_time, + ) + + async def close(self) -> None: + try: + await super().close() + finally: + await self._finalize() + + @property + def logging_obj(self) -> TranscriptionStreamLogging: + return self._logging_obj + + +def wrap_transcription_stream( + stream: Stream[TranscriptionStreamEvent] | AsyncStream[TranscriptionStreamEvent], + logging_obj: TranscriptionStreamLogging, + start_time: datetime.datetime, + map_exception: TranscriptionStreamExceptionMapper, +) -> LoggingTranscriptionStream | LoggingAsyncTranscriptionStream: + if isinstance(stream, AsyncStream): + return LoggingAsyncTranscriptionStream(stream, logging_obj, start_time, map_exception) + return LoggingTranscriptionStream(stream, logging_obj, start_time, map_exception) diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 640559ac7e2..29b10ac3c83 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -21,6 +21,16 @@ from litellm.types.files import ( from litellm.types.utils import FileTypes +def normalized_audio_duration_seconds(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + try: + seconds: Final = float(value) + except OverflowError: + return None + return seconds if math.isfinite(seconds) and seconds >= 0 else None + + @dataclass class ProcessedAudioFile: """ diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 2ba47cbb400..1f29b5d7a4d 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -8,6 +8,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm._internal_context import internal_sub_call from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs @@ -204,12 +205,13 @@ async def _execute_chat_completion_agentic_plan( ) try: - response_followup = await litellm.acompletion( - model=full_model_name, - messages=patch.messages, - **optional_params_for_followup, - **kwargs_for_followup, - ) + with internal_sub_call(logging_obj.litellm_call_id if isinstance(logging_obj, LiteLLMLoggingObject) else None): + response_followup = await litellm.acompletion( + model=full_model_name, + messages=patch.messages, + **optional_params_for_followup, + **kwargs_for_followup, + ) if _post_hook_overridden(callback): try: response_followup = await callback.async_post_agentic_loop_response_hook( diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 58189140a98..43c586e7245 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -4,7 +4,7 @@ import copy import logging import math import re -from collections.abc import Collection, Iterable, Iterator, Mapping +from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable @@ -97,13 +97,17 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non return float(seconds / denominator) -_DROP_PARAMS_BOOL: Final = TypeAdapter(bool) +_LAX_BOOL: Final = TypeAdapter(bool) def as_str_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, Mapping) else None +def as_object_list(value: object) -> Sequence[object] | None: + return value if isinstance(value, list) else None + + def decoded_json_object(value: object) -> Mapping[str, object]: mapping: Final = as_str_mapping(value) if mapping is not None: @@ -156,15 +160,19 @@ def attribute_items(source: object, name: str) -> Iterable[object]: return iter_items(value) if value else () -def normalize_drop_params(value: object) -> bool | None: +def lax_bool(value: object) -> bool | None: if value is None or isinstance(value, bool): return value try: - return _DROP_PARAMS_BOOL.validate_python(value.strip() if isinstance(value, str) else value) + return _LAX_BOOL.validate_python(value.strip() if isinstance(value, str) else value) except ValidationError: return None +def normalize_drop_params(value: object) -> bool | None: + return lax_bool(value) + + def drops_unsupported_params(drop_params: object) -> bool: return litellm.drop_params is True or normalize_drop_params(drop_params) is True diff --git a/litellm/litellm_core_utils/decisions_pricing.py b/litellm/litellm_core_utils/decisions_pricing.py index 9199a8097cc..accd315da23 100644 --- a/litellm/litellm_core_utils/decisions_pricing.py +++ b/litellm/litellm_core_utils/decisions_pricing.py @@ -12,6 +12,7 @@ from typing import Final DECISIONS_PRICING_SUFFIX: Final = "_decisions" DECISIONS_CALL_TYPES: Final = frozenset({"decisions", "adecisions", "systemone", "asystemone"}) +QUERY_PRICED_DECISIONS_PROVIDERS: Final = frozenset({"parallel_ai"}) def is_decisions_call_type(call_type: str) -> bool: @@ -27,6 +28,17 @@ def decisions_pricing_fields(row: Mapping[str, object]) -> Mapping[str, object]: } +def decisions_request_price(model_info: Mapping[str, object], overrides: Mapping[str, object]) -> float | None: + if model_info.get("litellm_provider") not in QUERY_PRICED_DECISIONS_PROVIDERS: + return None + rates: Final = {**model_info, **overrides} + request_price: Final = rates.get("input_cost_per_query") + if not isinstance(request_price, (int, float)): + return None + bills_tokens: Final = any(rates.get(field) for field in ("input_cost_per_token", "output_cost_per_token")) + return None if bills_tokens and "input_cost_per_query" not in overrides else float(request_price) + + def apply_decisions_pricing(row: Mapping[str, object]) -> Mapping[str, object]: """Return ``row`` priced for a Decisions call: ``_decisions`` values replace their plain fields.""" plain: Final = {key: value for key, value in row.items() if not key.endswith(DECISIONS_PRICING_SUFFIX)} diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 85c232ad3f4..f1195802fc5 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -107,6 +107,7 @@ class ExceptionCheckers: "exceeds the available context size", # llama.cpp/Lemonade "exceeds the maximum number of tokens allowed", # Gemini '"code": "context_length_exceeded"', + "prompt is too long", ] for substring in known_exception_substrings: if substring in _error_str_lowercase: diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 38242291b71..94ba02bfbd7 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -146,6 +146,10 @@ def handle_anthropic_text_model_custom_llm_provider( return model, custom_llm_provider +def declared_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: + return custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) + + def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: """The authenticating provider this pair already names, or None. @@ -154,7 +158,7 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str and for a declared pair the resolver's answer is the declaration itself, so metadata callers adopt the declaration instead of resolving. """ - declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) + declared: Final = declared_provider(model, custom_llm_provider) return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 453d4358b35..138a01c6524 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -26,7 +26,12 @@ from litellm.constants import ( MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, NON_INFERENCE_CALL_TYPES, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, @@ -35,6 +40,8 @@ from litellm.types.utils import ( BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" + ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" @@ -151,6 +158,26 @@ def is_unbilled_non_inference_call_from_params( return is_unbilled_non_inference_call(call_type, metadata, response) +def settles_background_response(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> bool: + return kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES or ( + call_type in NON_INFERENCE_CALL_TYPES + and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + ) + + +def logged_event_key(kwargs: Mapping[str, object], event_id: str) -> str: + call_type: Final = kwargs.get("call_type") + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = ( + as_str_mapping(litellm_params.get("litellm_metadata")) or as_str_mapping(litellm_params.get("metadata")) or {} + ) + if not event_id or not settles_background_response( + call_type if isinstance(call_type, str) else "", metadata, kwargs + ): + return event_id + return f"{event_id}{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}" + + def sanitize_user_api_key_auth(auth: object) -> object: """Copy of the auth object with its budget reservation removed; the cost callback falls back to reading the reservation from inside the auth object.""" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 639003c9a6a..97c019a2180 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,9 +19,11 @@ from datetime import datetime as dt_object from enum import Enum from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, TypeVar, cast from httpx import Response, Timeout +from openai import AsyncStream as OpenAIAsyncStream +from openai import Stream as OpenAIStream from pydantic import BaseModel, JsonValue, TypeAdapter import litellm @@ -97,6 +99,7 @@ from litellm.litellm_core_utils.hidden_params import ( get_hidden_params, served_from_cache, set_hidden_param, + set_hidden_params, ) from litellm.litellm_core_utils.internal_call_metadata import ( MODEL_ACCESS_GROUP_METADATA_KEY, @@ -122,6 +125,7 @@ from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import ( diagnose_zero_cost, zero_cost_warning, ) +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import LiteLLMResponseObjectHandler from litellm.litellm_core_utils.logging_utils import ( truncate_base64_in_messages, truncate_base64_in_messages_async, @@ -142,7 +146,7 @@ from litellm.litellm_core_utils.served_output_texts import ( from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.search.transformation import SearchResponse -from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.agents import LiteLLMSendMessageResponse from litellm.types.containers.main import ContainerObject from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse @@ -300,6 +304,11 @@ class _EnterpriseAlertingLoggers(NamedTuple): smtp_email: type[CustomLogger] +class BridgedChatRequest(NamedTuple): + optional_params: Mapping[str, object] + litellm_params: Mapping[str, object] + + @functools.cache def _enterprise_alerting_loggers() -> _EnterpriseAlertingLoggers: try: @@ -580,6 +589,7 @@ _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) _SERVED_ANTHROPIC_BODY: Final = TypeAdapter(dict[str, object]) _GOOGLE_USAGE_FIELDS: Final = frozenset(UsageMetadata.__annotations__) _REDACTED_RESULT: Final = TypeAdapter(object) +_LoggedResponse: Final = TypeVar("_LoggedResponse") _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" @@ -872,6 +882,8 @@ class Logging(LiteLLMLoggingBaseClass): # The model the proxy shows the client on streamed chunks. The logged streamed response carries it # once that response is priced, the same way a non-streamed response is logged self.client_facing_stream_model: str | None = None + self.client_facing_response_id: str | None = None + self.bridged_chat_request: BridgedChatRequest | None = None self.zero_cost_warned: bool = False self.llm_caching_handler: LLMCachingHandler | None = None @@ -918,6 +930,7 @@ class Logging(LiteLLMLoggingBaseClass): # post_call guardrails have run; the @client decorator then stores the # enqueue closure here instead of firing it immediately. self.defer_async_logging: bool = False + self.drops_stream_token_ids: bool = False self.enqueue_deferred_logging: Callable[[], None] | None = None self._async_success_scheduled: bool = False self._async_success_enqueued: bool = False @@ -2128,6 +2141,9 @@ class Logging(LiteLLMLoggingBaseClass): if resolved_cache_hit is True: return 0.0 + if isinstance(result, (OpenAIStream, OpenAIAsyncStream)): + return None + if is_unbilled_non_inference_call( self.call_type, StandardLoggingPayloadSetup.merge_litellm_metadata(self.litellm_params), result ): @@ -3228,6 +3244,40 @@ class Logging(LiteLLMLoggingBaseClass): return response return response.model_copy(update={"model": model}) + def _as_served_chat_response(self, response: _LoggedResponse) -> _LoggedResponse | ModelResponse: + request: Final = self.bridged_chat_request + if request is None or not isinstance(response, ResponsesAPIResponse): + return response + return self._bridged_chat_response(response, request) + + def _as_served_response( + self, response: _LoggedResponse + ) -> _LoggedResponse | ModelResponse | TextCompletionResponse: + request: Final = self.bridged_chat_request + if request is None or not isinstance(response, ResponsesAPIResponse): + return response + served: Final = self._bridged_chat_response(response, request) + if self.call_type not in (CallTypes.text_completion.value, CallTypes.atext_completion.value): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped + return served + return self._as_text_completion_priced_as_chat(served) + + def _as_text_completion_priced_as_chat(self, served: ModelResponse) -> TextCompletionResponse: + text_completion: Final = LiteLLMResponseObjectHandler.convert_chat_to_text_completion( + served, TextCompletionResponse() + ) + set_hidden_params(text_completion, dict(get_hidden_params(served) or {})) + return text_completion + + def _bridged_chat_response(self, response: ResponsesAPIResponse, request: BridgedChatRequest) -> ModelResponse: + translated: Final = self._translate_responses_api_response_to_model_response(response, request) + served: Final = translated.model_copy( + update={ + "id": ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response.id) + } + ) + set_hidden_params(served, {**(get_hidden_params(translated) or {}), **(get_hidden_params(response) or {})}) + return served + def _with_continued_response_id( self, response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, @@ -3237,6 +3287,23 @@ class Logging(LiteLLMLoggingBaseClass): return response return response.model_copy(update={"id": response_id}) + def _with_client_facing_response_id( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + response_id: Final = self.client_facing_response_id + if response_id is None or getattr(response, "id", None) == response_id: + return response + return response.model_copy(update={"id": response_id}) + + def _client_facing_streaming_response( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + return self._with_client_facing_response_id( + self._with_continued_response_id(self._with_client_facing_stream_model(response)) + ) + def _success_handler_helper_fn( self, result: object = None, @@ -3292,6 +3359,8 @@ class Logging(LiteLLMLoggingBaseClass): ) result = self.normalize_endpoint_logging_result(result) + if self.stream is not True: + result = self._as_served_response(result) # rebind-ok: bridged chat calls log the served response logging_result: Final = self.normalize_logging_result(result=result) if cache_hit is None and _served_from_cache_at_no_charge(logging_result): @@ -3606,6 +3675,7 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit) if not self.should_run_logging(event_type="sync_success"): # prevent double logging return + provider_transcription: Final = result if isinstance(result, TranscriptionResponse) else None start_time, end_time, result, request_logging_payload = self.prepare_success_logging( start_time=start_time, end_time=end_time, @@ -3620,7 +3690,12 @@ class Logging(LiteLLMLoggingBaseClass): try: ## BUILD COMPLETE STREAMED RESPONSE complete_streaming_response: ( - ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None + ModelResponse + | TextCompletionResponse + | ResponsesAPIResponse + | InteractionsAPIResponse + | TranscriptionResponse + | None ) = None if "complete_streaming_response" in self.model_call_details and not for_async_request: return # break out of this. @@ -3636,19 +3711,21 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["complete_streaming_response"] = complete_streaming_response self._surface_response_headers_from_result(complete_streaming_response) self.model_call_details["response_cost"] = self.response_cost_calculator( - result=complete_streaming_response + result=complete_streaming_response if provider_transcription is None else provider_transcription ) self._add_continued_partial_spend(complete_streaming_response) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) - logged_streaming_response: Final = self._with_continued_response_id( - self._with_client_facing_stream_model(complete_streaming_response) + logged_streaming_response: Final = self._client_facing_streaming_response( + self._as_served_chat_response(complete_streaming_response) ) self.model_call_details["complete_streaming_response"] = logged_streaming_response self.capture_usage_before_logging(logged_streaming_response, start_time, end_time) ## STANDARDIZED LOGGING PAYLOAD if not for_async_request: self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - logged_streaming_response, start_time, end_time + logged_streaming_response if provider_transcription is None else provider_transcription, + start_time, + end_time, ) standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( "standard_logging_object" @@ -4175,6 +4252,7 @@ class Logging(LiteLLMLoggingBaseClass): kwargs=self.model_call_details, messages=self.model_call_details.get("messages") ) ) + provider_transcription: Final = result if isinstance(result, TranscriptionResponse) else None start_time, end_time, result = self._success_handler_helper_fn( start_time=start_time, end_time=end_time, @@ -4194,9 +4272,7 @@ class Logging(LiteLLMLoggingBaseClass): ## BUILD COMPLETE STREAMED RESPONSE if "async_complete_streaming_response" in self.model_call_details: return # break out of this. - complete_streaming_response: Final[ - ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None - ] = self._get_assembled_streaming_response( + complete_streaming_response: Final = self._get_assembled_streaming_response( result=result, start_time=start_time, end_time=end_time, @@ -4218,7 +4294,7 @@ class Logging(LiteLLMLoggingBaseClass): get_base_model_from_metadata(model_call_details=self.model_call_details) # base_model defaults to None if not set on model_info self.model_call_details["response_cost"] = self.response_cost_calculator( - result=complete_streaming_response + result=complete_streaming_response if provider_transcription is None else provider_transcription ) verbose_logger.debug("Model=%s; cost=%s", self.model, self.model_call_details["response_cost"]) @@ -4239,15 +4315,17 @@ class Logging(LiteLLMLoggingBaseClass): await self._prepare_baseline_cache_estimate(complete_streaming_response) - logged_streaming_response: Final = self._with_continued_response_id( - self._with_client_facing_stream_model(complete_streaming_response) + logged_streaming_response: Final = self._client_facing_streaming_response( + self._as_served_chat_response(complete_streaming_response) ) self.model_call_details["async_complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD try: self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - logged_streaming_response, start_time, end_time + logged_streaming_response if provider_transcription is None else provider_transcription, + start_time, + end_time, ) except Exception: # noqa: BLE001 # payload build must never block later callbacks (slot release) verbose_logger.exception( @@ -5117,10 +5195,17 @@ class Logging(LiteLLMLoggingBaseClass): end_time: datetime.datetime, is_async: bool, streaming_chunks: list[object], - ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None: + ) -> ( + ModelResponse + | TextCompletionResponse + | ResponsesAPIResponse + | InteractionsAPIResponse + | TranscriptionResponse + | None + ): if self.stream is not True: return None - if isinstance(result, ModelResponse) or isinstance(result, TextCompletionResponse): + if isinstance(result, (ModelResponse, TextCompletionResponse, TranscriptionResponse)): return result elif isinstance( result, @@ -5191,6 +5276,13 @@ class Logging(LiteLLMLoggingBaseClass): row is built from this response rather than re-parsed from the served body.""" self.bridged_model_response = response + def record_bridged_chat_request( + self, optional_params: Mapping[str, object], litellm_params: Mapping[str, object] + ) -> None: + self.bridged_chat_request = BridgedChatRequest( + MappingProxyType(dict(optional_params)), MappingProxyType(dict(litellm_params)) + ) + def record_bridged_responses_api_response(self, response: ResponsesAPIResponse) -> None: self.record_bridged_model_response(self._translate_responses_api_response_to_model_response(response)) @@ -5262,6 +5354,9 @@ class Logging(LiteLLMLoggingBaseClass): ): result = result.response if isinstance(result, ResponsesAPIResponse): + bridged_request: Final = self.bridged_chat_request + if bridged_request is not None: + return self._bridged_chat_response(result, bridged_request) return self._translate_responses_api_response_to_model_response(result) raw_body: Final[object] = result @@ -5313,7 +5408,9 @@ class Logging(LiteLLMLoggingBaseClass): set_response_cost_in_hidden_params(result, reported_cost) return result - def _translate_responses_api_response_to_model_response(self, result: ResponsesAPIResponse) -> ModelResponse: + def _translate_responses_api_response_to_model_response( + self, result: ResponsesAPIResponse, request: BridgedChatRequest | None = None + ) -> ModelResponse: """ Convert a Responses API response into a ModelResponse for spend_logs. @@ -5336,8 +5433,8 @@ class Logging(LiteLLMLoggingBaseClass): logging_obj=self, request_data={}, messages=[], - optional_params={}, - litellm_params={}, + optional_params={} if request is None else dict(request.optional_params), + litellm_params={} if request is None else dict(request.litellm_params), encoding=litellm.encoding, ) translated.id = served_id or translated.id diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 50c26f4d4ca..c553bfeaad8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1040,6 +1040,11 @@ def calculate_cache_writing_cost( class PromptTokensDetailsResult(TypedDict): cache_hit_tokens: int cache_hit_audio_tokens: ReadOnly[int] + + cached_text_tokens: ReadOnly[int] + cached_audio_tokens: ReadOnly[int] + cached_image_tokens: ReadOnly[int] + has_cached_tokens_details: ReadOnly[bool] cache_creation_tokens: int cache_creation_token_details: CacheCreationTokenDetails | None text_tokens: int @@ -1126,6 +1131,10 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, cache_hit_audio_tokens=cached_audio_tokens, + cached_text_tokens=cached_text_tokens, + cached_audio_tokens=cached_audio_tokens, + cached_image_tokens=cached_image_tokens, + has_cached_tokens_details=cached_tokens_details is not None, cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, text_tokens=text_tokens, @@ -1234,15 +1243,11 @@ def _calculate_input_cost( prompt_cost = 0.0 if bills_characters else float(prompt_tokens_details["text_tokens"]) * prompt_base_cost ### CACHE READ COST - Now uses tiered pricing - cache_hit_audio_tokens: Final = prompt_tokens_details["cache_hit_audio_tokens"] - audio_cache_read_rate: Final = get_cost_per_unit( - model_info, - get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier), - None, - ) - prompt_cost += float(prompt_tokens_details["cache_hit_tokens"] - cache_hit_audio_tokens) * cache_read_cost - prompt_cost += float(cache_hit_audio_tokens) * ( - audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost + prompt_cost += _calculate_cache_read_cost( + prompt_tokens_details=prompt_tokens_details, + model_info=model_info, + cache_read_cost=cache_read_cost, + service_tier=service_tier, ) ### AUDIO COST @@ -1331,6 +1336,38 @@ def _calculate_input_cost( return prompt_cost +def _calculate_cache_read_cost( + prompt_tokens_details: PromptTokensDetailsResult, + model_info: ModelInfo, + cache_read_cost: float, + service_tier: str | None, +) -> float: + cached_text_tokens: Final = prompt_tokens_details["cached_text_tokens"] + cached_audio_tokens: Final = prompt_tokens_details["cached_audio_tokens"] + cached_image_tokens: Final = prompt_tokens_details["cached_image_tokens"] + classified_cached_tokens: Final = cached_text_tokens + cached_audio_tokens + cached_image_tokens + unclassified_cached_tokens: Final = max(prompt_tokens_details["cache_hit_tokens"] - classified_cached_tokens, 0) + total_cost = ( # rebind-ok: cached modality components accumulate into one cache-read cost + float(cached_text_tokens + unclassified_cached_tokens) * cache_read_cost + ) + + if cached_audio_tokens: + cached_audio_cost_key: Final = get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier) + cached_audio_cost: Final = get_cost_per_unit(model_info, cached_audio_cost_key, cache_read_cost) + total_cost += ( # rebind-ok: cached audio contributes to cache-read cost + float(cached_audio_tokens) * float(cached_audio_cost or 0.0) + ) + + if cached_image_tokens: + cached_image_cost_key: Final = get_service_tier_cost_key("cache_read_input_image_token_cost", service_tier) + cached_image_cost: Final = get_cost_per_unit(model_info, cached_image_cost_key, cache_read_cost) + total_cost += ( # rebind-ok: cached images contribute to cache-read cost + float(cached_image_tokens) * float(cached_image_cost or 0.0) + ) + + return total_cost + + def get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str | None) -> float: """ Resolve the per-model regional-processing uplift multiplier for a given @@ -1502,6 +1539,10 @@ def generic_cost_per_token( prompt_tokens_details = PromptTokensDetailsResult( cache_hit_tokens=0, cache_hit_audio_tokens=0, + cached_text_tokens=0, + cached_audio_tokens=0, + cached_image_tokens=0, + has_cached_tokens_details=False, cache_creation_tokens=0, cache_creation_token_details=None, text_tokens=usage.prompt_tokens, @@ -1680,6 +1721,7 @@ class BilledTokenRates: cache_creation_input_token_cost: float cache_creation_input_token_cost_by_ttl: Mapping[str, float] output_cost_per_reasoning_token: float + cache_read_input_image_token_cost: float | None = None def scaled(self, multiplier: float) -> "BilledTokenRates": if multiplier == 1.0: @@ -1694,6 +1736,11 @@ class BilledTokenRates: {ttl: rate * multiplier for ttl, rate in self.cache_creation_input_token_cost_by_ttl.items()} ), output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier, + cache_read_input_image_token_cost=( + self.cache_read_input_image_token_cost * multiplier + if self.cache_read_input_image_token_cost is not None + else None + ), ) @@ -1802,6 +1849,11 @@ def _cost_map_billed_rates( cache_creation_input_token_cost=base_costs.cache_creation_rate, cache_creation_input_token_cost_by_ttl=base_costs.cache_creation_rates_by_ttl, output_cost_per_reasoning_token=reasoning_rate, + cache_read_input_image_token_cost=_get_cost_per_unit( + model_info, + _get_service_tier_cost_key("cache_read_input_image_token_cost", service_tier), + None, + ), ).scaled(multiplier) @@ -1882,6 +1934,12 @@ def get_token_type_cost_breakdown( cache_read_tokens, cached_audio_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts( usage ) + cached_image_tokens: Final = parse_prompt_tokens_details(usage)["cached_image_tokens"] + image_cache_read_rate: Final = ( + rates.cache_read_input_image_token_cost + if rates.cache_read_input_image_token_cost is not None + else rates.cache_read_input_token_cost + ) cache_creation_cost: Final = ( float(cache_creation_tokens) * rates.cache_creation_input_token_cost if custom_cost_per_token is not None @@ -1895,8 +1953,9 @@ def get_token_type_cost_breakdown( return TokenTypeCostBreakdown( reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token, cache_read_cost=( - float(cache_read_tokens - cached_audio_tokens) * rates.cache_read_input_token_cost + float(cache_read_tokens - cached_audio_tokens - cached_image_tokens) * rates.cache_read_input_token_cost + float(cached_audio_tokens) * rates.cache_read_input_audio_token_cost + + float(cached_image_tokens) * image_cache_read_rate ), cache_creation_cost=cache_creation_cost, rates=rates, diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 58fe5b4297f..9a0a35979db 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,8 @@ import json import re import time import traceback -from collections.abc import Generator, Iterator, Mapping, Sequence +from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import Final, Literal, cast from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config @@ -11,7 +12,12 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger -from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.constants import ( + RESPONSE_FORMAT_TOOL_NAME, + STREAM_CHUNK_CHOICE_FIELDS, + STREAM_FINISH_CHOICE_FIELDS, + VLLM_CHOICE_FIELDS, +) from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_reasoning_content, @@ -170,6 +176,28 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]: return _REPLAY_CONTENT_SLICE_RE.findall(content) +def _drop_choice_fields(choice: StreamingChoices, fields: Iterable[str]) -> None: + for field in fields: + if hasattr(choice, field): + delattr(choice, field) + + +def _vllm_choice_fields(choice: Mapping[str, object]) -> Mapping[str, object]: + cached_fields: Final = _choice_provider_specific_fields(choice) + return MappingProxyType( + { + field: choice[field] if field in choice else cached_fields[field] + for field in VLLM_CHOICE_FIELDS + if field in choice or field in cached_fields + } + ) + + +def _choice_provider_specific_fields(choice: Mapping[str, object]) -> Mapping[str, object]: + fields: Final = choice.get("provider_specific_fields") + return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else MappingProxyType({}) + + _REPLAY_REASONING_FIELDS: Final = frozenset({"reasoning_content", "thinking_blocks", "reasoning_items"}) _REPLAY_LEAD_FIELDS: Final = _REPLAY_REASONING_FIELDS | {"role", "provider_specific_fields"} _REPLAY_TOOL_FIELDS: Final = frozenset({"tool_calls", "function_call"}) @@ -220,8 +248,9 @@ def _replay_frame( choice.finish_reason = finish_reason if not keeps_choice_metadata: choice.logprobs = None - if hasattr(choice, "enhancements"): - del choice.enhancements + _drop_choice_fields(choice, ("enhancements", *STREAM_CHUNK_CHOICE_FIELDS)) + if finish_reason is None: + _drop_choice_fields(choice, STREAM_FINISH_CHOICE_FIELDS) return frame @@ -325,7 +354,9 @@ async def convert_to_streaming_response_async( logprobs = choice.get("logprobs", None) - choice = StreamingChoices(finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs) + choice = StreamingChoices( + finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs, **_vllm_choice_fields(choice) + ) choice_list.append(choice) model_response_object.choices = choice_list @@ -388,6 +419,7 @@ def convert_to_streaming_response( delta=delta, logprobs=logprobs, enhancements=enhancements, + **_vllm_choice_fields(choice), ) choice_list.append(choice) @@ -872,6 +904,7 @@ def convert_to_model_response_object( logprobs=logprobs, enhancements=enhancements, provider_specific_fields=provider_specific_fields, + **_vllm_choice_fields(choice), ) if provider_sent_logprobs and logprobs is None: choice.logprobs = None @@ -937,9 +970,11 @@ def convert_to_model_response_object( model_response_object.data = response_object["data"] if "usage" in response_object and response_object["usage"] is not None: - model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) - model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) - model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) + embedding_usage: Final = model_response_object.usage or Usage() + embedding_usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) + embedding_usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) + embedding_usage.total_tokens = response_object["usage"].get("total_tokens", 0) + model_response_object.usage = embedding_usage if start_time is not None and end_time is not None: model_response_object._response_ms = ( @@ -982,6 +1017,10 @@ def convert_to_model_response_object( if key in response_object: setattr(model_response_object, key, response_object[key]) + detected_languages: Final = TranscriptionResponse.detected_languages(response_object.get("languages")) + if detected_languages is not None: + model_response_object.languages = detected_languages + if "usage" in response_object and response_object["usage"] is not None: tr_usage_object: TranscriptionUsageDurationObject | TranscriptionUsageTokensObject | None = None diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index b08596e5244..3e82077f9d6 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -1,6 +1,6 @@ from collections.abc import Mapping from functools import lru_cache -from typing import ClassVar, Final +from typing import ClassVar, Final, get_args from openai.types.chat.completion_create_params import ( CompletionCreateParamsNonStreaming, @@ -21,6 +21,7 @@ from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.types.llms.anthropic import AnthropicMessagesRequest +from litellm.types.llms.openai import OpenAIAudioTranscriptionOptionalParams from litellm.types.rerank import RerankRequest @@ -178,6 +179,7 @@ class ModelParamHelper: This follows the OpenAI API Spec """ + litellm_transcription_kwargs: Final = set(get_args(OpenAIAudioTranscriptionOptionalParams)) try: from openai.types.audio.transcription_create_params import ( TranscriptionCreateParamsNonStreaming, @@ -187,11 +189,11 @@ class ModelParamHelper: non_streaming_kwargs = set(getattr(TranscriptionCreateParamsNonStreaming, "__annotations__", {}).keys()) streaming_kwargs: Final = set(getattr(TranscriptionCreateParamsStreaming, "__annotations__", {}).keys()) - all_transcription_kwargs: Final = non_streaming_kwargs.union(streaming_kwargs) + all_transcription_kwargs: Final = non_streaming_kwargs.union(streaming_kwargs, litellm_transcription_kwargs) return all_transcription_kwargs except Exception as e: verbose_logger.debug("Error getting transcription kwargs %s", str(e)) - return set() + return litellm_transcription_kwargs @staticmethod def _get_litellm_supported_responses_api_kwargs() -> set[str]: diff --git a/litellm/litellm_core_utils/pattern_match_deployments.py b/litellm/litellm_core_utils/pattern_match_deployments.py index 2b8d1e9c11b..a9091348c50 100644 --- a/litellm/litellm_core_utils/pattern_match_deployments.py +++ b/litellm/litellm_core_utils/pattern_match_deployments.py @@ -5,6 +5,7 @@ Class to handle llm wildcard routing and regex pattern matching import copy import itertools import re +from collections.abc import Mapping from re import Match from typing import Final @@ -22,6 +23,16 @@ class _RoutedDeployment(BaseModel): _ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...]) +_PATTERN_WIDE_CACHE_FLAGS: Final = frozenset(("supports_prompt_cache_breakpoint", "supports_prompt_caching")) +_AZURE_AI_PATTERN_CACHE_FLAGS: Final = frozenset(("supports_prompt_caching",)) + + +def _cache_flags_left_to_the_catalog(litellm_params: Mapping[str, object]) -> frozenset[str]: + model: Final = litellm_params.get("model") + provider: Final = litellm_params.get("custom_llm_provider") or inferred_provider( + model if isinstance(model, str) else None + ) + return _AZURE_AI_PATTERN_CACHE_FLAGS if provider == "azure_ai" else _PATTERN_WIDE_CACHE_FLAGS class PatternUtils: @@ -133,6 +144,11 @@ class PatternMatchRouter: matched_pattern=matched_pattern, litellm_deployment_litellm_model=deployment["litellm_params"]["model"], ) + if "*" in deployment["litellm_params"]["model"] and isinstance(new_deployment.get("model_info"), dict): + left_to_the_catalog = _cache_flags_left_to_the_catalog(deployment["litellm_params"]) + new_deployment["model_info"] = { + key: value for key, value in new_deployment["model_info"].items() if key not in left_to_the_catalog + } new_deployments.append(new_deployment) return new_deployments diff --git a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py index 60134456577..49ffae03c40 100644 --- a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py +++ b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py @@ -10,12 +10,18 @@ class ChatTemplateToolSection: preamble_tokens: int implicit_system_turn_tokens: int render_tool_schemas: Callable[[Sequence[Mapping[str, object]]], str] + counts_tool_choice: bool = False + renders_strict_only_when_true: bool = False -def _deepseek_v4_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: +def _json_line_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: return "\n".join(json.dumps(function, ensure_ascii=False) for function in functions) +def _compact_json_line_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: + return "\n".join(json.dumps(function, ensure_ascii=False, separators=(",", ":")) for function in functions) + + def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: return "".join( "\n" + json.dumps({"type": "function", "function": function}, ensure_ascii=False) for function in functions @@ -23,13 +29,29 @@ def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: _TOOL_SECTIONS: Final = ( + ( + re.compile(r"deepseek-v4(?:[-._]|p)1(?!\d)", re.IGNORECASE), + ChatTemplateToolSection(211, 0, _json_line_tool_schemas), + ), ( re.compile( r"deepseek-v4(?![.\d]|p\d)|^(?:deepseek/)?deepseek-(?:chat|reasoner|flash)(?::free)?$", re.IGNORECASE ), - ChatTemplateToolSection(222, 0, _deepseek_v4_tool_schemas), + ChatTemplateToolSection(222, 0, _json_line_tool_schemas), ), (re.compile(r"qwen3[._]?5(?!\d)", re.IGNORECASE), ChatTemplateToolSection(196, 5, _qwen3_5_tool_schemas)), + ( + re.compile(r"grok-4[-.]1(?![.\d])", re.IGNORECASE), + ChatTemplateToolSection( + 0, 0, _compact_json_line_tool_schemas, counts_tool_choice=True, renders_strict_only_when_true=True + ), + ), + ( + re.compile(r"grok-4[-.]3(?![.\d])", re.IGNORECASE), + ChatTemplateToolSection( + 0, 0, _json_line_tool_schemas, counts_tool_choice=True, renders_strict_only_when_true=True + ), + ), ) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 938ab585b72..6ecab166dcd 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -21,14 +21,17 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm import verbose_logger +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.batch_utils import InMemoryFile from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionAudioObject, + ChatCompletionCachedContent, ChatCompletionFileObject, ChatCompletionImageObject, + ChatCompletionPromptCacheBreakpoint, ChatCompletionReasoningItem, ChatCompletionReasoningSummaryTextBlock, ChatCompletionRedactedThinkingBlock, @@ -2182,6 +2185,9 @@ def check_is_function_call(logging_obj: "LoggingClass") -> bool: _MarkedT: Final = TypeVar("_MarkedT", bound=Mapping[str, object]) +PROMPT_CACHE_BREAKPOINT_DROP_PARAMS_HINT: Final = ( + "set `drop_params: true` on the request or in litellm_settings (`litellm.drop_params = True` in the SDK)" +) def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT: @@ -2193,16 +2199,62 @@ def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT: _PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(PromptCacheBreakpoint) +_CHAT_COMPLETION_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(ChatCompletionPromptCacheBreakpoint) +_AcceptedBreakpointT: Final = TypeVar("_AcceptedBreakpointT", bound=Mapping[str, object]) + + +def _accepted_prompt_cache_breakpoint( + accepted: TypeAdapter[_AcceptedBreakpointT], marker: object +) -> _AcceptedBreakpointT | None: + try: + return accepted.validate_python(marker) + except ValidationError: + return None + + +def prompt_cache_breakpoint_within( + accepted: TypeAdapter[_AcceptedBreakpointT], marker: object +) -> _AcceptedBreakpointT | None: + """The client's marker cut down to the ``accepted`` shape: an unsupported ttl is dropped and the mode kept, + and a marker whose mode the target cannot read is dropped whole.""" + whole: Final = _accepted_prompt_cache_breakpoint(accepted, marker) + if whole is not None: + return whole + fields: Final = _breakpoint_fields_or_none(marker) + without_ttl: Final = ( + None + if fields is None + else _accepted_prompt_cache_breakpoint(accepted, {key: item for key, item in fields.items() if key != "ttl"}) + ) + verbose_logger.debug( + "Dropping malformed prompt_cache_breakpoint %r under drop_params" + if without_ttl is None + else "Dropping the unsupported ttl of prompt_cache_breakpoint %r under drop_params", + marker, + ) + return without_ttl def prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object: if marker is None or not drop_params: return marker - try: - return _PROMPT_CACHE_BREAKPOINT.validate_python(marker) - except ValidationError: - verbose_logger.debug("Dropping malformed prompt_cache_breakpoint %r under drop_params", marker) - return None + return prompt_cache_breakpoint_within(_PROMPT_CACHE_BREAKPOINT, marker) + + +def prompt_cache_breakpoint_with_readable_mode(marker: object, drop_params: bool) -> object: + """The client's marker as written, or None under drop_params when no target could read its mode, so the + target's own config decides which ttl and keys it keeps.""" + return None if prompt_cache_breakpoint_for_wire(marker, drop_params) is None else marker + + +def whole_prompt_cache_breakpoint_or_none(marker: object) -> PromptCacheBreakpoint | None: + return _accepted_prompt_cache_breakpoint(_PROMPT_CACHE_BREAKPOINT, marker) + + +def chat_completion_prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object: + if marker is None or not drop_params: + return marker + return prompt_cache_breakpoint_within(_CHAT_COMPLETION_PROMPT_CACHE_BREAKPOINT, marker) def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT: @@ -2215,11 +2267,63 @@ def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT: ) +def prompt_cache_breakpoint_as_cache_control(block: _MarkedT, drop_params: bool) -> _MarkedT: + """OpenAI's ``prompt_cache_breakpoint`` mark rewritten as Anthropic's ``cache_control`` on the same block. + A ``cache_control`` the client already wrote on the block wins, TTL included.""" + marker: Final = block.get("prompt_cache_breakpoint") + if marker is None: + return block + unmarked: Final = {key: value for key, value in block.items() if key != "prompt_cache_breakpoint"} + control: Final = ( + None + if block.get("cache_control") is not None + else _cache_control_for_prompt_cache_breakpoint(marker, drop_params) + ) + marked: Final = unmarked if control is None else {**unmarked, "cache_control": control} + return cast(_MarkedT, marked) # cast-ok: same block shape as the input with the marker key renamed + + +def _cache_control_for_prompt_cache_breakpoint(marker: object, drop_params: bool) -> ChatCompletionCachedContent | None: + try: + breakpoint_mark: Final = _PROMPT_CACHE_BREAKPOINT.validate_python(marker) + except ValidationError: + if not drop_params: + raise UnsupportedParamsError( + message=( + f"prompt_cache_breakpoint must be {{'mode': 'explicit'}} with an optional ttl of '30m' or '12h', got " + f"{marker!r}. Fix the marker, or " + f"{PROMPT_CACHE_BREAKPOINT_DROP_PARAMS_HINT} to drop it." + ), + status_code=400, + ) + verbose_logger.debug("Dropping malformed prompt_cache_breakpoint %r under drop_params", marker) + return None + ttl: Final = breakpoint_mark.get("ttl") + if ttl is not None and not drop_params: + raise UnsupportedParamsError( + message=( + f"This model's prompt cache lasts 5m, or 1h where it offers one, so prompt_cache_breakpoint " + f"ttl={ttl!r} has no equivalent. Remove the ttl, {PROMPT_CACHE_BREAKPOINT_DROP_PARAMS_HINT} to send " + "a 5m breakpoint instead, or send cache_control {'type': 'ephemeral', 'ttl': '1h'} on the block " + "where the model offers a 1h cache." + ), + status_code=400, + ) + return ChatCompletionCachedContent(type="ephemeral") + + _BREAKPOINT_BLOCK_LIST_KEYS: Final = frozenset({"content", "tool_calls"}) _BREAKPOINT_FIELDS: Final = TypeAdapter(Mapping[str, object]) _BREAKPOINT_BLOCKS: Final = TypeAdapter(list[object]) +def _breakpoint_fields_or_none(value: object) -> Mapping[str, object] | None: + try: + return _BREAKPOINT_FIELDS.validate_python(value) + except ValidationError: + return None + + def _without_prompt_cache_breakpoint(value: object) -> object: if not isinstance(value, Mapping): return value @@ -2242,6 +2346,27 @@ def without_prompt_cache_breakpoints(block: _MarkedT) -> _MarkedT: return cast(_MarkedT, stripped) # cast-ok: same block shape as the input minus the marker keys +def _part_with_chat_completion_prompt_cache_breakpoint(part: object) -> object: + if not isinstance(part, Mapping): + return part + fields: Final = _BREAKPOINT_FIELDS.validate_python(part) + marker: Final = fields.get("prompt_cache_breakpoint") + if marker is None: + return fields + unmarked: Final = {key: item for key, item in fields.items() if key != "prompt_cache_breakpoint"} + return with_prompt_cache_breakpoint(unmarked, chat_completion_prompt_cache_breakpoint_for_wire(marker, True)) + + +def chat_completion_prompt_cache_breakpoints_for_wire(message: _MarkedT, drop_params: bool) -> _MarkedT: + content: Final = message.get("content") + if not drop_params or not isinstance(content, list): + return message + parts: Final = [ + _part_with_chat_completion_prompt_cache_breakpoint(part) for part in _BREAKPOINT_BLOCKS.validate_python(content) + ] + return cast(_MarkedT, {**message, "content": parts}) # cast-ok: same message shape with its parts' markers trimmed + + LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS: Final = frozenset( {"thinking_blocks", "reasoning_items", "provider_specific_fields"} ) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 98348e434b2..4c3907dc52f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import split_leading_system_run, string_field, ) +from litellm.litellm_core_utils.tool_fields import tools_without_hosted_mcp from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client, header_value from litellm.types.files import get_file_extension_from_mime_type @@ -76,7 +77,7 @@ from .common_utils import ( parse_tool_call_arguments, salvage_concatenated_tool_arguments, ) -from .image_handling import convert_url_to_base64 +from .image_handling import convert_url_to_base64, is_pdf_media_type def default_pt(messages): @@ -910,8 +911,17 @@ def convert_to_anthropic_image_obj(openai_image_url: str, format: str | None) -> ) +def _anthropic_https_url_block( + image_url: str, image_format: str | None +) -> AnthropicMessagesImageParam | AnthropicMessagesDocumentParam: + source: Final = AnthropicContentParamSourceUrl(type="url", url=image_url) + if image_format and is_pdf_media_type(image_format): + return AnthropicMessagesDocumentParam(type="document", source=source) + return AnthropicMessagesImageParam(type="image", source=source) + + def create_anthropic_image_param( - image_url_input: str | dict, + image_url_input: str | Mapping[str, object], format: str | None = None, is_bedrock_invoke: bool = False, ) -> AnthropicMessagesImageParam | AnthropicMessagesDocumentParam: @@ -920,13 +930,11 @@ def create_anthropic_image_param( Supports both URL references (for HTTP/HTTPS URLs) and base64 encoding. """ - # Extract URL and format from input - if isinstance(image_url_input, str): - image_url = image_url_input - else: - image_url = image_url_input.get("url", "") - if format is None: - format = image_url_input.get("format") + url_value: Final = image_url_input if isinstance(image_url_input, str) else image_url_input.get("url") + format_value: Final = None if isinstance(image_url_input, str) else image_url_input.get("format") + image_url: Final = url_value if isinstance(url_value, str) else "" + explicit_format: Final = format if format is not None else format_value + image_format: Final = explicit_format if isinstance(explicit_format, str) else None # Check if the image URL is an HTTP/HTTPS URL if image_url.startswith("http://") or image_url.startswith("https://"): @@ -934,21 +942,14 @@ def create_anthropic_image_param( # as these providers don't support URL sources for images if is_bedrock_invoke or image_url.startswith("http://"): base64_url: Final = convert_url_to_base64(url=image_url) - image_chunk = convert_to_anthropic_image_obj(openai_image_url=base64_url, format=format) - return _anthropic_content_element_factory(image_chunk) + image_chunk = convert_to_anthropic_image_obj(openai_image_url=base64_url, format=image_format) + return anthropic_content_element_factory(image_chunk) else: - # HTTPS URL - pass directly for regular Anthropic - return AnthropicMessagesImageParam( - type="image", - source=AnthropicContentParamSourceUrl( - type="url", - url=image_url, - ), - ) + return _anthropic_https_url_block(image_url, image_format) else: # Convert to base64 for data URIs or other formats - image_chunk = convert_to_anthropic_image_obj(openai_image_url=image_url, format=format) - return _anthropic_content_element_factory(image_chunk) + image_chunk = convert_to_anthropic_image_obj(openai_image_url=image_url, format=image_format) + return anthropic_content_element_factory(image_chunk) # The following XML functions will be deprecated once JSON schema support is available on Bedrock and Vertex @@ -1328,7 +1329,7 @@ def get_thought_signature_from_tool(tool: Mapping[str, object]) -> str | None: _get_thought_signature_from_tool = get_thought_signature_from_tool -def _get_dummy_thought_signature() -> str: +def get_dummy_thought_signature() -> str: """Generate a dummy thought signature for models that require it. This is used when transferring conversation history from older models @@ -1415,7 +1416,7 @@ def convert_to_gemini_tool_call_invoke( # https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/thinking/thought-signatures#parallel_function_calling_example is_first_function_call = len(_parts_list) == 0 if not thought_signature and is_first_function_call and needs_dummy_signature: - thought_signature = _get_dummy_thought_signature() + thought_signature = get_dummy_thought_signature() if thought_signature: part_dict["thoughtSignature"] = thought_signature @@ -1439,7 +1440,7 @@ def convert_to_gemini_tool_call_invoke( # If no signature found and model is gemini-3, use dummy signature if not thought_signature and needs_dummy_signature: - thought_signature = _get_dummy_thought_signature() + thought_signature = get_dummy_thought_signature() if thought_signature: part_dict_function["thoughtSignature"] = thought_signature @@ -2254,16 +2255,16 @@ def add_cache_control_to_content( return anthropic_content_element -def _anthropic_content_element_factory( +def anthropic_content_element_factory( image_chunk: GenericImageParsingChunk, ) -> AnthropicMessagesImageParam | AnthropicMessagesDocumentParam: - if image_chunk["media_type"] == "application/pdf": + if is_pdf_media_type(image_chunk["media_type"]): _anthropic_content_element: AnthropicMessagesDocumentParam | AnthropicMessagesImageParam = ( AnthropicMessagesDocumentParam( type="document", source=AnthropicContentParamSource( type="base64", - media_type=image_chunk["media_type"], + media_type="application/pdf", data=image_chunk["data"], ), ) @@ -4253,6 +4254,7 @@ def _convert_to_bedrock_tool_call_invoke( tool_calls: list, model: str | None = None, deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> list[BedrockContentBlock]: """ OpenAI tool invokes: @@ -4337,6 +4339,7 @@ def _convert_to_bedrock_tool_call_invoke( block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if _cache_point_block is not None: _parts_list.append(_cache_point_block) @@ -4357,6 +4360,7 @@ def _convert_to_bedrock_tool_call_invoke( block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if cache_point_block is not None: _parts_list.append(cache_point_block) @@ -5028,6 +5032,7 @@ class BedrockConverseMessagesProcessor: deployment_model: str | None = None, base_model: str | None = None, drop_params: bool = False, + model_info: Mapping[str, object] | None = None, ) -> list[BedrockMessageBlock]: from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking @@ -5090,6 +5095,7 @@ class BedrockConverseMessagesProcessor: block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if _cache_point_block is not None: _parts.append(_cache_point_block) @@ -5097,7 +5103,11 @@ class BedrockConverseMessagesProcessor: elif message_content and isinstance(message_content, str): _part = BedrockContentBlock(text=message_content) _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block( - message_block, block_type="content_block", model=model, deployment_model=deployment_model + message_block, + block_type="content_block", + model=model, + deployment_model=deployment_model, + model_info=model_info, ) user_content.append(_part) if _cache_point_block is not None: @@ -5148,6 +5158,7 @@ class BedrockConverseMessagesProcessor: block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if cache_point_block is not None: tool_content.append(cache_point_block) @@ -5213,6 +5224,7 @@ class BedrockConverseMessagesProcessor: block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if _cache_point_block is not None and not ( drop_params and _as_thinking_block(element) is not None @@ -5230,6 +5242,7 @@ class BedrockConverseMessagesProcessor: block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if _cache_point_block is not None: assistant_content.append(_cache_point_block) @@ -5238,7 +5251,7 @@ class BedrockConverseMessagesProcessor: if _tool_calls: assistant_content.extend( _convert_to_bedrock_tool_call_invoke( - _tool_calls, model=model, deployment_model=deployment_model + _tool_calls, model=model, deployment_model=deployment_model, model_info=model_info ) ) @@ -5433,6 +5446,7 @@ def _bedrock_converse_messages_pt( deployment_model: str | None = None, base_model: str | None = None, drop_params: bool = False, + model_info: Mapping[str, object] | None = None, ) -> list[BedrockMessageBlock]: """ Converts given messages from OpenAI format to Bedrock format @@ -5501,6 +5515,7 @@ def _bedrock_converse_messages_pt( block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if _cache_point_block is not None: _parts.append(_cache_point_block) @@ -5508,7 +5523,11 @@ def _bedrock_converse_messages_pt( elif message_content and isinstance(message_content, str): _part = BedrockContentBlock(text=message_content) _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block( - message_block, block_type="content_block", model=model, deployment_model=deployment_model + message_block, + block_type="content_block", + model=model, + deployment_model=deployment_model, + model_info=model_info, ) user_content.append(_part) if _cache_point_block is not None: @@ -5561,6 +5580,7 @@ def _bedrock_converse_messages_pt( block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if cache_point_block is not None: tool_content.append(cache_point_block) @@ -5627,6 +5647,7 @@ def _bedrock_converse_messages_pt( block_type="content_block", model=model, deployment_model=deployment_model, + model_info=model_info, ) if _cache_point_block is not None and not ( drop_params and _as_thinking_block(element) is not None @@ -5639,14 +5660,20 @@ def _bedrock_converse_messages_pt( assistant_content.append(BedrockContentBlock(text=_assistant_content)) # Add cache point block for assistant string content _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block( - assistant_message_block, block_type="content_block", model=model, deployment_model=deployment_model + assistant_message_block, + block_type="content_block", + model=model, + deployment_model=deployment_model, + model_info=model_info, ) if _cache_point_block is not None: assistant_content.append(_cache_point_block) _tool_calls = assistant_message_block.get("tool_calls", []) if _tool_calls: assistant_content.extend( - _convert_to_bedrock_tool_call_invoke(_tool_calls, model=model, deployment_model=deployment_model) + _convert_to_bedrock_tool_call_invoke( + _tool_calls, model=model, deployment_model=deployment_model, model_info=model_info + ) ) msg_i += 1 @@ -5702,7 +5729,10 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str: def add_cache_point_tool_block( - tool: dict, model: str | None = None, deployment_model: str | None = None + tool: dict, + model: str | None = None, + deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> BedrockToolBlock | None: from litellm.llms.bedrock.common_utils import ( bedrock_cache_point_ttl_supported, @@ -5710,7 +5740,7 @@ def add_cache_point_tool_block( ) cache_control: Final = tool.get("cache_control", None) - if cache_control is not None and bedrock_model_accepts_cache_points(model, deployment_model): + if cache_control is not None and bedrock_model_accepts_cache_points(model, deployment_model, model_info): cache_point: Final = cache_control.get("type", "ephemeral") if cache_point == "ephemeral": cache_point_block: Final[CachePointBlock] = {"type": "default"} @@ -5747,7 +5777,10 @@ def _is_bedrock_tool_block(tool: dict) -> bool: def _bedrock_tools_pt( - tools: list, model: str | None = None, deployment_model: str | None = None + tools: list, + model: str | None = None, + deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> list[BedrockToolBlock]: """ OpenAI tools looks like: @@ -5868,7 +5901,9 @@ def _bedrock_tools_pt( tool_block_list.append(tool_block) ## ADD CACHE POINT TOOL BLOCK ## - cache_point_tool_block = add_cache_point_tool_block(tool, model=model, deployment_model=deployment_model) + cache_point_tool_block = add_cache_point_tool_block( + tool, model=model, deployment_model=deployment_model, model_info=model_info + ) if cache_point_tool_block is not None: tool_block_list.append(cache_point_tool_block) @@ -5916,8 +5951,13 @@ def function_call_prompt( llm_provider: str, drop_params: bool, ) -> list[dict[str, object]]: + prompt_functions: Final = tools_without_hosted_mcp( + functions, model=model, route=llm_provider, drop_params=drop_params + ) + if not prompt_functions: + return messages header: Final = """To call a function, reply with JSON ONLY in this format {"name": "function_name", "arguments":{"argument_name": "argument_value"}}. Once a function result answers the request, reply to the user in plain text instead of calling a function again. The following functions are available to you:""" - function_prompt: Final = header + "".join(f"\n{function}\n" for function in functions) + function_prompt: Final = header + "".join(f"\n{function}\n" for function in prompt_functions) if not any(_is_system_message(message) for message in messages): return [*messages, {"role": "system", "content": function_prompt}] return [ diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 3287515c81d..ffc0e82e410 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -5,28 +5,33 @@ Helper functions to handle images passed in messages import asyncio import base64 import time -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from itertools import chain +from functools import partial +from itertools import accumulate, chain, takewhile, tee from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar +from urllib.parse import urlparse import httpx -from httpx import Response +from httpx import USE_CLIENT_DEFAULT, Response import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache -from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB +from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB, NON_LLM_CONNECTION_TIMEOUT from litellm.litellm_core_utils.prompt_templates.common_utils import infer_content_type_from_url_and_content from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20 +_Fetched = TypeVar("_Fetched") _FETCH_ATTEMPTS: Final = 3 in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) +pdf_probe_verdicts: Final = InMemoryCache() def _header_value(response: Response, name: str) -> str | None: @@ -209,6 +214,173 @@ def inline_remote_image_urls(media: RemoteMedia) -> bool: return media.part_type == "image_url" +_PDF_MEDIA_TYPES: Final = frozenset({"application/pdf", "application/x-pdf"}) +_GENERIC_MEDIA_TYPES: Final = frozenset({"application/octet-stream", "binary/octet-stream"}) +_PDF_SIGNATURE: Final = b"%PDF-" +_PDF_SIGNATURE_WINDOW: Final = 1024 +_PDF_PROBE_HEADERS: Final = MappingProxyType({"Range": f"bytes=0-{_PDF_SIGNATURE_WINDOW - 1}"}) +_PDF_PROBE_STATUSES: Final = frozenset({200, 206}) +_PDF_PROBE_FAILURES: Final = (httpx.HTTPError, httpx.InvalidURL, ValueError) +_ASYNC_PDF_PROBE_FAILURES: Final = (*_PDF_PROBE_FAILURES, asyncio.TimeoutError) +_PDF_PROBE_VERDICT_TTL: Final = 600 +_PDF_PROBE_FAILURE_TTL: Final = 60 +_DECODED_BODY_HEADERS: Final = frozenset({"content-encoding", "content-length", "transfer-encoding"}) + + +def _bare_media_type(media_type: str) -> str: + return media_type.split(";", 1)[0].strip().lower() + + +def is_pdf_media_type(media_type: str) -> bool: + return _bare_media_type(media_type) in _PDF_MEDIA_TYPES + + +def is_unformatted_pdf_link(url: str, image_format: str | None) -> bool: + lowered: Final = url.lower() + return not image_format and (urlparse(lowered).path.endswith(".pdf") or lowered.endswith(".pdf")) + + +def is_unformatted_pdf_link_image(media: RemoteMedia) -> bool: + image_format: Final = media.fields.get("format") + return inline_remote_image_urls(media) and is_unformatted_pdf_link( + media.url, image_format if isinstance(image_format, str) else None + ) + + +def _serves_pdf(response: Response) -> bool: + if response.status_code not in _PDF_PROBE_STATUSES: + return False + served_type: Final = _bare_media_type(_header_value(response, "Content-Type") or "") + if served_type and served_type not in _GENERIC_MEDIA_TYPES: + return served_type in _PDF_MEDIA_TYPES + return _PDF_SIGNATURE in response.content[:_PDF_SIGNATURE_WINDOW] + + +def _signature_window_chunks(chunks: Iterable[bytes], deadline: float) -> Iterator[bytes]: + reading, counting = tee(chunks) + read_so_far: Final = chain((0,), accumulate(map(len, counting))) + still_reading: Final = takewhile( + lambda total: total < _PDF_SIGNATURE_WINDOW and time.monotonic() < deadline, read_so_far + ) + return (chunk for _, chunk in zip(still_reading, reading)) + + +def _with_body_prefix(response: Response, prefix: bytes) -> Response: + return Response( + response.status_code, + headers=[(name, value) for name, value in response.headers.multi_items() if name not in _DECODED_BODY_HEADERS], + content=prefix, + request=response.request, + ) + + +@dataclass(frozen=True, slots=True) +class _BodyPrefixReader: + client: httpx.Client + deadline: float + + def get( + self, + url: str, + *, + headers: Mapping[str, str] | None = None, + follow_redirects: bool = False, + timeout: float | httpx.Timeout | None = None, + ) -> Response: + with self.client.stream( + "GET", + url, + headers=headers, + follow_redirects=follow_redirects, + timeout=USE_CLIENT_DEFAULT if timeout is None else timeout, + ) as response: + prefix: Final = b"".join(_signature_window_chunks(response.iter_bytes(), self.deadline)) + return _with_body_prefix(response, prefix[:_PDF_SIGNATURE_WINDOW]) + + +@dataclass(frozen=True, slots=True) +class _AsyncBodyPrefixReader: + client: httpx.AsyncClient + + async def get( + self, + url: str, + *, + headers: Mapping[str, str] | None = None, + follow_redirects: bool = False, + timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, + ) -> Response: + async with self.client.stream( + "GET", + url, + headers=headers, + follow_redirects=follow_redirects, + timeout=USE_CLIENT_DEFAULT if timeout is None else timeout, + ) as response: + return _with_body_prefix(response, await anext(response.aiter_bytes(max_response_bytes), b"")) + + +def _cached_pdf_verdict(url: str) -> bool | None: + cached: Final = pdf_probe_verdicts.get_cache(url) + return cached if isinstance(cached, bool) else None + + +def _remember_pdf_verdict(url: str, verdict: bool, ttl: int) -> bool: + pdf_probe_verdicts.set_cache(url, verdict, ttl=ttl) + return verdict + + +def _remember_probe_response(url: str, response: Response) -> bool: + ttl: Final = _PDF_PROBE_VERDICT_TTL if response.status_code in _PDF_PROBE_STATUSES else _PDF_PROBE_FAILURE_TTL + return _remember_pdf_verdict(url, _serves_pdf(response), ttl) + + +def _failed_pdf_probe(url: str, error: Exception) -> bool: + verbose_logger.debug("PDF probe of %s failed, sending it as an image link: %s", url, error) + return _remember_pdf_verdict(url, False, _PDF_PROBE_FAILURE_TTL) + + +def serves_pdf(url: str, timeout: float = NON_LLM_CONNECTION_TIMEOUT) -> bool: + if MAX_IMAGE_URL_DOWNLOAD_SIZE_MB == 0: + return False + cached: Final = _cached_pdf_verdict(url) + if cached is not None: + return cached + try: + response: Final = safe_get( + _BodyPrefixReader(litellm.module_level_client.client, deadline=time.monotonic() + timeout), + url, + headers=dict(_PDF_PROBE_HEADERS), + timeout=timeout, + ) + except _PDF_PROBE_FAILURES as e: + return _failed_pdf_probe(url, e) + return _remember_probe_response(url, response) + + +async def async_serves_pdf(url: str, timeout: float = NON_LLM_CONNECTION_TIMEOUT) -> bool: + if MAX_IMAGE_URL_DOWNLOAD_SIZE_MB == 0: + return False + cached: Final = _cached_pdf_verdict(url) + if cached is not None: + return cached + try: + response: Final = await asyncio.wait_for( + async_safe_get( + _AsyncBodyPrefixReader(litellm.module_level_aclient.client), + url, + headers=dict(_PDF_PROBE_HEADERS), + timeout=timeout, + max_response_bytes=_PDF_SIGNATURE_WINDOW, + ), + timeout, + ) + except _ASYNC_PDF_PROBE_FAILURES as e: + return _failed_pdf_probe(url, e) + return _remember_probe_response(url, response) + + def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None: if fields.get("type") != "image_url": return None @@ -292,15 +464,40 @@ def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Call return _inline(remote, data_url) if data_url is not None else part -def _inline_message( - message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool] -) -> AllMessageValues: +def _with_parts(message: AllMessageValues, rewrite: Callable[[object], object]) -> AllMessageValues: parts: Final = _content_parts(message) if not parts: return message - inlined_parts: Final = [_inline_part(part, data_urls, should_inline) for part in parts] - inlined_message: Final = {**message, "content": inlined_parts} - return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined + rewritten: Final = {**message, "content": [rewrite(part) for part in parts]} + return rewritten # pyright: ignore[reportReturnType] # the same message with some of its parts rewritten + + +def _inline_message( + message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool] +) -> AllMessageValues: + return _with_parts(message, partial(_inline_part, data_urls=data_urls, should_inline=should_inline)) + + +def _pdf_formatted_part(part: object, pdf_links: frozenset[str], should_probe: Callable[[RemoteMedia], bool]) -> object: + remote: Final = _parse_remote_part(part) + match remote: + case _RemoteImage(fields, image_url, url) if url in pdf_links and should_probe(_remote_media(remote)): + image_url_fields: Final = image_url if image_url is not None else _NO_FIELDS + return {**fields, "image_url": {**image_url_fields, "url": url, "format": "application/pdf"}} + case _: + return part + + +def _with_pdf_formats( + messages: Iterable[AllMessageValues], + probed: Iterable[tuple[str, bool]], + should_probe: Callable[[RemoteMedia], bool], +) -> tuple[AllMessageValues, ...]: + pdf_links: Final = frozenset(url for url, is_pdf in probed if is_pdf) + if not pdf_links: + return tuple(messages) + rewrite: Final = partial(_pdf_formatted_part, pdf_links=pdf_links, should_probe=should_probe) + return tuple(_with_parts(message, rewrite) for message in messages) async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore, timeout: float | httpx.Timeout | None) -> str: @@ -308,11 +505,12 @@ async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore, timeout: float return await async_convert_url_to_base64(url, timeout=timeout) -async def async_convert_urls_to_base64( - remote_urls: tuple[str, ...], timeout: float | httpx.Timeout | None = None -) -> tuple[str, ...]: - in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) - fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight, timeout)) for url in remote_urls) +async def _serves_pdf_within(url: str, in_flight: asyncio.Semaphore) -> bool: + async with in_flight: + return await async_serves_pdf(url) + + +async def _gather_or_cancel_all(fetches: tuple["asyncio.Task[_Fetched]", ...]) -> tuple[_Fetched, ...]: try: return tuple(await asyncio.gather(*fetches)) except BaseException: @@ -322,6 +520,15 @@ async def async_convert_urls_to_base64( raise +async def async_convert_urls_to_base64( + remote_urls: tuple[str, ...], timeout: float | httpx.Timeout | None = None +) -> tuple[str, ...]: + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + return await _gather_or_cancel_all( + tuple(asyncio.create_task(_fetch_data_url(url, in_flight, timeout)) for url in remote_urls) + ) + + def _remote_urls_to_inline( messages: Iterable[AllMessageValues], should_inline: Callable[[RemoteMedia], bool] ) -> tuple[str, ...]: @@ -351,3 +558,27 @@ async def async_inline_remote_media( data_urls: Final = await async_convert_urls_to_base64(remote_urls) inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True))) return [_inline_message(message, inlined, should_inline) for message in messages] + + +def format_pdf_links( + messages: Sequence[AllMessageValues], should_probe: Callable[[RemoteMedia], bool] +) -> tuple[AllMessageValues, ...]: + links: Final = _remote_urls_to_inline(messages, should_probe) + if not links: + return tuple(messages) + with ThreadPoolExecutor(max_workers=min(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES, len(links))) as probes: + verdicts: Final = tuple(probes.map(serves_pdf, links)) + return _with_pdf_formats(messages, zip(links, verdicts, strict=True), should_probe) + + +async def async_format_pdf_links( + messages: Sequence[AllMessageValues], should_probe: Callable[[RemoteMedia], bool] +) -> tuple[AllMessageValues, ...]: + links: Final = _remote_urls_to_inline(messages, should_probe) + if not links: + return tuple(messages) + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + verdicts: Final = await _gather_or_cancel_all( + tuple(asyncio.create_task(_serves_pdf_within(url, in_flight)) for url in links) + ) + return _with_pdf_formats(messages, zip(links, verdicts, strict=True), should_probe) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 46f5ca87a93..8e6e71a31d3 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,4 +1,6 @@ import asyncio +import base64 +import functools import json import time import traceback @@ -8,7 +10,7 @@ from enum import Enum, auto from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast, runtime_checkable -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly import litellm @@ -19,9 +21,12 @@ from litellm.constants import ( REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS, REALTIME_LAZY_SETUP_DEADLINE_SECONDS, REALTIME_METERED_USAGE_KEY, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) +from litellm.litellm_core_utils.audio_utils.utils import normalized_audio_duration_seconds +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError from litellm.llms.base_llm.realtime.transformation import ( @@ -35,6 +40,8 @@ from litellm.types.llms.openai import ( OpenAIRealtimeResponseDelta, OpenAIRealtimeStreamResponseBaseObject, OpenAIRealtimeStreamSessionEvents, + OpenAIRealtimeTranslationClosedEvent, + OpenAIRealtimeTranslationDurationUsage, ) from litellm.types.realtime import ALL_DELTA_TYPES @@ -42,9 +49,11 @@ from .litellm_logging import Logging as LiteLLMLogging from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason if TYPE_CHECKING: + from openai.types.realtime.realtime_audio_formats import RealtimeAudioFormats from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks CLIENT_CONNECTION_CLASS = ClientConnection @@ -57,6 +66,10 @@ class BackendClose: code: int reason: str + @property + def ended_cleanly(self) -> bool: + return self.code == 1000 + @property def message(self) -> str: if not self.reason: @@ -67,6 +80,7 @@ class BackendClose: class ClientLoopExit(Enum): CLIENT_DISCONNECTED = auto() BACKEND_CLOSED = auto() + RELAY_FAILED = auto() def backend_close_from(error: "ConnectionClosed") -> BackendClose: @@ -96,6 +110,7 @@ class _ClientEventFrame(TypedDict, total=False): type: ReadOnly[str] item: ReadOnly[_ClientEventItem] session: ReadOnly[Mapping[str, object]] + response: ReadOnly[Mapping[str, object]] class _ResponseDoneBody(TypedDict, total=False): @@ -212,6 +227,87 @@ def _mapping_field(container: Mapping[str, object], key: str) -> Mapping[str, ob return _json_object(container.get(key)) +_BETA_AUDIO_FORMAT_BY_GA_TYPE: Final[Mapping[str, str]] = MappingProxyType( + {"audio/pcm": "pcm16", "audio/pcmu": "g711_ulaw", "audio/pcma": "g711_alaw"} +) +_GA_TO_BETA_TOP_LEVEL_SESSION_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"output_modalities": "modalities", "max_output_tokens": "max_response_output_tokens"} +) +_GA_TO_BETA_AUDIO_INPUT_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + { + "format": "input_audio_format", + "transcription": "input_audio_transcription", + "noise_reduction": "input_audio_noise_reduction", + "turn_detection": "turn_detection", + } +) +_GA_TO_BETA_AUDIO_OUTPUT_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"format": "output_audio_format", "voice": "voice", "speed": "speed"} +) +_GA_ONLY_SESSION_FIELDS: Final = frozenset(("type", "audio", *_GA_TO_BETA_TOP_LEVEL_SESSION_FIELDS)) + + +def _beta_session_value(ga_field: str, value: object) -> object: + if ga_field == "output_modalities": + return ["audio", "text"] if value == ["audio"] else value + if ga_field != "format": + return value + ga_format: Final = _json_object(value) + format_type: Final = None if ga_format is None else ga_format.get("type") + return _BETA_AUDIO_FORMAT_BY_GA_TYPE.get(format_type, value) if isinstance(format_type, str) else value + + +def _beta_session_fields(ga_fields: Mapping[str, object], beta_names: Mapping[str, str]) -> Mapping[str, object]: + return { + beta_names[name]: _beta_session_value(name, value) for name, value in ga_fields.items() if name in beta_names + } + + +def ga_session_to_beta(session: Mapping[str, object]) -> Mapping[str, object]: + audio: Final = _mapping_field(session, "audio") or _EMPTY_SESSION + return { + **{key: value for key, value in session.items() if key not in _GA_ONLY_SESSION_FIELDS}, + **_beta_session_fields(session, _GA_TO_BETA_TOP_LEVEL_SESSION_FIELDS), + **_beta_session_fields(_mapping_field(audio, "input") or _EMPTY_SESSION, _GA_TO_BETA_AUDIO_INPUT_FIELDS), + **_beta_session_fields(_mapping_field(audio, "output") or _EMPTY_SESSION, _GA_TO_BETA_AUDIO_OUTPUT_FIELDS), + } + + +_GA_ONLY_RESPONSE_FIELDS: Final = frozenset(("audio", "output_modalities")) +_GA_TO_BETA_RESPONSE_FIELDS: Final[Mapping[str, str]] = MappingProxyType({"output_modalities": "modalities"}) + + +def ga_response_to_beta(response: Mapping[str, object]) -> Mapping[str, object]: + audio: Final = _mapping_field(response, "audio") or _EMPTY_SESSION + return { + **{key: value for key, value in response.items() if key not in _GA_ONLY_RESPONSE_FIELDS}, + **_beta_session_fields(response, _GA_TO_BETA_RESPONSE_FIELDS), + **_beta_session_fields(_mapping_field(audio, "output") or _EMPTY_SESSION, _GA_TO_BETA_AUDIO_OUTPUT_FIELDS), + } + + +_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...]) +_GA_ASSISTANT_CONTENT_TYPE_BY_BETA_TYPE: Final[Mapping[str, str]] = MappingProxyType( + {"text": "output_text", "audio": "output_audio"} +) + + +def _ga_assistant_content_block(block: object) -> object: + content_block: Final = _json_object(block) + block_type: Final = None if content_block is None else content_block.get("type") + if content_block is None or not isinstance(block_type, str): + return block + ga_type: Final = _GA_ASSISTANT_CONTENT_TYPE_BY_BETA_TYPE.get(block_type) + return block if ga_type is None else {**content_block, "type": ga_type} + + +def beta_item_to_ga(item: Mapping[str, object]) -> Mapping[str, object]: + content: Final = item.get("content") + if item.get("role") != "assistant" or not isinstance(content, list): + return item + return {**item, "content": [_ga_assistant_content_block(block) for block in _JSON_ARRAY.validate_python(content)]} + + def _guardrail_event_id() -> str: return f"{_GUARDRAIL_EVENT_ID_PREFIX}{uuid.uuid4().hex}" @@ -270,6 +366,26 @@ def _gate_session_update( return {**client_session, "turn_detection": gate(None)} +@functools.cache +def _translation_audio_format_adapter() -> "TypeAdapter[str | RealtimeAudioFormats | None]": + from openai.types.realtime.realtime_audio_formats import RealtimeAudioFormats + + return TypeAdapter(str | RealtimeAudioFormats | None) + + +def _translation_audio_bytes_per_second(audio_format: "str | RealtimeAudioFormats | None") -> float | None: + from openai.types.realtime.realtime_audio_formats import AudioPCM, AudioPCMA, AudioPCMU + + match audio_format: + case "pcm16" | AudioPCM(): + return 48000.0 + case "g711_ulaw" | "g711_alaw" | AudioPCMU() | AudioPCMA(): + return 8000.0 + case _: + pass + return None + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... def normalize(self, event: dict) -> dict: ... @@ -294,6 +410,79 @@ DefaultLoggedRealTimeEventTypes: Final = [ ] +TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE: Final = "Translation sessions cannot enforce configured realtime guardrails" + + +def _guardrails_run_for_event_hooks( + request_data: Mapping[str, object], + event_hooks: Sequence["GuardrailEventHooks"], + callbacks: Sequence[object], +) -> bool: + from litellm.integrations.custom_guardrail import CustomGuardrail + + return any( + isinstance(cb, CustomGuardrail) + and any(cb.should_run_guardrail(data={**request_data, "stream": True}, event_type=et) for et in event_hooks) + for cb in callbacks + ) + + +def _opted_out_of_always_on_guardrail(callback: object, request_data: Mapping[str, object]) -> bool: + from litellm.integrations.custom_guardrail import CustomGuardrail + + return ( + isinstance(callback, CustomGuardrail) + and callback.event_hook is None + and callback.get_disable_global_guardrail(dict(request_data)) is True + ) + + +def _realtime_runner_guardrail( + callback: object, + check_data: Mapping[str, object], + event_hooks: Sequence["GuardrailEventHooks"], +) -> "CustomGuardrail | None": + from litellm.integrations.custom_guardrail import CustomGuardrail + + if not isinstance(callback, CustomGuardrail) or callback.use_native_lifecycle_hooks: + return None + runs: Final = any( + callback.should_run_guardrail(data=check_data, event_type=event_hook) for event_hook in event_hooks + ) + return callback if runs else None + + +def translation_guardrails_apply(request_data: Mapping[str, object]) -> bool: + from litellm.types.guardrails import GuardrailEventHooks + + callbacks: Final = tuple( + callback for callback in litellm.callbacks if not _opted_out_of_always_on_guardrail(callback, request_data) + ) + if _guardrails_run_for_event_hooks(request_data, (GuardrailEventHooks.post_call,), callbacks): + return True + if any( + isinstance(bucket, Mapping) and "_guardrail_pipelines" in bucket + for bucket in (request_data.get("metadata"), request_data.get("litellm_metadata")) + ): + from litellm.proxy.utils import pipeline_managed_guardrail_names + + if pipeline_managed_guardrail_names(request_data, "pre_call") or pipeline_managed_guardrail_names( + request_data, "post_call" + ): + return True + return realtime_guardrails_apply(request_data, callbacks) + + +def realtime_guardrails_apply(request_data: Mapping[str, object], callbacks: Sequence[object]) -> bool: + from litellm.types.guardrails import GuardrailEventHooks + + check_data: Final = {**request_data, "stream": True} + runner_event_hooks: Final = (GuardrailEventHooks.realtime_input_transcription, GuardrailEventHooks.pre_call) + return any( + _realtime_runner_guardrail(callback, check_data, runner_event_hooks) is not None for callback in callbacks + ) + + class RealTimeStreaming: def __init__( self, @@ -303,12 +492,15 @@ class RealTimeStreaming: provider_config: BaseRealtimeConfig | None = None, model: str = "", user_api_key_dict: object | None = None, - request_data: dict | None = None, + request_data: Mapping[str, object] | None = None, backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, usage_meter: RealtimeUsageMeter | None = None, logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, + translation_session: bool = False, + bridges_beta_client_to_ga: bool = False, + drop_params: bool = False, lazy_setup_deadline_seconds: float = REALTIME_LAZY_SETUP_DEADLINE_SECONDS, lazy_setup_content_grace_seconds: float = REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS, clock: Callable[[], float] = time.monotonic, @@ -326,14 +518,24 @@ class RealTimeStreaming: self._logging_worker = logging_worker self.messages: list[OpenAIRealtimeEvents] = [] self._backend_sent_frames: bool = False + self._backend_sent_non_error_frames: bool = False self.input_message: dict = {} self.input_messages: list[dict[str, str]] = [] self.session_tools: list[dict] = [] self.tool_calls: list[dict] = [] + self._is_translation_session = translation_session + self._bridges_beta_client_to_ga: Final = bridges_beta_client_to_ga + self._remaps_all_beta_fields_to_ga: Final = bridges_beta_client_to_ga or translation_session + self._drop_params: Final = drop_params + self._translation_input_seconds = 0.0 + self._translation_input_bytes_per_second = 48000.0 + self._translation_output_seconds = 0.0 + self._translation_output_bytes_per_second = 48000.0 + self._translation_usage_finalized = False # Detect whether the client is explicitly opting into the beta protocol. self._client_wants_beta = self.detect_beta_header(websocket) - self._backend_uses_beta_protocol = ( + self._backend_uses_beta_protocol = not translation_session and ( self._client_wants_beta if backend_uses_beta_protocol is None else backend_uses_beta_protocol ) @@ -352,7 +554,7 @@ class RealTimeStreaming: self.current_delta_type: ALL_DELTA_TYPES | None = None self.session_configuration_request: str | None = None self.user_api_key_dict = user_api_key_dict - self.request_data: dict = request_data or {} + self.request_data: Mapping[str, object] = request_data or {} # Violation counter for end_session_after_n_fails support self._violation_count: int = 0 # When a text message is blocked, hold the guardrail reason so the next @@ -384,6 +586,7 @@ class RealTimeStreaming: # their input_audio_transcription.completed usage drives duration-based cost. self._force_transcription_model = force_transcription_model self._is_transcription_session: bool = force_transcription_model is not None + self._bound_nested_transcription_model: str | None = None # Optional per-provider GA event normalizer (e.g. XAIRealtimeNormalizer). self._event_normalizer = event_normalizer self._gate_turn_detection: Final[_TurnDetectionGate] = ( @@ -407,6 +610,7 @@ class RealTimeStreaming: _MAX_BUFFERED_BYTES: int = 10 * 1024 * 1024 # 10 MB _SESSION_EVENT_TYPES = frozenset(["session.created", "session.updated"]) + _RESPONSE_EVENT_TYPES: Final = frozenset(("response.created", "response.done")) _CLIENT_AUDIO_BUFFER_TYPES = frozenset( [ "input_audio_buffer.append", @@ -429,8 +633,8 @@ class RealTimeStreaming: _DEFERRED_SETUP_HELD_CONTENT_TYPES: Final = frozenset(["conversation.item.create"]) _AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = { "pcm16": {"type": "audio/pcm", "rate": 24000}, - "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000}, - "g711_alaw": {"type": "audio/G711-alaw", "rate": 8000}, + "g711_ulaw": {"type": "audio/pcmu"}, + "g711_alaw": {"type": "audio/pcma"}, } # GA name → beta name (when client WebSocket includes OpenAI-Beta: realtime=v1) _GA_TO_BETA_EVENT_TYPES: dict[str, str] = { @@ -449,12 +653,12 @@ class RealTimeStreaming: def _should_store_message( self, - message_obj: dict | OpenAIRealtimeEvents, + message_obj: Mapping[str, object], ) -> bool: _msg_type: Final = message_obj["type"] if "type" in message_obj else None if self.logged_real_time_event_types == "*": return True - if _msg_type and _msg_type in self.logged_real_time_event_types: + if isinstance(_msg_type, str) and _msg_type in self.logged_real_time_event_types: return True return False @@ -555,7 +759,7 @@ class RealTimeStreaming: input_messages, so it is not duplicated into the response log here. """ try: - usage: Final = event_obj.get("usage") + usage: Final = as_str_mapping(event_obj.get("usage")) if usage is None: return # If this event type is already captured by store_message (e.g. the user @@ -623,6 +827,7 @@ class RealTimeStreaming: async def log_messages(self): """Log messages in list""" + self._finalize_translation_usage() if self.logging_obj: if self._usage_meter is not None: self.logging_obj.model_call_details[REALTIME_METERED_USAGE_KEY] = self._usage_meter.snapshot() @@ -641,6 +846,142 @@ class RealTimeStreaming: ) self.logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + def _capture_translation_output_audio(self, event_obj: Mapping[str, object]) -> None: + if not self._is_translation_session: + return + if event_obj.get("type") == "session.closed": + usage: Final = as_str_mapping(event_obj.get("usage")) + output_seconds: Final = ( + normalized_audio_duration_seconds(usage.get("output_seconds")) if usage is not None else None + ) + reported_input_seconds: Final = ( + normalized_audio_duration_seconds(usage.get("input_seconds")) if usage is not None else None + ) + input_seconds: Final = ( + reported_input_seconds + if reported_input_seconds is not None + else self._translation_input_seconds or None + ) + synthetic_output_seconds: Final = self._translation_output_seconds or None + resolved_output_seconds: Final = output_seconds if output_seconds is not None else synthetic_output_seconds + if input_seconds is not None or resolved_output_seconds is not None: + if self._should_store_message(event_obj): + supplemental_usage: Final = OpenAIRealtimeTranslationDurationUsage( + type="duration", + input_seconds=float(input_seconds or 0.0) if reported_input_seconds is None else 0.0, + output_seconds=float(synthetic_output_seconds or 0.0) if output_seconds is None else 0.0, + ) + if ( + supplemental_usage.get("input_seconds", 0.0) > 0 + or supplemental_usage.get("output_seconds", 0.0) > 0 + ): + self.messages.append( + OpenAIRealtimeTranslationClosedEvent( + type="session.closed", + usage=supplemental_usage, + ) + ) + else: + normalized_usage: Final = ( + OpenAIRealtimeTranslationDurationUsage( + type="duration", + input_seconds=input_seconds, + output_seconds=float(resolved_output_seconds or 0.0), + ) + if input_seconds is not None + else OpenAIRealtimeTranslationDurationUsage( + type="duration", output_seconds=float(resolved_output_seconds or 0.0) + ) + ) + self.messages.append( + OpenAIRealtimeTranslationClosedEvent(type="session.closed", usage=normalized_usage) + ) + self._translation_usage_finalized = True + return + self._capture_translation_audio_formats(event_obj) + if event_obj.get("type") not in ( + "session.output_audio.delta", + "response.output_audio.delta", + "response.audio.delta", + ): + return + delta: Final = event_obj.get("delta") + if not isinstance(delta, str): + return + try: + decoded: Final = base64.b64decode(delta, validate=True) + except (ValueError, TypeError): + return + self._translation_output_seconds += len(decoded) / self._translation_output_bytes_per_second + + def _capture_translation_input_audio(self, message: str) -> None: + if not self._is_translation_session: + return + try: + event: Final = _decode_json_object(message) + except (json.JSONDecodeError, TypeError): + return + if event.get("type") != "session.input_audio_buffer.append" or not isinstance(audio := event.get("audio"), str): + return + try: + decoded: Final = base64.b64decode(audio, validate=True) + except (ValueError, TypeError): + return + self._translation_input_seconds += len(decoded) / self._translation_input_bytes_per_second + + def _capture_translation_audio_formats(self, event_obj: Mapping[str, object]) -> None: + if event_obj.get("type") not in ("session.created", "session.updated"): + return + session: Final = as_str_mapping(event_obj.get("session")) + if session is None: + return + audio: Final = as_str_mapping(session.get("audio")) + if audio is None: + return + for direction, raw_configuration in (("input", audio.get("input")), ("output", audio.get("output"))): + if (configuration := as_str_mapping(raw_configuration)) is None: + continue + try: + if ( + bytes_per_second := _translation_audio_bytes_per_second( + _translation_audio_format_adapter().validate_python(configuration.get("format")) + ) + ) is None: + continue + except ValidationError: + continue + if direction == "input": + self._translation_input_bytes_per_second = bytes_per_second + else: + self._translation_output_bytes_per_second = bytes_per_second + + def _finalize_translation_usage(self) -> None: + if self._translation_usage_finalized: + return + for event in self.messages: + if event.get("type") != "session.closed": + continue + if (event_usage := as_str_mapping(event.get("usage"))) is not None and normalized_audio_duration_seconds( + event_usage.get("output_seconds") + ) is not None: + self._translation_usage_finalized = True + return + if self._translation_output_seconds == 0 and self._translation_input_seconds == 0: + return + synthetic_usage: Final = ( + OpenAIRealtimeTranslationDurationUsage( + type="duration", + input_seconds=self._translation_input_seconds, + output_seconds=self._translation_output_seconds, + ) + if self._translation_input_seconds > 0 + else OpenAIRealtimeTranslationDurationUsage( + type="duration", output_seconds=self._translation_output_seconds + ) + ) + self.messages.append(OpenAIRealtimeTranslationClosedEvent(type="session.closed", usage=synthetic_usage)) + self._translation_usage_finalized = True + async def _send_to_backend(self, message: str) -> bool: """Send a message to the backend WebSocket. @@ -653,10 +994,25 @@ class RealTimeStreaming: backend, False if the provider transformation produced no output and the message was effectively dropped. """ - message = self._enforce_transcription_session_model(message) + from litellm.proxy._types import ProxyException + + try: + authorized_message: Final = await self._apply_nested_transcription_model_policy(message) + except ProxyException as denied: + denial_message: Final = redact_internal_details_from_client_message(denied.message) + try: + await self.websocket.send_text(realtime_error_event(denial_message, error_type="permission_error")) + finally: + await self.websocket.close( + code=1008, + reason=websocket_close_reason( + denial_message, fallback="Nested transcription model is not authorized" + ), + ) + return False if self.provider_config: transformed: Final = self.provider_config.transform_realtime_request( - message, self.model, self.session_configuration_request + authorized_message, self.model, self.session_configuration_request ) sent = False for msg in transformed: @@ -691,14 +1047,94 @@ class RealTimeStreaming: if is_content_message: self._content_sent_after_setup = True sent = True + if sent: + self._capture_translation_input_audio(authorized_message) return sent - await self.backend_ws.send(message) + await self.backend_ws.send(authorized_message) if self._usage_meter is not None: - sent_event: Final = self._parse_backend_event(message) + sent_event: Final = self._parse_backend_event(authorized_message) if sent_event is not None: self._usage_meter.observe_client_event(sent_event) + self._capture_translation_input_audio(authorized_message) return True + async def _apply_nested_transcription_model_policy(self, message: str) -> str: + if self._force_transcription_model is not None: + return self._enforce_transcription_session_model(message) + if self._is_translation_session: + return await self._enforce_translation_nested_transcription_model(message) + return message + + def _session_update_message_obj(self, message: str) -> Mapping[str, object] | None: + try: + message_obj: Final = _decode_json_object(message) + except (json.JSONDecodeError, TypeError): + return None + if message_obj.get("type") not in ( + "session.update", + "transcription_session.update", + ): + return None + return message_obj + + def _nested_transcription_models_from_session( + self, + session: Mapping[str, object], + ) -> tuple[str, ...]: + audio: Final = as_str_mapping(session.get("audio")) + audio_input: Final = as_str_mapping(audio.get("input")) if audio is not None else None + nested_transcription: Final = audio_input.get("transcription") if audio_input is not None else None + nested_model: Final = self._transcription_model_value(nested_transcription) + flat_model: Final = self._transcription_model_value(session.get("input_audio_transcription")) + return tuple(dict.fromkeys(model for model in (nested_model, flat_model) if model is not None)) + + def _transcription_model_value(self, transcription_config: object) -> str | None: + if (fields := as_str_mapping(transcription_config)) is None: + return None + model: Final = fields.get("model") + if isinstance(model, str) and model: + return model + return None + + def _rewrite_session_update_transcription_model(self, message: str, authorized_model: str) -> str: + message_obj: Final = self._session_update_message_obj(message) + if message_obj is None: + return message + session: Final = as_str_mapping(message_obj.get("session")) + if session is None: + return message + transcription: Final = as_str_mapping(session.get("input_audio_transcription")) + rewrite_flat: Final = transcription is not None and transcription.get("model") != authorized_model + audio: Final = as_str_mapping(session.get("audio")) + audio_input: Final = as_str_mapping(audio.get("input")) if audio is not None else None + nested_transcription: Final = ( + as_str_mapping(audio_input.get("transcription")) if audio_input is not None else None + ) + rewrite_nested: Final = ( + nested_transcription is not None and nested_transcription.get("model") != authorized_model + ) + if not rewrite_flat and not rewrite_nested: + return message + updated_session: Final = { + **session, + **( + {"input_audio_transcription": {**transcription, "model": authorized_model}} + if transcription is not None and rewrite_flat + else {} + ), + **( + { + "audio": { + **audio, + "input": {**audio_input, "transcription": {**nested_transcription, "model": authorized_model}}, + } + } + if audio is not None and audio_input is not None and nested_transcription is not None and rewrite_nested + else {} + ), + } + return json.dumps({**message_obj, "session": updated_session}) + def _builds_setup_lazily(self) -> bool: if self.provider_config is None: return False @@ -877,56 +1313,52 @@ class RealTimeStreaming: if self._force_transcription_model is None: return message - try: - message_obj: Final = _decode_json_object(message) - except (json.JSONDecodeError, TypeError): + message_obj: Final = self._session_update_message_obj(message) + if message_obj is None: return message - - if message_obj.get("type") not in ( - "session.update", - "transcription_session.update", - ): - return message - - session: Final = message_obj.get("session") - if not isinstance(session, dict): - return message - - if session.get("type") == "transcription": + session: Final = as_str_mapping(message_obj.get("session")) + if session is not None and session.get("type") == "transcription": self._is_transcription_session = True + return self._rewrite_session_update_transcription_model(message, self._force_transcription_model) - authorized_model: Final = self._force_transcription_model - changed = False + async def _enforce_translation_nested_transcription_model(self, message: str) -> str: + message_obj: Final = self._session_update_message_obj(message) + session: Final = as_str_mapping(message_obj.get("session")) if message_obj is not None else None + bound_model: Final = self._bound_nested_transcription_model + requested_models: Final = ( + tuple(model for model in self._nested_transcription_models_from_session(session) if model != bound_model) + if session is not None + else () + ) + if not requested_models: + return ( + message + if bound_model is None + else self._rewrite_session_update_transcription_model(message, bound_model) + ) - transcription: Final = session.get("input_audio_transcription") - if isinstance(transcription, dict) and transcription.get("model") != authorized_model: - session["input_audio_transcription"] = { - **transcription, - "model": authorized_model, - } - changed = True + valid_token: Final = self.user_api_key_dict - audio: Final = session.get("audio") - if isinstance(audio, dict): - audio_input: Final = audio.get("input") - if isinstance(audio_input, dict): - nested_transcription: Final = audio_input.get("transcription") - if isinstance(nested_transcription, dict) and nested_transcription.get("model") != authorized_model: - session["audio"] = { - **audio, - "input": { - **audio_input, - "transcription": { - **nested_transcription, - "model": authorized_model, - }, - }, - } - changed = True + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + from litellm.proxy.proxy_server import llm_model_list, llm_router + from litellm.proxy.realtime_endpoints.endpoints import nested_transcription_upstream_model - if not changed: + if not isinstance(valid_token, UserAPIKeyAuth): return message - return json.dumps(message_obj) + + for requested_model in requested_models: + await can_key_call_resolved_model( + model=requested_model, + valid_token=valid_token, + llm_model_list=llm_model_list, + llm_router=llm_router, + ) + upstream_model: Final = nested_transcription_upstream_model( + requested_models[0], valid_token.team_id, llm_router + ) + self._bound_nested_transcription_model = upstream_model + return self._rewrite_session_update_transcription_model(message, upstream_model) def _uses_deferred_backend_setup(self) -> bool: """True when setup is deferred until the client's first session.update.""" @@ -1047,7 +1479,19 @@ class RealTimeStreaming: return event def _event_to_client_json(self, event: dict) -> str: - return json.dumps(self._normalize_event_for_ga_client(event)) + return json.dumps(self._session_event_for_client(self._normalize_event_for_ga_client(event))) + + def _session_event_for_client(self, event: Mapping[str, object]) -> Mapping[str, object]: + if not self._bridges_beta_client_to_ga: + return event + event_type: Final = event.get("type") + if event_type in self._SESSION_EVENT_TYPES: + session: Final = _mapping_field(event, "session") + return event if session is None else {**event, "session": ga_session_to_beta(session)} + if event_type in self._RESPONSE_EVENT_TYPES: + response: Final = _mapping_field(event, "response") + return event if response is None else {**event, "response": ga_response_to_beta(response)} + return event async def _send_event_to_client(self, event: object, event_str: str) -> bool: if self._should_drop_event_from_client(event): @@ -1060,7 +1504,7 @@ class RealTimeStreaming: translated: Final = self._translate_event_to_beta(event) if translated is None: return False - await self.websocket.send_text(json.dumps(translated)) + await self.websocket.send_text(json.dumps(self._session_event_for_client(translated))) return True except Exception as e: verbose_logger.warning( @@ -1183,32 +1627,7 @@ class RealTimeStreaming: self, event_hooks: Sequence["GuardrailEventHooks"], ) -> bool: - """Return True if any callback would run for one of ``event_hooks``.""" - from litellm.integrations.custom_guardrail import CustomGuardrail - - return any( - isinstance(cb, CustomGuardrail) - and any( - cb.should_run_guardrail( - data={**self.request_data, "stream": True}, - event_type=et, - ) - for et in event_hooks - ) - for cb in litellm.callbacks - ) - - def _has_realtime_guardrails(self) -> bool: - """Return True if any callback is registered for realtime guardrail event types.""" - from litellm.types.guardrails import GuardrailEventHooks - - return self._has_realtime_guardrails_for_event_hooks( - [ - GuardrailEventHooks.realtime_input_transcription, - GuardrailEventHooks.pre_call, - GuardrailEventHooks.post_call, - ] - ) + return _guardrails_run_for_event_hooks(self.request_data, event_hooks, litellm.callbacks) def _has_audio_transcription_guardrails(self) -> bool: """Return True when a guardrail is configured for the audio/VAD transcript path. @@ -1222,8 +1641,12 @@ class RealTimeStreaming: return self._has_realtime_guardrails_for_event_hooks([GuardrailEventHooks.realtime_input_transcription]) def _should_disable_vad_auto_response(self) -> bool: - """Transcription-only sessions have no assistant turn to gate and reject realtime session updates.""" - return not self._is_transcription_session and self._has_audio_transcription_guardrails() + """Transcription and translation sessions have no assistant turn to gate and reject realtime session updates.""" + return ( + not self._is_transcription_session + and not self._is_translation_session + and self._has_audio_transcription_guardrails() + ) def _refuses_unscreened_audio(self, msg_type: str | None) -> bool: return ( @@ -1264,7 +1687,6 @@ class RealTimeStreaming: transcript completion uses ``realtime_input_transcription`` only; typed user messages and tool outputs use ``pre_call``. """ - from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks if event_hooks is None: @@ -1273,14 +1695,11 @@ class RealTimeStreaming: _check_data: Final = {**self.request_data, "transcript": transcript, "stream": True} _already_run: Final[set] = set() - for callback in litellm.callbacks: - if not isinstance(callback, CustomGuardrail): - continue - if callback.use_native_lifecycle_hooks: - continue - if id(callback) in _already_run: - continue - if not any(callback.should_run_guardrail(data=_check_data, event_type=et) for et in _realtime_event_types): + runner_guardrails: Final = ( + _realtime_runner_guardrail(callback, _check_data, _realtime_event_types) for callback in litellm.callbacks + ) + for callback in runner_guardrails: + if callback is None or id(callback) in _already_run: continue _already_run.add(id(callback)) try: @@ -1369,7 +1788,10 @@ class RealTimeStreaming: async def _handle_provider_config_message(self, raw_response: str) -> None: """Process a backend message when a provider_config is set (transformed path).""" - returned_object: Final = self.provider_config.transform_realtime_response( + provider_config: Final = self.provider_config + if provider_config is None: + raise RuntimeError("Provider response handling requires a provider configuration") + returned_object: Final = provider_config.transform_realtime_response( raw_response, self.model, self.logging_obj, @@ -1398,6 +1820,7 @@ class RealTimeStreaming: continue if isinstance(event, dict): self._detect_transcription_session_from_backend(event) + self._capture_translation_output_audio(event) is_session_created_event = isinstance(event, dict) and event.get("type") == "session.created" if is_session_created_event: if self._uses_deferred_backend_setup() and not self._backend_setup_complete: @@ -1514,6 +1937,7 @@ class RealTimeStreaming: continue if self.provider_config: + self._backend_sent_non_error_frames = True try: await self._handle_provider_config_message(raw_response) except Exception as e: @@ -1521,6 +1945,8 @@ class RealTimeStreaming: continue else: event = self._parse_backend_event(raw_response) + if event is None or event.get("type") != "error": + self._backend_sent_non_error_frames = True if event is None: await self.websocket.send_text(raw_response) continue @@ -1535,6 +1961,7 @@ class RealTimeStreaming: verbose_logger.debug("Realtime: not forwarding an error caused by a guardrail event: %s", event) continue + self._capture_translation_output_audio(event) if await self._handle_raw_backend_message(event, raw_response): continue @@ -1548,7 +1975,7 @@ class RealTimeStreaming: translated = self._translate_event_to_beta(event) if translated is None: continue - await self.websocket.send_text(json.dumps(translated)) + await self.websocket.send_text(json.dumps(self._session_event_for_client(translated))) async def backend_to_client_send_messages(self) -> BackendClose: import websockets @@ -1573,8 +2000,8 @@ class RealTimeStreaming: def _backend_refused_session(self, close: BackendClose) -> bool: return ( - close.code != 1000 - and not self._backend_sent_frames + not self._backend_sent_non_error_frames + and (close.code != 1000 or self._backend_sent_frames) and not self._holds_session_updates_until_first_content() ) @@ -1617,7 +2044,7 @@ class RealTimeStreaming: _detect_beta_header = detect_beta_header @staticmethod - def _remap_beta_session_to_ga(session: dict) -> dict: + def _remap_beta_session_to_ga(session: dict, all_beta_fields: bool = False, drop_temperature: bool = False) -> dict: """ Convert a beta-style session.update payload to the GA nested schema. @@ -1630,6 +2057,11 @@ class RealTimeStreaming: session.output_audio_format → session.audio.output.format (with type/rate) session.turn_detection → session.audio.input.turn_detection session.input_audio_transcription → session.audio.input.transcription + Only with all_beta_fields: + session.input_audio_noise_reduction → session.audio.input.noise_reduction + session.speed → session.audio.output.speed + session.max_response_output_tokens → session.max_output_tokens + session.temperature (dropped when drop_temperature) ───────────────────────────────────────────────────────────────────── Fields not in the mapping (instructions, tools, etc.) are passed through. GA clients that already use the nested shape are unaffected. @@ -1681,6 +2113,19 @@ class RealTimeStreaming: if "input_audio_transcription" in session: inp["transcription"] = session.pop("input_audio_transcription") + if all_beta_fields and "input_audio_noise_reduction" in session: + inp["noise_reduction"] = session.pop("input_audio_noise_reduction") + + if all_beta_fields and "speed" in session: + out["speed"] = session.pop("speed") + + if all_beta_fields and "max_response_output_tokens" in session: + max_tokens: Final = session.pop("max_response_output_tokens") + session.setdefault("max_output_tokens", max_tokens) + + if all_beta_fields and drop_temperature: + session.pop("temperature", None) + if inp: audio["input"] = inp if out: @@ -1699,6 +2144,53 @@ class RealTimeStreaming: return session + @staticmethod + def _remap_beta_response_to_ga(response: Mapping[str, object], drop_temperature: bool) -> dict[str, object]: + moved: Final = frozenset( + ("modalities", "voice", "output_audio_format", "max_response_output_tokens") + + (("temperature",) if drop_temperature else ()) + ) + raw_format: Final = response.get("output_audio_format") + output_audio: Final = { + **({} if "voice" not in response else {"voice": response["voice"]}), + **( + {} + if "output_audio_format" not in response + else { + "format": RealTimeStreaming._AUDIO_FORMAT_MAP.get(raw_format, raw_format) + if isinstance(raw_format, str) + else raw_format + } + ), + } + modalities: Final = response.get("modalities") + modality_names: Final = ( + frozenset(m.lower() for m in modalities if isinstance(m, str)) + if isinstance(modalities, list) + else frozenset() + ) + output_modalities: Final = ( + ["audio"] if "audio" in modality_names else ["text"] if "text" in modality_names else None + ) + existing_audio: Final = as_str_mapping(response.get("audio")) or {} + existing_output: Final = as_str_mapping(existing_audio.get("output")) or {} + return { + **{key: value for key, value in response.items() if key not in moved}, + **( + {} + if output_modalities is None or "output_modalities" in response + else {"output_modalities": output_modalities} + ), + **( + {} + if "max_response_output_tokens" not in response or "max_output_tokens" in response + else {"max_output_tokens": response["max_response_output_tokens"]} + ), + **( + {} if not output_audio else {"audio": {**existing_audio, "output": {**existing_output, **output_audio}}} + ), + } + @staticmethod def _translate_event_to_beta(event: dict) -> dict | None: """Translate a single GA event dict to its beta equivalent. @@ -1932,10 +2424,34 @@ class RealTimeStreaming: if msg_type == "session.update" and not self._backend_uses_beta_protocol: session = client_event.get("session", {}) if isinstance(session, dict): - session = self._remap_beta_session_to_ga(session) + session = self._remap_beta_session_to_ga( + session, + all_beta_fields=self._remaps_all_beta_fields_to_ga, + drop_temperature=self._drop_params, + ) + if self._is_translation_session: + session.pop("type", None) msg_obj["session"] = session message = json.dumps(msg_obj) + if ( + msg_type == "response.create" + and self._remaps_all_beta_fields_to_ga + and (beta_response := as_str_mapping(client_event.get("response"))) is not None + ): + msg_obj["response"] = self._remap_beta_response_to_ga( + beta_response, drop_temperature=self._drop_params + ) + message = json.dumps(msg_obj) + + if ( + msg_type == "conversation.item.create" + and self._bridges_beta_client_to_ga + and (beta_item := as_str_mapping(client_event.get("item"))) is not None + ): + msg_obj["item"] = beta_item_to_ga(beta_item) + message = json.dumps(msg_obj) + if msg_type == "session.update" and self._event_normalizer: session = client_event.get("session") if isinstance(session, dict): @@ -2016,16 +2532,28 @@ class RealTimeStreaming: return ClientLoopExit.BACKEND_CLOSED except Exception as e: verbose_logger.debug("Error in client ack messages: %s", e) - return ClientLoopExit.CLIENT_DISCONNECTED + return ClientLoopExit.RELAY_FAILED async def bidirectional_forward(self) -> None: + if self._is_translation_session and translation_guardrails_apply(self.request_data): + await self.websocket.send_text( + realtime_error_event(TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE, error_type="guardrail_error") + ) + await self.websocket.close(code=1008, reason=TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE) + return forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages()) client_task: Final = asyncio.create_task(self.client_ack_messages()) try: await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED) - if client_task.done() and client_task.result() is ClientLoopExit.CLIENT_DISCONNECTED: + client_exit: Final = client_task.result() if client_task.done() else None + if client_exit is ClientLoopExit.RELAY_FAILED and self.logging_obj: + self.logging_obj.model_call_details[REALTIME_SESSION_ENDED_ABNORMALLY_KEY] = True + if client_exit in (ClientLoopExit.CLIENT_DISCONNECTED, ClientLoopExit.RELAY_FAILED): return - await self._close_client(await forward_task) + close: Final = await forward_task + if not close.ended_cleanly and self.logging_obj: + self.logging_obj.model_call_details[REALTIME_SESSION_ENDED_ABNORMALLY_KEY] = True + await self._close_client(close) finally: forward_task.cancel() client_task.cancel() @@ -2035,7 +2563,7 @@ class RealTimeStreaming: redacted_message: Final = redact_internal_details_from_client_message(close.message) redacted_reason: Final = redact_internal_details_from_client_message(close.reason) try: - if close.code != 1000: + if not close.ended_cleanly: await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error")) await self.websocket.close( code=client_close_code(close.code), diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 981859c620b..9da0e35bdda 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -10,7 +10,7 @@ import asyncio import copy import inspect -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace from typing import TYPE_CHECKING, Any, Final @@ -107,7 +107,13 @@ def redacted_response_obj_for_custom_logger(response_obj: object, custom_logger: def _is_redactable_response(response_obj: object) -> bool: return isinstance( - response_obj, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse) + response_obj, + ( + litellm.ModelResponse, + litellm.TextCompletionResponse, + litellm.ResponsesAPIResponse, + litellm.EmbeddingResponse, + ), ) or (isinstance(response_obj, dict) and ("choices" in response_obj or "output" in response_obj)) @@ -171,6 +177,8 @@ def _redact_streaming_response(streaming_response): _redact_bedrock_guardrail_trace(streaming_response) elif isinstance(streaming_response, litellm.ResponsesAPIResponse): return perform_redaction({}, streaming_response) + elif isinstance(streaming_response, litellm.TranscriptionResponse): + streaming_response.text = REDACTED_BY_LITELLM elif hasattr(streaming_response, "output"): _redact_responses_api_output(streaming_response.output) if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None: @@ -196,6 +204,9 @@ def _redact_function_call(function_call) -> None: def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" + if isinstance(choice, (litellm.Choices, litellm.utils.StreamingChoices, litellm.utils.TextChoices)): + drop_token_ids(choice) + _redact_provider_specific_fields(choice) if isinstance(choice, litellm.Choices): if choice.message.content is not None: choice.message.content = REDACTED_BY_LITELLM @@ -203,6 +214,7 @@ def _redact_choice_content(choice): choice.message.reasoning_content = REDACTED_BY_LITELLM if hasattr(choice.message, "thinking_blocks"): choice.message.thinking_blocks = None + _drop_field_when_present(choice.message, "reasoning_items") _redact_provider_specific_fields(choice.message) _redact_tool_calls(getattr(choice.message, "tool_calls", None)) _redact_function_call(getattr(choice.message, "function_call", None)) @@ -213,9 +225,14 @@ def _redact_choice_content(choice): choice.delta.reasoning_content = REDACTED_BY_LITELLM if hasattr(choice.delta, "thinking_blocks"): choice.delta.thinking_blocks = None + _drop_field_when_present(choice.delta, "reasoning_items") _redact_provider_specific_fields(choice.delta) _redact_tool_calls(getattr(choice.delta, "tool_calls", None)) _redact_function_call(getattr(choice.delta, "function_call", None)) + elif isinstance(choice, litellm.utils.TextChoices): + _redact_field_when_present(choice, "text") + _redact_field_when_present(choice, "reasoning_content") + _drop_field_when_present(choice, "logprobs") def _as_mapping(value: object) -> Mapping[str, object] | None: @@ -227,14 +244,27 @@ def _redact_provider_specific_fields(message: object) -> None: setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields)) # noqa: B010 # typed as object +def _drop_provider_specific_token_ids(holder: object) -> None: + fields: Final = _as_mapping(getattr(holder, "provider_specific_fields", None)) + if fields is not None and fields.get("token_ids") is not None: + setattr(holder, "provider_specific_fields", {**fields, "token_ids": None}) # noqa: B010 # typed as object + + +def drop_token_ids(choice: object) -> None: + if getattr(choice, "token_ids", None) is not None: + setattr(choice, "token_ids", None) # noqa: B010 # pydantic extra, not a declared field + for holder in (choice, getattr(choice, "delta", None), getattr(choice, "message", None)): + _drop_provider_specific_token_ids(holder) + + def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None: - for message in (choice.get("message"), choice.get("delta")): + for holder in (choice, choice.get("message"), choice.get("delta")): if ( - (readable := _as_mapping(message)) is not None + (readable := _as_mapping(holder)) is not None and (fields := _as_mapping(readable.get("provider_specific_fields"))) is not None - and isinstance(message, dict) + and isinstance(holder, dict) ): - message["provider_specific_fields"] = _redacted_provider_specific_fields(fields) + holder["provider_specific_fields"] = _redacted_provider_specific_fields(fields) def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[str, object]: @@ -251,6 +281,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ "compaction_blocks", "compaction_start", "compaction_delta", + "token_ids", ) if key in fields }, @@ -262,6 +293,11 @@ def _redact_field_when_present(item: object, field: str) -> None: setattr(item, field, REDACTED_BY_LITELLM) +def _drop_field_when_present(item: object, field: str) -> None: + if getattr(item, field, None) is not None: + setattr(item, field, None) + + def _redact_responses_api_output(output_items): """Helper to redact ResponsesAPIResponse output items.""" for output_item in output_items: @@ -418,30 +454,22 @@ def _redact_tool_calls_dict(message: Mapping[str, object]) -> None: function_call["arguments"] = REDACTED_BY_LITELLM -def _redact_model_response_dict_choices(choices, redacted_str: str): +def _redact_model_response_dict_choices(choices: Iterable[object], redacted_str: str) -> None: for choice in choices: if isinstance(choice, dict): _redact_provider_specific_fields_dict(choice) - if "message" in choice and isinstance(choice["message"], dict): - if choice["message"].get("content") is not None: - choice["message"]["content"] = redacted_str - if choice["message"].get("reasoning_content") is not None: - choice["message"]["reasoning_content"] = redacted_str - if "thinking_blocks" in choice["message"]: - choice["message"]["thinking_blocks"] = None - if "audio" in choice["message"]: - choice["message"]["audio"] = None - _redact_tool_calls_dict(choice["message"]) - elif "delta" in choice and isinstance(choice["delta"], dict): - if choice["delta"].get("content") is not None: - choice["delta"]["content"] = redacted_str - if choice["delta"].get("reasoning_content") is not None: - choice["delta"]["reasoning_content"] = redacted_str - if "thinking_blocks" in choice["delta"]: - choice["delta"]["thinking_blocks"] = None - if "audio" in choice["delta"]: - choice["delta"]["audio"] = None - _redact_tool_calls_dict(choice["delta"]) + if choice.get("token_ids") is not None: + choice["token_ids"] = None + if ( + body := next((choice[key] for key in ("message", "delta") if isinstance(choice.get(key), dict)), None) + ) is not None: + for field in ("content", "reasoning_content"): + if body.get(field) is not None: + body[field] = redacted_str + for field in ("thinking_blocks", "reasoning_items", "audio"): + if field in body: + body[field] = None + _redact_tool_calls_dict(body) elif choice.get("text") is not None: choice["text"] = redacted_str else: @@ -564,7 +592,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons return {"text": REDACTED_BY_LITELLM} _result: Final = copy.deepcopy(result) - if isinstance(_result, litellm.ModelResponse): + if isinstance(_result, (litellm.ModelResponse, litellm.TextCompletionResponse)): if hasattr(_result, "choices") and _result.choices is not None: for choice in _result.choices: _redact_choice_content(choice) diff --git a/litellm/litellm_core_utils/stream_usage_chunks.py b/litellm/litellm_core_utils/stream_usage_chunks.py index eda57c7ea12..d073ddc2724 100644 --- a/litellm/litellm_core_utils/stream_usage_chunks.py +++ b/litellm/litellm_core_utils/stream_usage_chunks.py @@ -1,6 +1,7 @@ from collections.abc import Iterable from typing import Final +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.types.utils import ModelResponseStream, StreamingChoices @@ -9,12 +10,19 @@ def is_empty_streaming_choice(choice: StreamingChoices) -> bool: return False if getattr(choice, "logprobs", None) is not None: return False + if not _all_values_none(_vllm_choice_values(choice)): + return False delta: Final = choice.delta if hasattr(choice, "delta") else None if delta is None: return True return _all_values_none(delta.model_dump().values()) +def _vllm_choice_values(choice: StreamingChoices) -> Iterable[object]: + choice_extras: Final = choice.model_extra or {} + return (choice_extras.get(field) for field in VLLM_CHOICE_FIELDS) + + def _all_values_none(values: Iterable[object]) -> bool: return all(value is None for value in values) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 7191b2ce8e8..c87357122ce 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.core_helpers import ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, UPSTREAM_USAGE_COST, @@ -31,6 +32,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict ) +from litellm.litellm_core_utils.redact_messages import drop_token_ids from litellm.litellm_core_utils.token_counter import usage_tokenizer_model from litellm.types.llms.openai import ( AllMessageValues, @@ -434,6 +436,8 @@ class ChunkProcessor: chunks, unreported="tool_calls" if _carries_tool_calls(chunks) else "stop" ), message={"role": ChunkProcessor._get_role_from_chunks(chunks), "content": ""}, + provider_specific_fields=None, + **_streamed_vllm_choice_fields(chain.from_iterable(chunk.get("choices") or () for chunk in chunks)), ) @staticmethod @@ -504,6 +508,7 @@ class ChunkProcessor: "index": 0, "message": {"role": role, "content": ""}, "finish_reason": finish_reason, + **_streamed_vllm_choice_fields(_first_choices(chunks)), } ] ), @@ -1371,6 +1376,7 @@ def _rebuilt_text_choice(index: int, choices: Sequence[TextChoices]) -> TextChoi index=index, logprobs=_merged_logprobs(choices), finish_reason=next((choice.finish_reason for choice in reversed(choices) if choice.finish_reason), None), + **_streamed_vllm_choice_fields(choices), ) @@ -1504,6 +1510,31 @@ def _field_of(value: object, name: str) -> object: return getattr(value, name, None) +def _first_choices(chunks: Iterable[object]) -> tuple[object, ...]: + return tuple(next(iter(object_items(_field_of(chunk, "choices"))), None) for chunk in chunks) + + +def _choice_extras(choice: object) -> Mapping[str, object]: + if isinstance(choice, BaseModel): + return choice.model_extra or {} + return object_mapping(choice) or {} + + +def _streamed_vllm_choice_fields(choices: Iterable[object]) -> Mapping[str, object]: + extras: Final = tuple(map(_choice_extras, choices)) + token_id_runs: Final = tuple(object_items(fields.get("token_ids")) for fields in extras) + stop_reasons: Final = tuple(reason for fields in extras if (reason := fields.get("stop_reason")) is not None) + has_token_ids: Final = any(fields.get("token_ids") is not None for fields in extras) + streamed_values: Final[Mapping[str, object]] = MappingProxyType( + { + "token_ids": list(chain.from_iterable(token_id_runs)) if has_token_ids else None, + "stop_reason": stop_reasons[-1] if stop_reasons else None, + } + ) + present_fields: Final = frozenset(chain.from_iterable(fields.keys() for fields in extras)) + return MappingProxyType({field: streamed_values[field] for field in VLLM_CHOICE_FIELDS if field in present_fields}) + + def _delta_reasoning_items(delta: object) -> Sequence[object]: items: Final = _field_of(delta, "reasoning_items") return cast(Sequence[object], items) if isinstance(items, list) else () # cast-ok: items forwarded verbatim @@ -1776,6 +1807,19 @@ def stream_chunk_builder( end_time: object = None, logging_obj: Optional["Logging"] = None, count_prompt_tokens: Callable[[], int] | None = None, +) -> ModelResponse | TextCompletionResponse | None: + response: Final = _assembled_stream_response(chunks, messages, logging_obj, count_prompt_tokens) + if response is not None and getattr(logging_obj, "drops_stream_token_ids", False) is True: + for choice in response.choices: + drop_token_ids(choice) + return response + + +def _assembled_stream_response( + chunks: Sequence[object], + messages: Sequence[AllMessageValues | Message] | None, + logging_obj: Optional["Logging"], + count_prompt_tokens: Callable[[], int] | None, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -1788,7 +1832,7 @@ def stream_chunk_builder( if not chunks: return None - processor: Final = ChunkProcessor(chunks, messages) + processor: Final = ChunkProcessor(list(chunks), messages) chunks = processor.chunks ### BASE-CASE ### diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 9bb5b761fd5..a9ef06c3b4f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_FIELDS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( UPSTREAM_USAGE_COST, @@ -37,10 +38,11 @@ from litellm.litellm_core_utils.hidden_params import ( set_hidden_param, stored_attribute, ) +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_items from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) -from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject +from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject, drop_token_ids from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel @@ -137,6 +139,70 @@ def _prepend_upstream_choice_extras( choice.__pydantic_fields_set__.update(upstream_extras) +_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) +_NO_CHOICE_FIELDS_BY_INDEX: Final[Mapping[int, Mapping[str, object]]] = MappingProxyType({}) + + +def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None: + match getattr(upstream_chunk, "choices", None): + case [BaseModel() as first_choice, *_]: + return first_choice + case _: + return None + + +def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + first_choice: Final = _upstream_first_choice(upstream_chunk) + if first_choice is None: + return _NO_UPSTREAM_CHOICE_FIELDS + upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {} + return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras}) + + +def _upstream_choice_index(upstream_chunk: object) -> int: + upstream_index: Final[object] = getattr(_upstream_first_choice(upstream_chunk), "index", 0) + return upstream_index if isinstance(upstream_index, int) else 0 + + +def _response_choice_index(response_obj: Mapping[str, object]) -> int: + choice_index: Final = response_obj.get("choice_index") + if isinstance(choice_index, int): + return choice_index + return _upstream_choice_index(response_obj.get("original_chunk")) + + +def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + first_choice: Final = _upstream_first_choice(upstream_chunk) + upstream_extras: Final = first_choice.model_extra if first_choice is not None else None + if not upstream_extras: + return _NO_UPSTREAM_CHOICE_FIELDS + return MappingProxyType( + {field: upstream_extras[field] for field in fields if upstream_extras.get(field) is not None} + ) + + +def _apply_choice_fields(choice: StreamingChoices, fields: Mapping[str, object]) -> None: + for field, value in fields.items(): + setattr(choice, field, value) + + +def _choice_fields_chunk( + model_response: ModelResponseStream, upstream_index: int, fields: Mapping[str, object] +) -> ModelResponseStream: + model_response.choices[0].delta = Delta(content="") # rebind-ok: chunks are built in place here + model_response.choices[0].index = upstream_index # rebind-ok: chunks are built in place here + _apply_choice_fields(model_response.choices[0], fields) + return model_response + + +def _without_token_ids_under_post_call_rules(chunk: ModelResponseStream | None) -> ModelResponseStream | None: + if chunk is None or not litellm.post_call_rules: + return chunk + for choice in chunk.choices: + drop_token_ids(choice) + return chunk + + def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool: return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) @@ -378,6 +444,8 @@ class CustomStreamWrapper: self._provider_response_model: str | None = None self._upstream_provider: str | None = None self.received_finish_reason: str | None = None + self.received_finish_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX + self.carried_chunk_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.received_any_finish_reason = False # a later chunk can reset intermittent_finish_reason self.special_tokens = [ @@ -823,6 +891,7 @@ class CustomStreamWrapper: "logprobs": logprobs, "original_chunk": str_line, "usage": usage, + "chunk_choice_fields": _non_null_upstream_choice_fields(str_line, STREAM_CHUNK_CHOICE_FIELDS), } except Exception as e: raise e @@ -847,7 +916,7 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_openai_text_completion_chunk(self, chunk): + def handle_openai_text_completion_chunk(self, chunk: object): try: text = "" is_finished = False @@ -860,11 +929,15 @@ class CustomStreamWrapper: is_finished = True finish_reason = choices[0].finish_reason usage = getattr(chunk, "usage", None) + choice_index: Final = _upstream_choice_index(chunk) return { "text": text, "is_finished": is_finished, "finish_reason": finish_reason, "usage": usage, + "chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS), + "choice_index": choice_index, + "sibling_finish_reason": finish_reason if choice_index != 0 else None, } except Exception as e: @@ -1183,6 +1256,13 @@ class CustomStreamWrapper: ) is_chunk_non_empty: Final = self.is_chunk_non_empty(completion_obj, model_response, response_obj) + upstream_chunk_choice_fields: Final[Mapping[str, object]] = ( + _NO_UPSTREAM_CHOICE_FIELDS + if litellm.post_call_rules + else response_obj.get("chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS) + ) + chunk_choice_fields: Final = self._take_carried_chunk_choice_fields(response_obj, upstream_chunk_choice_fields) + sibling_finish_reason: Final[str | None] = response_obj.get("sibling_finish_reason") if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string self.raise_on_model_repetition() @@ -1270,16 +1350,40 @@ class CustomStreamWrapper: if response_obj.get("provider_specific_fields") is not None: completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] model_response.choices[0].delta = Delta(**completion_obj) - _index: Final[int | None] = completion_obj.get("index") + _index: Final[int | None] = completion_obj.get("index", response_obj.get("choice_index")) if _index is not None: model_response.choices[0].index = _index + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) + if sibling_finish_reason is not None: + self._finish_sibling_choice(model_response.choices[0], sibling_finish_reason) + elif response_obj.get("is_finished"): + self._apply_finish_choice_fields(model_response.choices[0]) self._optional_combine_thinking_block_in_choices(model_response=model_response) return model_response else: return + elif sibling_finish_reason is not None: + sibling_chunk: Final = _choice_fields_chunk( + model_response, _response_choice_index(response_obj), chunk_choice_fields + ) + self._finish_sibling_choice(sibling_chunk.choices[0], sibling_finish_reason) + return sibling_chunk elif self.received_finish_reason is not None: + chunk_upstream_index: Final = _response_choice_index(response_obj) + if chunk_choice_fields and self._streams_past_the_pending_finish( + model_response, response_obj, chunk_upstream_index + ): + self.carried_chunk_choice_fields = MappingProxyType( + {**self.carried_chunk_choice_fields, chunk_upstream_index: chunk_choice_fields} + ) + else: + another_choice_fields_chunk: Final = self._another_choice_fields_chunk( + model_response, response_obj, chunk_choice_fields + ) + if another_choice_fields_chunk is not None: + return another_choice_fields_chunk if self.sent_last_chunk is True: # Bedrock returns the guardrail trace in the last chunk - we want to return this here if self.custom_llm_provider == "bedrock" and "trace" in model_response: @@ -1317,12 +1421,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason( map_finish_reason(finish_reason=self.received_finish_reason) ) # ensure consistent output to openai + self._apply_finish_choice_fields(model_response.choices[0]) + if ( + response_obj.get("is_finished") + and _response_choice_index(response_obj) == model_response.choices[0].index + ): + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) self.sent_last_chunk = True return model_response elif self._has_special_delta_content(model_response): return self._handle_special_delta_content(model_response) + elif chunk_choice_fields: + return _choice_fields_chunk(model_response, _response_choice_index(response_obj), chunk_choice_fields) else: if hasattr(model_response, "usage"): self.chunks.append(model_response) @@ -1545,6 +1657,8 @@ class CustomStreamWrapper: completion_obj["content"] = response_obj["text"] print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: + self._record_finish_choice_fields(chunk) + if response_obj["is_finished"] and response_obj["sibling_finish_reason"] is None: self.received_finish_reason = response_obj["finish_reason"] if response_obj["usage"] is not None: _text_completion_usage: Final[Usage] = response_obj["usage"] @@ -1608,6 +1722,7 @@ class CustomStreamWrapper: self.system_fingerprint = cached_chunk.system_fingerprint if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] + self._record_finish_choice_fields(cached_chunk) else: # openai / azure chat model self.uses_openai_chat_stream = True if isinstance(chunk, ChatCompletionChunk): @@ -1637,6 +1752,7 @@ class CustomStreamWrapper: f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}" ) self.received_finish_reason = response_obj["finish_reason"] + self._record_finish_choice_fields(chunk) if response_obj.get("original_chunk", None) is not None: if hasattr(response_obj["original_chunk"], "id"): model_response = self.set_model_id(response_obj["original_chunk"].id, model_response) @@ -1704,7 +1820,7 @@ class CustomStreamWrapper: completion_obj=completion_obj, ) if isinstance(dispatch_result, _ProviderChunkEarlyReturn): - return dispatch_result.value + return _without_token_ids_under_post_call_rules(dispatch_result.value) response_obj = dispatch_result.response_obj model_response.model = self._served_openrouter_model() or self.model @@ -1794,7 +1910,7 @@ class CustomStreamWrapper: response_obj=response_obj, plain_delta_created=plain_delta_created, ) - return result + return _without_token_ids_under_post_call_rules(result) except StopIteration: raise StopIteration @@ -1987,6 +2103,75 @@ class CustomStreamWrapper: return "tool_calls" return finish_reason + def _record_finish_choice_fields(self, upstream_chunk: object) -> None: + self.received_finish_choice_fields = MappingProxyType( + { + **self.received_finish_choice_fields, + _upstream_choice_index(upstream_chunk): _upstream_choice_fields( + upstream_chunk, STREAM_FINISH_CHOICE_FIELDS + ), + } + ) + + def _apply_finish_choice_fields(self, choice: StreamingChoices) -> None: + _apply_choice_fields(choice, self.received_finish_choice_fields.get(choice.index, _NO_UPSTREAM_CHOICE_FIELDS)) + + def _finish_sibling_choice(self, choice: StreamingChoices, finish_reason: str) -> None: + choice.finish_reason = map_finish_reason(finish_reason) # rebind-ok: chunks are built in place here + self._apply_finish_choice_fields(choice) + + def _streams_past_the_pending_finish( + self, model_response: ModelResponseStream, response_obj: Mapping[str, object], upstream_index: int + ) -> bool: + finishing_index: Final = model_response.choices[0].index + return ( + not self.sent_last_chunk + and not response_obj.get("is_finished") + and upstream_index != finishing_index + and finishing_index in self.received_finish_choice_fields + ) + + def _take_carried_chunk_choice_fields( + self, response_obj: Mapping[str, object], chunk_choice_fields: Mapping[str, object] + ) -> Mapping[str, object]: + if not self.carried_chunk_choice_fields: + return chunk_choice_fields + upstream_index: Final = _response_choice_index(response_obj) + carried: Final = self.carried_chunk_choice_fields.get(upstream_index) + if carried is None: + return chunk_choice_fields + self.carried_chunk_choice_fields = MappingProxyType( + {index: fields for index, fields in self.carried_chunk_choice_fields.items() if index != upstream_index} + ) + return MappingProxyType( + { + field: [*object_items(carried.get(field)), *object_items(chunk_choice_fields.get(field))] + for field in STREAM_CHUNK_CHOICE_FIELDS + if field in carried or field in chunk_choice_fields + } + ) + + def _another_choice_fields_chunk( + self, + model_response: ModelResponseStream, + response_obj: Mapping[str, object], + chunk_choice_fields: Mapping[str, object], + ) -> ModelResponseStream | None: + upstream_index: Final = _response_choice_index(response_obj) + is_finished: Final = bool(response_obj.get("is_finished")) + if is_finished and upstream_index == model_response.choices[0].index and not self.sent_last_chunk: + return None + finish_fields: Final = ( + self.received_finish_choice_fields.get(upstream_index, _NO_UPSTREAM_CHOICE_FIELDS) + if is_finished + else _NO_UPSTREAM_CHOICE_FIELDS + ) + if not chunk_choice_fields and all(value is None for value in finish_fields.values()): + return None + return _choice_fields_chunk( + model_response, upstream_index, MappingProxyType({**finish_fields, **chunk_choice_fields}) + ) + def finish_reason_handler(self): model_response: Final = self.model_response_creator( hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} @@ -2005,6 +2190,7 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason( _finish_reason if _finish_reason is not None else "stop" ) + self._apply_finish_choice_fields(model_response.choices[0]) close_tag: Final = self._close_open_think_block() if close_tag is not None: model_response.choices[0].delta.content = close_tag diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 323af72940a..df8038e756c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -990,7 +990,11 @@ def _count_extra( tool_section (Optional[ChatTemplateToolSection]): How the model's chat template renders tools, if known. """ reply_primer_tokens: Final = 3 - tool_choice_tokens: Final = 0 if tool_section is not None else _count_tool_choice(count_function, tool_choice) + tool_choice_tokens: Final = ( + _count_tool_choice(count_function, tool_choice) + if tool_section is None or tool_section.counts_tool_choice + else 0 + ) return ( reply_primer_tokens + _count_tool_definitions(count_function, tools, includes_system_message, tool_section) @@ -1029,14 +1033,29 @@ def _count_tool_choice( return 0 +def _named_function_definitions(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: + return (function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name")) + + +def sends_function_tools(tools: Sequence[object] | None) -> bool: + return tools is not None and any(_named_function_definitions(tools)) + + +def _without_false_strict(function: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in function.items() if key != "strict" or value} + + def _count_chat_template_tools( count_function: TokenCounterFunction, tools: Sequence[object], includes_system_message: bool, tool_section: ChatTemplateToolSection, ) -> int: - functions: Final = tuple( - function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name") + named_functions: Final = tuple(_named_function_definitions(tools)) + functions: Final = ( + tuple(map(_without_false_strict, named_functions)) + if tool_section.renders_strict_only_when_true + else named_functions ) if not functions: return 0 diff --git a/litellm/litellm_core_utils/tool_fields.py b/litellm/litellm_core_utils/tool_fields.py index f90c8f5d174..5ee96ea9133 100644 --- a/litellm/litellm_core_utils/tool_fields.py +++ b/litellm/litellm_core_utils/tool_fields.py @@ -1,13 +1,16 @@ -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence +from copy import copy from dataclasses import dataclass from functools import reduce -from typing import Annotated, Final, TypeAlias +from typing import Annotated, Final, TypeAlias, TypeVar from pydantic import BeforeValidator, ConfigDict, Json, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm._logging import verbose_logger +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params ROUTES_WITH_TOOLS: Final = frozenset( { @@ -162,3 +165,145 @@ def resolve_deployment_tool_fields( if "proxy_server_request" not in kwargs or route not in ROUTES_WITH_TOOLS: return kwargs return resolve_tool_fields(kwargs, lambda: (drop_policy(kwargs),)) + + +_ToolT: Final = TypeVar("_ToolT") + +HOSTED_MCP_DROP_WARNING: Final = ( + "Dropping the hosted MCP tool(s) %s for model=%s on %s (drop_params=True) because this route cannot run a " + "hosted MCP tool." +) + + +def _hosted_mcp_server_label(tool: object) -> str | None: + match tool: + case {"type": "mcp", "server_label": str() as server_label}: + return server_label + case {"type": "mcp"}: + return "unlabeled" + case _: + return None + + +def _drop_params_hint(dropped: str) -> str: + return ( + f"To drop {dropped} instead, set `litellm.drop_params=True` or for proxy:\n\n" + "`litellm_settings:\n drop_params: true`\n" + ) + + +def _one_log_line(text: str) -> str: + return text.replace("\r", "").replace("\n", "") + + +def _unrunnable_mcp_error(message: str, model: str, route: str) -> UnsupportedParamsError: + return UnsupportedParamsError(message=message, model=model, llm_provider=route, status_code=400) + + +def tools_without_hosted_mcp(tools: Sequence[_ToolT], model: str, route: str, drop_params: bool) -> Sequence[_ToolT]: + server_labels: Final = tuple(label for label in map(_hosted_mcp_server_label, tools) if label is not None) + if not server_labels: + return tools + named_servers: Final = ", ".join(f"`{label}`" for label in server_labels) + if not drops_unsupported_params(drop_params): + raise _unrunnable_mcp_error( + f'{model} on {route} cannot run the hosted MCP tool {named_servers} (`type: "mcp"`). To have LiteLLM ' + "run the MCP server for this model, register it on the LiteLLM MCP gateway and send " + '`server_url: "litellm_proxy/mcp/"`. ' + _drop_params_hint("hosted MCP tools"), + model=model, + route=route, + ) + verbose_logger.warning( + HOSTED_MCP_DROP_WARNING, _one_log_line(named_servers), _one_log_line(model), _one_log_line(route) + ) + return [tool for tool in tools if _hosted_mcp_server_label(tool) is None] + + +MCP_SERVERS_DROP_WARNING: Final = ( + "Dropping mcp_servers %s and their mcp_toolset tools for model=%s on %s (drop_params=True) because LiteLLM " + "cannot pass mcp_servers to this route." +) + +ROUTES_WITHOUT_MCP_SERVERS: Final = frozenset( + {"bedrock", "databricks", "deepseek", "gemini", "gigachat", "ollama", "ollama_chat", "vertex_ai", "vertex_ai_beta"} +) + + +def _mcp_server_name(server: object) -> str: + match server: + case {"name": str() as name}: + return name + case _: + return "unnamed" + + +def _mcp_server_names(mcp_servers: object) -> tuple[str, ...]: + if not mcp_servers: + return () + if not isinstance(mcp_servers, list): + return ("unnamed",) + return tuple(map(_mcp_server_name, _TOOLS.validate_python(mcp_servers))) + + +def _is_mcp_toolset(tool: object) -> bool: + match tool: + case {"type": "mcp_toolset"}: + return True + case _: + return False + + +def drops_mcp_servers(request_kwargs: Mapping[str, object]) -> bool: + additional_drop_params: Final = request_kwargs.get("additional_drop_params") + return drops_unsupported_params(request_kwargs.get("drop_params")) or ( + isinstance(additional_drop_params, list) and "mcp_servers" in additional_drop_params + ) + + +def tools_without_mcp_servers( + mcp_servers: object, tools: Sequence[_ToolT], model: str, route: str, drop_params: bool +) -> Sequence[_ToolT]: + server_names: Final = _mcp_server_names(mcp_servers) + if not server_names: + return tools + named_servers: Final = ", ".join(f"`{name}`" for name in server_names) + if not drops_unsupported_params(drop_params): + raise _unrunnable_mcp_error( + f"LiteLLM cannot pass `mcp_servers` to {model} on {route}, so the MCP server {named_servers} would not " + 'run. To have LiteLLM run it for this model, register it on the LiteLLM MCP gateway and send a `type: "mcp"` ' + 'tool with `server_url: "litellm_proxy/mcp/"`. ' + _drop_params_hint("`mcp_servers`"), + model=model, + route=route, + ) + verbose_logger.warning( + MCP_SERVERS_DROP_WARNING, _one_log_line(named_servers), _one_log_line(model), _one_log_line(route) + ) + return [tool for tool in tools if not _is_mcp_toolset(tool)] + + +_TOOL_SELECTION_PARAMS: Final = frozenset({"tools", "tool_choice", "parallel_tool_calls"}) +_ParamsT: Final = TypeVar("_ParamsT", bound=dict[str, object]) + + +def tool_params_without_hosted_mcp( + params: _ParamsT, model: str, route: str, drop_params: bool, runs_hosted_mcp: Callable[[], bool] | None = None +) -> _ParamsT: + tools: Final = params.get("tools") + if not isinstance(tools, list) or not tools: + return params + sent_tools: Final = _TOOLS.validate_python(tools) + if all(_hosted_mcp_server_label(tool) is None for tool in sent_tools) or ( + runs_hosted_mcp is not None and runs_hosted_mcp() + ): + return params + kept_tools: Final = tools_without_hosted_mcp(tools=sent_tools, model=model, route=route, drop_params=drop_params) + if len(kept_tools) == len(sent_tools): + return params + kept_params: Final = copy(params) + if kept_tools: + kept_params["tools"] = list(kept_tools) + return kept_params + dropped_keys: Final = ("tools",) if params.get("functions") else _TOOL_SELECTION_PARAMS + for key in dropped_keys: + _ = kept_params.pop(key, None) + return kept_params diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7f688ae733b..7e3fda06b1a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -34,6 +34,7 @@ from litellm.llms.anthropic.messages.adapters.transformation import ( ) from litellm.llms.anthropic.messages.convertible_blocks import ( keep_non_string_tool_result_texts, + keep_url_documents, scannable_tool_result_request, ) from litellm.llms.anthropic.messages.native.utils import filtered_safeguard_results, redacted_stop_details @@ -1187,7 +1188,10 @@ class AnthropicMessagesHandler(BaseTranslation): cast(Sequence[object], converted), # cast-ok: Anthropic message payload original_tool_names, ) - data["messages"] = keep_non_string_tool_result_texts(converted, data.get("messages")) + original_messages: Final = data.get("messages") + data["messages"] = keep_url_documents( + keep_non_string_tool_result_texts(converted, original_messages), original_messages + ) @staticmethod def _extract_midturn_system_text( diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 75f88355652..4eae06dbcdf 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -516,6 +516,7 @@ class AnthropicChatCompletion(BaseLLM): headers=headers, request_data=request_data, provider=custom_llm_provider, + api_base=api_base, ) data: Final = dict( strip_beta_gated_request_body( @@ -1194,6 +1195,9 @@ class ModelResponseIterator: start_message: Final = message_start_block["message"] if "model" in start_message: self.served_model = start_message["model"] + input_transformations: Final = start_message.get("input_transformations") + if input_transformations is not None: + provider_specific_fields["input_transformations"] = input_transformations if "usage" in message_start_block["message"]: usage = self._handle_usage(anthropic_usage_chunk=message_start_block["message"]["usage"]) elif type_chunk == "error": diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index f198d76f72a..6b96df8f879 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -13,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm +from litellm.anthropic_beta_headers_manager import get_unsupported_headers from litellm.constants import ( ANTHROPIC_BILLING_METADATA_PREFIX, ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, @@ -26,6 +27,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, RESPONSE_FORMAT_TOOL_NAME, ) +from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, map_finish_reason from litellm.litellm_core_utils.prompt_templates.common_utils import ( DEFS_MAX_INLINED_BYTES, @@ -46,8 +48,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( RemoteMedia, + async_format_pdf_links, async_inline_remote_media, + format_pdf_links, inline_remote_image_urls, + is_unformatted_pdf_link_image, ) from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( message_field, @@ -56,6 +61,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import split_leading_system_run, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_non_string_reasoning_effort +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.base_llm.base_utils import ( ResponseFormatSchema, as_json_object, @@ -1315,6 +1321,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): *(tool for tool in tools[first_search_tool_index + 1 :] if not _is_openai_tool_search_tool(tool)), ] + def runs_hosted_mcp(self) -> bool: + return ANTHROPIC_BETA_HEADER_VALUES.MCP_CLIENT_2025_04_04.value not in get_unsupported_headers( + self._resolved_provider + ) + def _tools_param_with_openai_tool_search_mapped( self, non_default_params: Mapping[str, object], model: str, drop_params: bool ) -> Mapping[str, object]: @@ -1903,7 +1914,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): model: str, drop_params: bool, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params=non_default_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, + model=model, + route=self._resolved_provider, + drop_params=drop_params, + runs_hosted_mcp=self.runs_hosted_mcp, + ) + is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params=request_params) # NB: ``map_openai_params`` deliberately does NOT sanitize tool names # here. Names are the *original* OpenAI names at this stage, and must @@ -1916,7 +1934,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # via ``data = {**optional_params}`` and would 400 with # ``Extra inputs are not permitted``). - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["max_tokens"] = value if isinstance(value, int) else max(1, int(round(value))) elif param == "tools": @@ -1930,8 +1948,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params["mcp_servers"] = mcp_servers elif param == "tool_choice" or param == "parallel_tool_calls": _tool_choice: AnthropicMessagesToolChoice | None = self._map_tool_choice( - tool_choice=non_default_params.get("tool_choice"), - parallel_tool_use=non_default_params.get("parallel_tool_calls"), + tool_choice=request_params.get("tool_choice"), + parallel_tool_use=request_params.get("parallel_tool_calls"), ) if _tool_choice is not None: @@ -1957,7 +1975,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if self.uses_native_structured_output_for_schema( model, self._extract_json_schema_from_response_format(value), - REQUEST_PARAMS.validate_python(non_default_params), + REQUEST_PARAMS.validate_python(request_params), ): _output_format = self.map_response_format_to_anthropic_output_format(value) if _output_format is not None: @@ -1993,7 +2011,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # `thinking={type: enabled, budget_tokens}` interface a # pre-4.6 model actually supports instead of forwarding a # shape the model will reject. - max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens") + max_tokens = request_params.get("max_completion_tokens") or request_params.get("max_tokens") legacy_thinking = AnthropicConfig._map_reasoning_effort( reasoning_effort="medium", model=model, @@ -2087,7 +2105,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## handle thinking tokens self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=optional_params + non_default_params=request_params, optional_params=optional_params ) return optional_params @@ -2319,6 +2337,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def inlines_remote_media(self, media: RemoteMedia) -> bool: return inline_remote_image_urls(media) and media.url.startswith("http://") + def probes_pdf_link(self, media: RemoteMedia) -> bool: + return is_unformatted_pdf_link_image(media) and not self.inlines_remote_media(media) + + def _with_pdf_links_formatted(self, messages: Sequence[AllMessageValues]) -> Sequence[AllMessageValues]: + return messages if is_event_loop_running() else format_pdf_links(messages, should_probe=self.probes_pdf_link) + async def async_transform_request( self, model: str, @@ -2327,9 +2351,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): litellm_params: dict[str, object], # mutable-ok: BaseConfig signature headers: dict[str, object], # mutable-ok: BaseConfig signature ) -> dict[str, object]: # mutable-ok: BaseConfig signature + inlined: Final = await async_inline_remote_media(messages, should_inline=self.inlines_remote_media) return self.transform_request( model=model, - messages=await async_inline_remote_media(messages, should_inline=self.inlines_remote_media), + messages=list(await async_format_pdf_links(inlined, should_probe=self.probes_pdf_link)), optional_params=optional_params, litellm_params=litellm_params, headers=headers, @@ -2345,6 +2370,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): messages, optional_params.get("tools"), optional_params.get("cache_control"), optional_params ) + def translates_prompt_cache_breakpoint(self, model: str) -> bool: + return True + def transform_request( self, model: str, @@ -2358,19 +2386,32 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ ## VALIDATE REQUEST """Anthropic requires ``tools`` when messages include tool blocks; LiteLLM injects a dummy tool if omitted (no ``modify_params`` needed).""" + from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook, cache_control_count from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_messages_pt, ) drop_params: Final = drops_unsupported_params(litellm_params.get("drop_params")) + translated_messages: Final = ( + list( + AnthropicCacheControlHook.prompt_cache_breakpoints_as_capped_cache_control( + messages, + drop_params, + cache_control_count, + lambda: AnthropicCacheControlHook.count_chat_external_cache_breakpoints(optional_params), + ) + ) + if self.translates_prompt_cache_breakpoint(model) + else messages + ) capped_messages: Final = ( list( self._messages_within_cache_control_cap( - messages, _REQUEST_PARAMS_ADAPTER.validate_python(optional_params) + translated_messages, _REQUEST_PARAMS_ADAPTER.validate_python(optional_params) ) ) if drop_params - else messages + else translated_messages ) if "tools" not in optional_params and has_tool_call_blocks(capped_messages): @@ -2463,7 +2504,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): try: anthropic_messages = anthropic_messages_pt( model=model, - messages=list(conversation), + messages=list(self._with_pdf_links_formatted(conversation)), llm_provider=self._resolved_provider, drop_params=drop_params, ) @@ -3028,7 +3069,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _build_provider_specific_fields( self, - completion_response: dict, + completion_response: Mapping[str, object], citations: Sequence[object] | None, thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None, web_search_results: Sequence[object] | None, @@ -3045,6 +3086,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if context_management is not None: provider_specific_fields["context_management"] = context_management + input_transformations: Final = completion_response.get("input_transformations") + if input_transformations is not None: + provider_specific_fields["input_transformations"] = input_transformations + if web_search_results is not None: provider_specific_fields["web_search_results"] = web_search_results @@ -3055,17 +3100,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if tool_results is not None: provider_specific_fields["tool_results"] = tool_results + container: Final = completion_response.get("container") + container_fields: Final[Mapping[str, object] | None] = container if isinstance(container, dict) else None if tool_results is not None and any(result.get("type") != "tool_search_tool_result" for result in tool_results): - container_id: Final = ( - completion_response.get("container", {}).get("id") - if isinstance(completion_response.get("container"), dict) - else None - ) + raw_container_id: Final = container_fields.get("id") if container_fields is not None else None + container_id: Final = raw_container_id if isinstance(raw_container_id, str) else None code_by_id: Final = self._build_code_by_id_map(tool_calls) code_interpreter_results = self._build_code_interpreter_results(tool_results, code_by_id, container_id) provider_specific_fields["code_interpreter_results"] = code_interpreter_results - container: Final = completion_response.get("container") if container is not None: provider_specific_fields["container"] = container @@ -3166,8 +3209,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): STOP_REASON_FIELDS.validate_python(completion_response) ) completion_response["stop_reason"] = "stop" - if stop_reason_fields: - json_mode_message.provider_specific_fields = dict(stop_reason_fields) + input_transformations: Final = provider_specific_fields.get("input_transformations") + json_mode_fields: Final = { + **stop_reason_fields, + **({} if input_transformations is None else {"input_transformations": input_transformations}), + } + if json_mode_fields: + json_mode_message.provider_specific_fields = json_mode_fields if "refusal" in refusal_fields: json_mode_message.refusal = refusal_fields["refusal"] if thinking_blocks: diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 4f5dc678bae..bfd1341711a 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -28,7 +28,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, ) from litellm.exceptions import UnsupportedParamsError -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, lax_bool from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, is_encrypted_reasoning_block, @@ -307,9 +307,10 @@ def count_tokens_feature_betas( ) -def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None) -> bool: +def supports_anthropic_cache_control( + model: str, custom_llm_provider: str | None, model_info: Mapping[str, object] | None = None +) -> bool: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - from litellm.utils import supports_prompt_caching try: provider: Final = custom_llm_provider if custom_llm_provider is not None else get_llm_provider(model=model)[1] @@ -318,10 +319,27 @@ def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None return ( provider in ("anthropic", "bedrock", "vertex_ai", "azure_ai") and "claude" in model.lower() - and supports_prompt_caching(model=model, custom_llm_provider=provider) + and _deployment_supports_prompt_caching(model, provider, model_info) ) +def _deployment_supports_prompt_caching(model: str, provider: str, model_info: Mapping[str, object] | None) -> bool: + from litellm.utils import catalog_with_runtime_registrations, get_model_info, supports_prompt_caching + + declared: Final = None if model_info is None else lax_bool(model_info.get("supports_prompt_caching")) + if declared is not None: + return declared + live: Final = supports_prompt_caching(model=model, custom_llm_provider=provider) + if model_info is None: + return live + try: + key: Final = get_model_info(model=model, custom_llm_provider=provider).get("key") + except Exception: # noqa: BLE001 # an unmapped model keeps the live answer + return live + catalog_entry: Final = catalog_with_runtime_registrations().get(key) + return live if catalog_entry is None else catalog_entry.get("supports_prompt_caching") is True + + def _tool_config_point_control(points: object) -> Mapping[object, object] | None: point_fields: Final = tuple(validated_object_mapping(point) for point in validated_object_list(points) or ()) tool_config_point: Final = next( @@ -1297,7 +1315,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): return None return anthropic_beta_header.split(",") - def get_computer_tool_beta_header(self, computer_tool_version: str) -> str: + def get_computer_tool_beta_header(self, computer_tool_version: str) -> str | None: """ Get the appropriate beta header for a given computer tool version. @@ -1305,16 +1323,14 @@ class AnthropicModelInfo(BaseLLMModelInfo): computer_tool_version: The computer tool version (e.g., 'computer_20250124', 'computer_20241022') Returns: - The corresponding beta header string + The corresponding beta header string, or None for a version with no known beta """ computer_tool_beta_mapping: Final = { + "computer_20251124": "computer-use-2025-11-24", "computer_20250124": "computer-use-2025-01-24", "computer_20241022": "computer-use-2024-10-22", } - return computer_tool_beta_mapping.get( - computer_tool_version, - "computer-use-2024-10-22", # Default fallback - ) + return computer_tool_beta_mapping.get(computer_tool_version) def get_anthropic_beta_list( self, @@ -1346,9 +1362,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): if effort_used: betas.append(ANTHROPIC_EFFORT_BETA_HEADER) # effort-2025-11-24 - if computer_tool_used: - beta_header: Final = self.get_computer_tool_beta_header(computer_tool_used) - betas.append(beta_header) + if computer_tool_used and (computer_tool_beta := self.get_computer_tool_beta_header(computer_tool_used)): + betas.append(computer_tool_beta) # Anthropic no longer requires the prompt-caching beta header # Prompt caching now works automatically when cache_control is used in messages @@ -1437,9 +1452,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): # Anthropic no longer requires the prompt-caching beta header # Prompt caching now works automatically when cache_control is used in messages # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching - if computer_tool_used: - beta_header: Final = self.get_computer_tool_beta_header(computer_tool_used) - betas.add(beta_header) + if computer_tool_used and (computer_tool_beta := self.get_computer_tool_beta_header(computer_tool_used)): + betas.add(computer_tool_beta) # if pdf_used: # betas.add("pdfs-2024-09-25") if file_id_used: diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 1c8f0c206cd..3034927c612 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -49,7 +49,10 @@ _PLAIN_MESSAGE_KEYS: Final = frozenset({"role", "content"}) def _count_image(block: JsonValue) -> JsonValue: if not isinstance(block, dict) or block.get("type") != "image_url": return block - from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_image_obj + from litellm.litellm_core_utils.prompt_templates.factory import ( + anthropic_content_element_factory, + convert_to_anthropic_image_obj, + ) image_block: Final = _IMAGE_BLOCK.validate_python(block) image_url: Final = image_block["image_url"] @@ -57,7 +60,7 @@ def _count_image(block: JsonValue) -> JsonValue: openai_image_url=image_url if isinstance(image_url, str) else image_url["url"], format=image_url.get("format") if isinstance(image_url, dict) else None, ) - image: Final = _COUNT_REQUEST.validate_python({"type": "image", "source": source}) + image: Final = _COUNT_REQUEST.validate_python(anthropic_content_element_factory(source)) return {**{key: value for key, value in block.items() if key not in {"type", "image_url"}}, **image} diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index d13fd4f47ff..31a4e9d9d73 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/litellm/llms/anthropic/messages/adapters/handler.py @@ -15,6 +15,11 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.tool_fields import ( + ROUTES_WITHOUT_MCP_SERVERS, + drops_mcp_servers, + tools_without_mcp_servers, +) from litellm.llms.anthropic.messages.adapters.transformation import ( AnthropicAdapter, ) @@ -610,6 +615,23 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) + provider_kwargs: Final = extra_kwargs or {} + custom_llm_provider: Final = provider_kwargs.get("custom_llm_provider") + target_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None + drop_params: Final = drops_unsupported_params(provider_kwargs.get("drop_params")) + strips_mcp_servers: Final = target_provider in ROUTES_WITHOUT_MCP_SERVERS + request_tools: Final = ( + tools_without_mcp_servers( + provider_kwargs.get("mcp_servers"), + tools or (), + model=model, + route=target_provider, + drop_params=drops_mcp_servers(provider_kwargs), + ) + if strips_mcp_servers + else tools + ) + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, @@ -626,10 +648,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: request_data["temperature"] = temperature if thinking: request_data["thinking"] = thinking - if tool_choice: + if tool_choice and (request_tools or not tools): request_data["tool_choice"] = tool_choice - if tools: - request_data["tools"] = tools + if request_tools: + request_data["tools"] = request_tools if top_k is not None: request_data["top_k"] = top_k if top_p is not None: @@ -647,16 +669,15 @@ class LiteLLMMessagesToCompletionTransformationHandler: if "output_config" in extra_kwargs: request_data["output_config"] = extra_kwargs["output_config"] - custom_llm_provider: Final = extra_kwargs.get("custom_llm_provider") api_base: Final = extra_kwargs.get("api_base") ( openai_request, tool_name_mapping, ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( request_data, - custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + custom_llm_provider=target_provider, api_base=api_base if isinstance(api_base, str) else None, - drop_params=drops_unsupported_params(extra_kwargs.get("drop_params")), + drop_params=drop_params, ) if openai_request is None: @@ -684,7 +705,11 @@ 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 | CONVERTED_STREAM_KEYS | {"anthropic_messages"} + excluded_keys: Final = ( + ANTHROPIC_ONLY_REQUEST_KEYS + | CONVERTED_STREAM_KEYS + | ({"anthropic_messages", "mcp_servers"} if strips_mcp_servers else {"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/messages/adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py index 61d23c2de0d..c41deb00f1d 100644 --- a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py @@ -493,10 +493,23 @@ class _CombinedChunkSplitter: for position, block in enumerate(blocks) ) + @staticmethod + def _without_sibling_finish(chunk: "ModelResponseStream") -> "ModelResponseStream": + choices: Final = _optional_attr_sequence(chunk, "choices") + if ( + len(choices) != 1 + or _optional_attr(choices[0], "index") in (None, 0) + or _optional_attr(choices[0], "finish_reason") is None + ): + return chunk + unfinished: Final = copy.deepcopy(chunk) + unfinished.choices[0].finish_reason = None + return unfinished + def _expand(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]": split_chunks: Final = tuple( split_chunk - for combined_chunk in self._split(chunk) + for combined_chunk in self._split(self._without_sibling_finish(chunk)) for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits ) block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks))) diff --git a/litellm/llms/anthropic/messages/context_management/inference_view.py b/litellm/llms/anthropic/messages/context_management/inference_view.py new file mode 100644 index 00000000000..004beebe838 --- /dev/null +++ b/litellm/llms/anthropic/messages/context_management/inference_view.py @@ -0,0 +1,138 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.types.llms.anthropic import AnthropicMessagesTextParam +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat + + +@dataclass(frozen=True, slots=True) +class InferenceRequestView: + messages: Sequence[Mapping[str, object]] + system: object | None + + +@dataclass(frozen=True, slots=True) +class InferenceRequestViewError: + error_message: str + status_code: int + + +@dataclass(frozen=True, slots=True) +class CountContextEdits: + spec: JsonValue + model: str + tools: tuple[Mapping[str, object], ...] | None + drop_invalid_edits: bool = False + + +_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter( + Mapping[str, object] | tuple[Mapping[str, object], ...] +) +_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...]) + + +def count_context_edits( + context_management: CountContextManagement | None, + tools: Sequence[Mapping[str, object]] | None, + litellm_params: Mapping[str, object], + fallback_model: str, +) -> CountContextEdits: + model: Final = litellm_params.get("model") + return CountContextEdits( + spec=None if context_management is None else context_management.spec, + model=model if isinstance(model, str) else fallback_model, + tools=None if tools is None else tuple(tools), + drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits, + ) + + +def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None: + try: + return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec) + except ValidationError: + return None + + +def _typed_system(system: object | None) -> str | tuple[AnthropicMessagesTextParam, ...] | None | ValidationError: + if system is None or isinstance(system, str): + return system + try: + return _SYSTEM_BLOCKS.validate_python(system) + except ValidationError as error: + return error + + +def _view_without_edits( + messages: Sequence[Mapping[str, object]], + system: object | None, + typed_system: str | Sequence[Mapping[str, object]] | None, +) -> InferenceRequestView: + from litellm.llms.anthropic.messages.context_management.editors.compact import ( + apply_client_compaction_block_history, + ) + + history: Final = apply_client_compaction_block_history( + messages=[dict(message) for message in messages], + system=typed_system + if typed_system is None or isinstance(typed_system, str) + else [dict(block) for block in typed_system], + ) + if history is None: + return InferenceRequestView(messages=messages, system=system) + return InferenceRequestView(messages=history.messages, system=history.system) + + +async def _anthropic_inference_view( + messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits +) -> InferenceRequestView | InferenceRequestViewError: + from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request + from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError + from litellm.llms.anthropic.messages.native.utils import normalize_native_messages + + typed_system: Final = _typed_system(system) + if isinstance(typed_system, ValidationError): + return InferenceRequestView(messages=messages, system=system) + normalized: Final = normalize_native_messages(messages) + spec: Final = _typed_spec(edits.spec) + try: + prepared: Final = await prepare_context_managed_request( + model=edits.model, + messages=normalized, + tools=None if edits.tools is None else [dict(tool) for tool in edits.tools], + system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system, + context_management_spec=( + None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec] + ), + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + count_only=True, + drop_invalid_edits=edits.drop_invalid_edits, + ) + except AnthropicContextManagementError as error: + if error.status_code >= 500: + return _view_without_edits(normalized, system, typed_system) + return InferenceRequestViewError(error_message=error.message, status_code=error.status_code) + if prepared is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=prepared.messages, system=prepared.system) + + +async def inference_request_view( + messages: Sequence[Mapping[str, object]], + system: object | None, + messages_format: TokenCountMessagesFormat, + edits: CountContextEdits, +) -> InferenceRequestView | InferenceRequestViewError: + """Apply the history rewrites and context edits /v1/messages runs before it converts a body for another provider.""" + match messages_format: + case "openai": + return InferenceRequestView(messages=messages, system=system) + case "anthropic": + return await _anthropic_inference_view(messages, system, edits) + case _: + assert_never(messages_format) diff --git a/litellm/llms/anthropic/messages/convertible_blocks.py b/litellm/llms/anthropic/messages/convertible_blocks.py index 09d662da543..b98b4cf3f75 100644 --- a/litellm/llms/anthropic/messages/convertible_blocks.py +++ b/litellm/llms/anthropic/messages/convertible_blocks.py @@ -314,6 +314,76 @@ def keep_non_string_tool_result_texts(rewritten: Sequence[object], original: obj return [_with_user_blocks(message, partial(_restored_block, kept=kept_in_order)) for message in rewritten] +def _url_source_url(block: object) -> str | None: + match as_str_mapping(block): + case {"type": "image" | "document", "source": {"type": "url", "url": str(url)}}: + return url + case _: + return None + + +def _is_document(block: object) -> bool: + match as_str_mapping(block): + case {"type": "document"}: + return True + case _: + return False + + +@dataclass(frozen=True, slots=True, order=True) +class _UrlBlockPlace: + in_tool_result: bool + url: str + + +def _url_sourced_blocks(block: object, in_tool_result: bool = False) -> tuple[tuple[_UrlBlockPlace, object], ...]: + url: Final = _url_source_url(block) + if url is not None: + return ((_UrlBlockPlace(in_tool_result, url), block),) + match as_str_mapping(block): + case {"type": "tool_result", "content": list() as content}: + in_this_tool_result: Final = partial(_url_sourced_blocks, in_tool_result=True) + return tuple(chain.from_iterable(map(in_this_tool_result, _list_items(content)))) + case _: + return () + + +def _entry_place(entry: tuple[_UrlBlockPlace, object]) -> _UrlBlockPlace: + return entry[0] + + +def _entry_block(entry: tuple[_UrlBlockPlace, object]) -> object: + return entry[1] + + +def _restored_url_document( + block: object, originals: Mapping[_UrlBlockPlace, Iterator[object]], in_tool_result: bool = False +) -> object: + url: Final = _url_source_url(block) + if url is not None: + remaining: Final = originals.get(_UrlBlockPlace(in_tool_result, url)) + original: Final = None if remaining is None else next(remaining, None) + return original if _is_document(original) else block + match as_str_mapping(block): + case {"type": "tool_result", "content": list() as content} as tool_result: + restore: Final = partial(_restored_url_document, originals=originals, in_tool_result=True) + return {**tool_result, "content": list(map(restore, _list_items(content)))} + case _: + return block + + +def keep_url_documents(rewritten: Sequence[object], original: object) -> Sequence[object]: + blocks: Final = chain.from_iterable(map(_user_content_blocks, _list_items(original))) + by_place: Final = tuple(sorted(chain.from_iterable(map(_url_sourced_blocks, blocks)), key=_entry_place)) + if not any(_is_document(block) for _, block in by_place): + return rewritten + originals_in_order: Final = MappingProxyType( + {place: iter(tuple(map(_entry_block, group))) for place, group in groupby(by_place, key=_entry_place)} + ) + restore: Final = partial(_restored_url_document, originals=originals_in_order) + return [_with_user_blocks(message, restore) for message in rewritten] + + def _without_non_string_text(block: object) -> object: match as_str_mapping(block): case {"type": "tool_result", "content": list()} as tool_result: diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 1bf00cbdb2f..922de386fe4 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -8,6 +8,7 @@ import asyncio import contextvars from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence +from dataclasses import dataclass from functools import partial from typing import Final, cast @@ -15,6 +16,7 @@ from pydantic import TypeAdapter import litellm from litellm.anthropic_beta_headers_manager import aload_beta_headers_config +from litellm.constants import PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.exception_mapping_utils import exception_type @@ -44,7 +46,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.utils import CallTypes, LlmProviders, LlmProvidersSet from litellm.utils import ProviderConfigManager, client from ..adapters.handler import LiteLLMMessagesToCompletionTransformationHandler, request_safeguards_evaluator @@ -94,6 +96,68 @@ def _responses_mode_is_lost_by_prefix_strip( ) and not _bridges_to_responses_api(resolved_model, custom_llm_provider, deployment_model_info) +@dataclass(frozen=True, slots=True) +class _BridgeTarget: + model: str + api_key: str | None + api_base: str | None + + +def _resolve_requested_provider( + model: str, custom_llm_provider: str | None, litellm_params: GenericLiteLLMParams +) -> tuple[str, str, str | None, str | None]: + resolved: Final = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + litellm_params=litellm_params, + ) + requested_provider: Final = model.partition("/")[0] + if ( + custom_llm_provider in (None, requested_provider) + or requested_provider not in LlmProvidersSet + or requested_provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO + ): + return resolved + own: Final = _resolve_by_model_prefix(model, litellm_params) + return own if own is not None and own[1] == custom_llm_provider else resolved + + +def _resolve_by_model_prefix( + model: str, litellm_params: GenericLiteLLMParams +) -> tuple[str, str, str | None, str | None] | None: + try: + return litellm.get_llm_provider( + model=model, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + litellm_params=litellm_params.model_copy(update={"custom_llm_provider": None}), + ) + except litellm.BadRequestError: + return None + + +def _sets_own_entra_credential(litellm_params: GenericLiteLLMParams) -> bool: + return bool( + litellm_params.azure_ad_token + or (litellm_params.model_extra or {}).get("azure_ad_token_provider") is not None + or (litellm_params.tenant_id and litellm_params.client_id and litellm_params.client_secret) + or (litellm_params.azure_username and litellm_params.azure_password and litellm_params.client_id) + ) + + +def _bridge_target(requested: _BridgeTarget, resolved: _BridgeTarget, custom_llm_provider: str) -> _BridgeTarget: + requested_provider: Final = requested.model.partition("/")[0] + if requested_provider == custom_llm_provider or requested_provider not in LlmProvidersSet: + return requested + return _BridgeTarget( + model=f"{custom_llm_provider}/{resolved.model}", + api_key=requested.api_key if requested.api_base else requested.api_key or resolved.api_key, + api_base=resolved.api_base, + ) + + def _should_route_to_responses_api( custom_llm_provider: str | None, requested_model: str | None = None, @@ -624,13 +688,7 @@ def anthropic_messages_handler( custom_llm_provider, dynamic_api_key, dynamic_api_base, - ) = litellm.get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=litellm_params.api_base, - api_key=litellm_params.api_key, - litellm_params=litellm_params, - ) + ) = _resolve_requested_provider(model, custom_llm_provider, litellm_params) # Store agentic loop params in logging object for agentic hooks # This provides original request context needed for follow-up calls @@ -656,6 +714,16 @@ def anthropic_messages_handler( mock_response=litellm_params.mock_response, ) + bridge: Final = _bridge_target( + _BridgeTarget(model=original_model, api_key=api_key, api_base=api_base), + _BridgeTarget( + model=model, + api_key=None if _sets_own_entra_credential(litellm_params) else dynamic_api_key, + api_base=dynamic_api_base, + ), + custom_llm_provider, + ) + # Expand litellm_proxy MCP references through the MCP gateway before dispatch, so every # downstream path (native passthrough and both bridges) gets real tools rather than a # reference the provider cannot resolve. Popped from kwargs so it never reaches the provider. @@ -682,8 +750,8 @@ def anthropic_messages_handler( top_k=top_k, top_p=top_p, container=container, - api_key=api_key, - api_base=api_base, + api_key=bridge.api_key, + api_base=bridge.api_base, client=client, custom_llm_provider=custom_llm_provider, **kwargs, @@ -700,7 +768,7 @@ def anthropic_messages_handler( return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( max_tokens=max_tokens, messages=messages, - model=original_model, + model=bridge.model, metadata=metadata, stop_sequences=stop_sequences, stream=stream, @@ -716,8 +784,8 @@ def anthropic_messages_handler( top_k=top_k, top_p=top_p, _is_async=is_async, - api_key=api_key, - api_base=api_base, + api_key=bridge.api_key, + api_base=bridge.api_base, client=client, custom_llm_provider=custom_llm_provider, **kwargs, @@ -730,7 +798,7 @@ def anthropic_messages_handler( return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( max_tokens=max_tokens, messages=messages, - model=original_model, + model=bridge.model, metadata=metadata, stop_sequences=stop_sequences, stream=stream, @@ -744,8 +812,8 @@ def anthropic_messages_handler( top_k=top_k, top_p=top_p, _is_async=is_async, - api_key=api_key, - api_base=api_base, + api_key=bridge.api_key, + api_base=bridge.api_base, client=client, custom_llm_provider=custom_llm_provider, **kwargs, diff --git a/litellm/llms/anthropic/messages/responses_adapters/handler.py b/litellm/llms/anthropic/messages/responses_adapters/handler.py index 0eef3fd87c0..ec09987837c 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/handler.py +++ b/litellm/llms/anthropic/messages/responses_adapters/handler.py @@ -13,6 +13,7 @@ import litellm from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, + routed_model_info, supports_openai_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.asyncify import ( @@ -21,9 +22,10 @@ from litellm.litellm_core_utils.asyncify import ( from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( cache_control_as_prompt_cache_breakpoint, - prompt_cache_breakpoint_for_wire, + prompt_cache_breakpoint_with_readable_mode, with_prompt_cache_breakpoint, ) +from litellm.litellm_core_utils.tool_fields import drops_mcp_servers, tools_without_mcp_servers from litellm.responses.utils import auto_reasoning_summary from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic import ( @@ -186,6 +188,7 @@ def _reads_prompt_cache_breakpoint(model: str, forwarded_kwargs: Mapping[str, ob provider if isinstance(provider, str) else None, forwarded_kwargs.get("api_base") or forwarded_kwargs.get("base_url"), forwarded_kwargs.get("prompt_cache_options"), + routed_model_info(forwarded_kwargs.get("model_info")), ) @@ -193,10 +196,11 @@ def _model_takes_prompt_cache_breakpoint(model: str, forwarded_kwargs: Mapping[s provider: Final = forwarded_kwargs.get("custom_llm_provider") provider_name: Final = provider if isinstance(provider, str) else None base_model: Final = forwarded_kwargs.get("base_model") - return supports_openai_prompt_cache_breakpoint(model, provider_name) or ( + model_info: Final = routed_model_info(forwarded_kwargs.get("model_info")) + return supports_openai_prompt_cache_breakpoint(model, provider_name, model_info) or ( isinstance(base_model, str) and bool(base_model) - and supports_openai_prompt_cache_breakpoint(base_model, provider_name) + and supports_openai_prompt_cache_breakpoint(base_model, provider_name, model_info) ) @@ -207,7 +211,7 @@ def _without_prompt_cache_breakpoint(block: Mapping[str, object]) -> Mapping[str def _cache_control_as_wire_safe_prompt_cache_breakpoint(block: Mapping[str, object]) -> Mapping[str, object]: if "prompt_cache_breakpoint" not in block: return cache_control_as_prompt_cache_breakpoint(block) - wire_marker: Final = prompt_cache_breakpoint_for_wire(block["prompt_cache_breakpoint"], drop_params=True) + wire_marker: Final = prompt_cache_breakpoint_with_readable_mode(block["prompt_cache_breakpoint"], drop_params=True) return cache_control_as_prompt_cache_breakpoint( with_prompt_cache_breakpoint(_without_prompt_cache_breakpoint(block), wire_marker) ) @@ -272,6 +276,14 @@ def _build_responses_kwargs( """ forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) reject_malformed_tool_result_text(messages, model, forwarded_kwargs) + target_provider: Final = forwarded_kwargs.get("custom_llm_provider") + request_tools: Final = tools_without_mcp_servers( + forwarded_kwargs.get("mcp_servers"), + tools or (), + model=model, + route=target_provider if isinstance(target_provider, str) else "this route", + drop_params=drops_mcp_servers(forwarded_kwargs), + ) block_rewrite: Final = _prompt_cache_block_rewrite(model, forwarded_kwargs) request_messages: Final[AnthropicRequestMessages] = ( [ @@ -304,10 +316,10 @@ def _build_responses_kwargs( request_data["temperature"] = temperature if thinking: request_data["thinking"] = thinking - if tool_choice: + if tool_choice and (request_tools or not tools): request_data["tool_choice"] = tool_choice - if tools: - request_data["tools"] = tools + if request_tools: + request_data["tools"] = list(request_tools) if top_p is not None: request_data["top_p"] = top_p if output_format: @@ -347,7 +359,7 @@ def _build_responses_kwargs( responses_kwargs["stream"] = True # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.) - excluded: Final = CONVERTED_STREAM_KEYS | {"anthropic_messages"} + excluded: Final = CONVERTED_STREAM_KEYS | {"anthropic_messages", "mcp_servers"} 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/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index df8a0614665..386815412b5 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -1,11 +1,17 @@ -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from typing import TYPE_CHECKING, Final -from openai import AsyncAzureOpenAI, AzureOpenAI +from openai import AsyncAzureOpenAI, AsyncStream, AzureOpenAI, Stream +from openai.types.audio import TranscriptionStreamEvent from pydantic import BaseModel from litellm._uuid import uuid from litellm.litellm_core_utils.audio_utils.utils import get_audio_file_name +from litellm.llms.base_llm.audio_transcription.transformation import ( + acreate_sdk_transcription, + create_sdk_transcription, + sdk_compatible_transcription_request_data, +) from litellm.types.utils import FileTypes from litellm.utils import ( TranscriptionResponse, @@ -38,13 +44,18 @@ class AzureAudioTranscription(AzureChatCompletion): atranscription: bool = False, litellm_params: dict | None = None, custom_llm_provider: str = "azure", - ) -> TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]: + ) -> ( + TranscriptionResponse + | Stream[TranscriptionStreamEvent] + | Coroutine[object, object, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] + ): data: Final = {"model": model, "file": audio_file, **optional_params} + sdk_data: Final = sdk_compatible_transcription_request_data(data) if atranscription is True: return self.async_audio_transcriptions( audio_file=audio_file, - data=data, + data=sdk_data, model_response=model_response, timeout=timeout, api_key=api_key, @@ -58,7 +69,7 @@ class AzureAudioTranscription(AzureChatCompletion): ) azure_client: Final = self.get_azure_openai_client( - api_version=api_version, + api_version=None if self.is_azure_v1_api_version(api_version) else api_version, api_base=api_base, api_key=api_key, model=model, @@ -84,10 +95,9 @@ class AzureAudioTranscription(AzureChatCompletion): }, ) - response: Final = azure_client.audio.transcriptions.create( - **data, - timeout=timeout, - ) + _, response = create_sdk_transcription(azure_client, sdk_data, timeout) + if isinstance(response, Stream): + return response if isinstance(response, BaseModel): stringified_response = response.model_dump() @@ -114,7 +124,7 @@ class AzureAudioTranscription(AzureChatCompletion): self, audio_file: FileTypes, model: str, - data: dict, + data: Mapping[str, object], model_response: TranscriptionResponse, timeout: float, logging_obj: "LiteLLMLoggingObj", @@ -125,7 +135,7 @@ class AzureAudioTranscription(AzureChatCompletion): max_retries=None, litellm_params: dict | None = None, custom_llm_provider: str = "azure", - ) -> TranscriptionResponse: + ) -> TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]: response = None try: async_azure_client: Final = self.get_azure_openai_client( @@ -155,12 +165,9 @@ class AzureAudioTranscription(AzureChatCompletion): }, ) - raw_response: Final = await async_azure_client.audio.transcriptions.with_raw_response.create( - **data, timeout=timeout - ) - - headers: Final = dict(raw_response.headers) - response = raw_response.parse() + headers, response = await acreate_sdk_transcription(async_azure_client, data, timeout) + if isinstance(response, AsyncStream): + return response if isinstance(response, BaseModel): stringified_response = response.model_dump() diff --git a/litellm/llms/azure/chat/gpt_reasoning_transformation.py b/litellm/llms/azure/chat/gpt_reasoning_transformation.py index 72a0500e0af..29305866e23 100644 --- a/litellm/llms/azure/chat/gpt_reasoning_transformation.py +++ b/litellm/llms/azure/chat/gpt_reasoning_transformation.py @@ -1,6 +1,6 @@ """Support for Azure OpenAI gpt-5 model family.""" -from collections.abc import Collection +from collections.abc import Collection, Mapping from typing import ClassVar, Final import litellm @@ -105,7 +105,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): model: str, messages: list[AllMessageValues], optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], headers: dict, ) -> dict: model = model.replace(self.GPT5_SERIES_ROUTE, "") diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 990c8869de3..0fe7caa1e16 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -1,16 +1,26 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import ( + TYPE_CHECKING, + Any, + Final, + cast, # noqa: TID251 # non_default_params reaches the mapper as an untyped dict +) from httpx._models import Headers, Response import litellm +from litellm.integrations.anthropic_cache_control_hook import deployment_prompt_cache_breakpoint_support +from litellm.litellm_core_utils.allowed_tools_tool_choice import is_allowed_tools_tool_choice +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( + chat_completion_prompt_cache_breakpoints_for_wire, drop_tool_reference_parts_from_tool_messages, flatten_combinators_and_drop_non_python_regex_patterns, hoist_images_from_tool_messages, system_messages_first, tool_with_sanitized_parameters, + without_prompt_cache_breakpoints, ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, @@ -241,12 +251,12 @@ class AzureOpenAIConfig(BaseConfig): or (api_version_year == "2023" and api_version_month < "12") or (api_version_year == "2023" and api_version_month == "12" and api_version_day < "01") ): - if litellm.drop_params is True or (drop_params is not None and drop_params is True): - pass - else: + if is_allowed_tools_tool_choice( + cast(object, value) # cast-ok: widening an untyped dict value to object is always sound + ) or not (litellm.drop_params is True or (drop_params is not None and drop_params is True)): raise UnsupportedParamsError( status_code=400, - message=f"""Azure does not support 'tool_choice', for api_version={api_version}. Bump your API version to '2023-12-01-preview' or later. This parameter requires 'api_version="2023-12-01-preview"' or later. Azure API Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions""", + message=f"""Azure does not support 'tool_choice', for api_version={api_version}. Bump your API version to '2023-12-01-preview' or later. This parameter requires 'api_version="2023-12-01-preview"' or later. drop_params never drops an allowed_tools tool_choice, since that would let the model call tools the request excluded. Azure API Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions""", ) elif value == "required" and ( api_version_year == "2024" and api_version_month <= "05" @@ -290,11 +300,16 @@ class AzureOpenAIConfig(BaseConfig): model: str, messages: list[AllMessageValues], optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], headers: dict, ) -> dict: ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages - stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages) + wire_messages: Final = ( + self._with_wire_safe_prompt_cache_breakpoints(model, ordered_messages, litellm_params) + if drops_unsupported_params(litellm_params.get("drop_params")) + else ordered_messages + ) + stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(wire_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) request_params: Final = MappingProxyType( { @@ -310,6 +325,14 @@ class AzureOpenAIConfig(BaseConfig): **sanitized_tools_update(optional_params), } + @staticmethod + def _with_wire_safe_prompt_cache_breakpoints( + model: str, messages: list[AllMessageValues], litellm_params: Mapping[str, object] + ) -> list[AllMessageValues]: + if deployment_prompt_cache_breakpoint_support(model, "azure", litellm_params) is False: + return [without_prompt_cache_breakpoints(message) for message in messages] + return [chat_completion_prompt_cache_breakpoints_for_wire(message, drop_params=True) for message in messages] + def transform_response( self, model: str, diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 3d4576fd7ac..16ec4249244 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -5,9 +5,11 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. """ from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import Final, Protocol, cast +import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.url_utils import join_api_base @@ -48,17 +50,55 @@ async def forward_messages(client_ws: _RealtimeTextSink, backend_ws: _RealtimeTe pass +def azure_realtime_requires_ga(model: str) -> bool: + try: + azure_model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="azure") + except Exception: # noqa: BLE001 # unmapped deployments can select a protocol explicitly + try: + openai_model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="openai") + except Exception: # noqa: BLE001 # unmapped deployments can select a protocol explicitly + return False + openai_entry: Final = openai_model_info.get("provider_specific_entry") + return openai_entry is not None and openai_entry.get("realtime_ga_only") == 1 + azure_entry: Final = azure_model_info.get("provider_specific_entry") + return azure_entry is not None and azure_entry.get("realtime_ga_only") == 1 + + +def _is_ga_protocol(protocol: str | None) -> bool: + return protocol is not None and protocol.upper() in ("GA", "V1") + + +@dataclass(frozen=True, slots=True) +class AzureRealtimeProtocol: + name: str + bridges_beta_client_to_ga: bool = False + + def azure_realtime_protocol_for_client( configured_protocol: object, *, + model: str, + base_model: str | None = None, + realtime_mode: str, query_params: RealtimeQueryParams | None, websocket: ScopedWebSocket, -) -> str: +) -> AzureRealtimeProtocol: + serves_a_ga_only_session: Final = ( + realtime_mode == "translation" or (query_params or {}).get("intent") == "transcription" + ) + client_wants_beta: Final = client_sent_openai_beta_realtime_header(websocket) if isinstance(configured_protocol, str) and configured_protocol: - return configured_protocol - if (query_params or {}).get("intent") == "transcription": - return "GA" - return "beta" if client_sent_openai_beta_realtime_header(websocket) else "GA" + return AzureRealtimeProtocol( + name=configured_protocol, + bridges_beta_client_to_ga=client_wants_beta + and not serves_a_ga_only_session + and _is_ga_protocol(configured_protocol), + ) + if serves_a_ga_only_session: + return AzureRealtimeProtocol(name="GA") + if azure_realtime_requires_ga(base_model or model): + return AzureRealtimeProtocol(name="GA", bridges_beta_client_to_ga=client_wants_beta) + return AzureRealtimeProtocol(name="beta" if client_wants_beta else "GA") class _ProxyClientWebSocket(Protocol): @@ -92,6 +132,7 @@ class AzureOpenAIRealtime(AzureChatCompletion): api_version: str | None, realtime_protocol: str | None = None, query_params: RealtimeQueryParams | None = None, + realtime_mode: str = "realtime", ) -> str: """ Construct Azure realtime WebSocket URL. @@ -104,6 +145,7 @@ class AzureOpenAIRealtime(AzureChatCompletion): - "GA" or "v1": Uses /openai/v1/realtime (GA path) - "beta" or None: Uses /openai/realtime (beta path, default) query_params: Extra query params to forward (e.g. intent=transcription). + realtime_mode: Realtime surface to target, "realtime" or "translation". Returns: WebSocket URL string @@ -117,24 +159,29 @@ class AzureOpenAIRealtime(AzureChatCompletion): api_base = api_base.replace("https://", "wss://") # Determine path based on realtime_protocol (case-insensitive) - _is_ga: Final = realtime_protocol is not None and realtime_protocol.upper() in ( - "GA", - "V1", - ) + _is_ga: Final = _is_ga_protocol(realtime_protocol) intent: Final = (query_params or {}).get("intent") - if _is_ga: - path = "/openai/v1/realtime" - query_parts = [] - if intent != "transcription" and (query_params is None or "model" in query_params): - query_parts.append(urlencode({"model": model})) - else: - # Default to beta path for backwards compatibility - path = "/openai/realtime" - query_parts = [urlencode({"api-version": api_version, "deployment": model})] + path: Final = ( + "/openai/v1/realtime/translations" + if realtime_mode == "translation" + else "/openai/v1/realtime" + if _is_ga + else "/openai/realtime" + ) + base_query_parts: Final = ( + (urlencode((("model", model),)),) + if realtime_mode == "translation" + else ( + (urlencode((("model", model),)),) + if intent != "transcription" and (query_params is None or "model" in query_params) + else () + ) + if _is_ga + else (urlencode((("api-version", api_version), ("deployment", model))),) + ) - if intent: - query_parts.append(urlencode({"intent": intent})) + query_parts: Final = (*base_query_parts, urlencode((("intent", intent),))) if intent else base_query_parts qs: Final = "&".join(query_parts) endpoint_url: Final = join_api_base(api_base, path) @@ -165,13 +212,16 @@ class AzureOpenAIRealtime(AzureChatCompletion): query_params: RealtimeQueryParams | None = None, user_api_key_dict: object | None = None, litellm_metadata: dict | None = None, + realtime_mode: str = "realtime", + bridges_beta_client_to_ga: bool = False, + drop_params: bool = False, ): import websockets from websockets.asyncio.client import ClientConnection if api_base is None: raise ValueError("api_base is required for Azure OpenAI calls") - backend_uses_beta_protocol: Final = realtime_protocol is None or realtime_protocol.upper() not in ("GA", "V1") + backend_uses_beta_protocol: Final = not _is_ga_protocol(realtime_protocol) if api_version is None and backend_uses_beta_protocol: raise ValueError("api_version is required for Azure OpenAI calls") @@ -181,6 +231,7 @@ class AzureOpenAIRealtime(AzureChatCompletion): api_version, realtime_protocol=realtime_protocol, query_params=query_params, + realtime_mode=realtime_mode, ) auth_headers: Final = self.get_auth_headers(api_key=api_key, azure_ad_token=azure_ad_token) @@ -204,6 +255,9 @@ class AzureOpenAIRealtime(AzureChatCompletion): force_transcription_model=( model if (query_params or {}).get("intent") == "transcription" else None ), + translation_session=realtime_mode == "translation", + bridges_beta_client_to_ga=bridges_beta_client_to_ga, + drop_params=drop_params, ) await realtime_streaming.bidirectional_forward() diff --git a/litellm/llms/azure/realtime/http_transformation.py b/litellm/llms/azure/realtime/http_transformation.py index 545c3348182..64051b2c3c9 100644 --- a/litellm/llms/azure/realtime/http_transformation.py +++ b/litellm/llms/azure/realtime/http_transformation.py @@ -1,11 +1,20 @@ """Azure OpenAI realtime HTTP transformation config (client_secrets + realtime_calls).""" +from collections.abc import Mapping from typing import Final import litellm from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.azure.common_utils import get_azure_ad_token from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams + +_CALLER_AUTH_HEADERS: Final = frozenset({"authorization", "api-key"}) + + +def _azure_resource_base(api_base: str) -> str: + return api_base.rstrip("/").removesuffix("/openai/v1").removesuffix("/openai") def _azure_realtime_v1_url(api_base: str, endpoint: str) -> str: @@ -21,6 +30,19 @@ class AzureRealtimeHTTPConfig(BaseRealtimeHTTPConfig): def get_api_key(self, api_key: str | None, **kwargs: object) -> str: return api_key or litellm.api_key or get_secret_str("AZURE_API_KEY") or "" + def get_extra_headers( + self, + litellm_params: GenericLiteLLMParams, + api_key: str, + extra_headers: Mapping[str, object] | None, + ) -> Mapping[str, object] | None: + if api_key or any(name.lower() in _CALLER_AUTH_HEADERS for name in extra_headers or {}): + return extra_headers + azure_ad_token: Final = get_azure_ad_token(litellm_params) + if azure_ad_token is None: + return extra_headers + return {**(extra_headers or {}), "Authorization": f"Bearer {azure_ad_token}"} + def get_complete_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: return _azure_realtime_v1_url(self.get_api_base(api_base), "client_secrets") @@ -30,20 +52,30 @@ class AzureRealtimeHTTPConfig(BaseRealtimeHTTPConfig): model: str, api_key: str | None = None, ) -> dict: - return { + validated_headers: Final = { **headers, - "api-key": api_key or "", "Content-Type": "application/json", } + if api_key: + validated_headers["api-key"] = api_key + return validated_headers def get_realtime_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: return _azure_realtime_v1_url(self.get_api_base(api_base), "calls") def get_transcription_session_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: - base: Final = self.get_api_base(api_base) + base: Final = _azure_resource_base(self.get_api_base(api_base)) version: Final = api_version or get_secret_str("AZURE_API_VERSION") or "2024-12-17" return f"{join_api_base(base, 'openai/realtime/transcription_sessions')}?api-version={version}" + def get_translation_client_secret_url( + self, api_base: str | None, model: str, api_version: str | None = None + ) -> str: + return _azure_realtime_v1_url(self.get_api_base(api_base), "translations/client_secrets") + + def get_translation_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: + return _azure_realtime_v1_url(self.get_api_base(api_base), "translations/calls") + def get_realtime_calls_headers(self, ephemeral_key: str) -> dict[str, str]: return { "Authorization": f"Bearer {ephemeral_key}", diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 2b757bb10d9..6c954d00e64 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -1,6 +1,7 @@ import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from copy import deepcopy +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal from urllib.parse import urlsplit @@ -11,6 +12,7 @@ from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION +from litellm.litellm_core_utils.prompt_templates.common_utils import prompt_cache_breakpoint_within from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.common_utils import BaseAzureLLM @@ -29,6 +31,7 @@ else: LiteLLMLoggingObj = Any _DECODED_JSON: Final = TypeAdapter(object) +_AZURE_RESPONSES_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(AzureResponsesPromptCacheBreakpoint) _DATED_GA_API_VERSION: Final = re.compile(r"\d{4}-\d{2}-\d{2}") @@ -81,6 +84,13 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def _is_gpt_5_model(model: str) -> bool: return AzureOpenAIGPT5Config.is_model_gpt_5_model(model) + def _prompt_cache_breakpoint_to_wire( + self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None + ) -> Callable[[object], object]: + if reads_breakpoints or known_support is True: + return partial(prompt_cache_breakpoint_within, _AZURE_RESPONSES_PROMPT_CACHE_BREAKPOINT) + return super()._prompt_cache_breakpoint_to_wire(litellm_params, reads_breakpoints, known_support) + @staticmethod def _supports_reasoning_effort_none(model: str) -> bool: return AzureOpenAIGPT5Config.supports_reasoning_effort_level(model, "none") diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index c324f0f0980..97ee44bd3bb 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -12,6 +12,7 @@ from pydantic import JsonValue import litellm from litellm._logging import verbose_logger +from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.common_utils import AnthropicError @@ -39,6 +40,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): tools: list[dict[str, JsonValue]] | None = None, system: object = None, optional_params: Mapping[str, JsonValue] | None = None, + context_management: JsonValue = None, ) -> dict[str, object]: """ Handle a CountTokens request using httpx with Azure authentication. @@ -70,6 +72,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): tools=tools, system=system, optional_params=optional_params, + context_management=context_management, drop_params=drops_unsupported_params((litellm_params or {}).get("drop_params")), ) @@ -80,11 +83,14 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): verbose_logger.debug("Making request to: %s", endpoint_url) + if context_management is not None: + await aload_beta_headers_config() # Get required headers with Azure authentication headers: Final = self.get_required_headers( api_key=api_key, litellm_params=litellm_params, messages=messages, + context_management=context_management, ) if not {"x-api-key", "Authorization"} & headers.keys(): raise AnthropicError( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index af7cf9d81ce..14f7d4bfd2b 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -135,6 +135,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): tools=tools, system=system, optional_params=count_options, + context_management=None if context_management is None else context_management.spec, ) if result is not None: diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index c02c92a95d6..91ed40d9271 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -8,9 +8,16 @@ from collections.abc import Sequence from typing import Final from urllib.parse import urlsplit +from pydantic import JsonValue + +from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.anthropic.common_utils import count_tokens_feature_betas, merge_anthropic_beta_headers +from litellm.llms.anthropic.common_utils import ( + context_management_betas, + count_tokens_feature_betas, + merge_anthropic_beta_headers, +) from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -31,6 +38,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): api_key: str | None, litellm_params: dict[str, object] | None = None, messages: Sequence[object] = (), + context_management: JsonValue = None, ) -> dict[str, str]: """ Get the required headers for the Azure AI Anthropic CountTokens API. @@ -42,6 +50,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): api_key: The Azure AI API key litellm_params: Optional LiteLLM parameters for additional auth config messages: The request messages, read for betas the counted features need + context_management: The edit spec the count forwards, read for the betas its edits need Returns: Dictionary of required headers with both x-api-key and Azure authentication @@ -51,7 +60,11 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": merge_anthropic_beta_headers( - ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, count_tokens_feature_betas("azure_ai", {}, messages) + ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + ( + *count_tokens_feature_betas("azure_ai", {}, messages), + *filter_and_transform_beta_headers(list(context_management_betas(context_management)), "azure_ai"), + ), ), **({"x-api-key": api_key} if api_key else {}), } diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 5214d31fb2e..487958d6d55 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -10,7 +10,11 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger -from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook +from litellm.integrations.anthropic_cache_control_hook import ( + AnthropicCacheControlHook, + declared_prompt_cache_breakpoint, + routed_model_info, +) from litellm.litellm_core_utils.prompt_templates.common_utils import ( content_parts, convert_content_list_to_str, @@ -125,18 +129,14 @@ def _with_breakpoints_for_foundry(message: AllMessageValues, drop_params: bool) def _deployment_reads_breakpoints(model: str, litellm_params: Mapping[str, object]) -> bool: + model_info: Final = routed_model_info(litellm_params.get("model_info")) + declared: Final = declared_prompt_cache_breakpoint(model_info) + if declared is not None: + return declared base_model: Final = litellm_params.get("base_model") named_models: Final = (base_model, model) if isinstance(base_model, str) and base_model else (model,) - model_info: Final = litellm_params.get("model_info") - if not isinstance(model_info, Mapping): - return any( - AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(name, "azure_ai") for name in named_models - ) - deployment_flag: Final = _FIELDS.validate_python(model_info).get("supports_prompt_cache_breakpoint") - if isinstance(deployment_flag, bool): - return deployment_flag return any( - AnthropicCacheControlHook.catalog_targets_openai_prompt_cache_breakpoint(name, "azure_ai") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(name, "azure_ai", model_info=model_info) for name in named_models ) diff --git a/litellm/llms/base_llm/audio_transcription/transformation.py b/litellm/llms/base_llm/audio_transcription/transformation.py index e4bf148abf3..68365372e7e 100644 --- a/litellm/llms/base_llm/audio_transcription/transformation.py +++ b/litellm/llms/base_llm/audio_transcription/transformation.py @@ -1,9 +1,30 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final +from types import MappingProxyType +from typing import ( + TYPE_CHECKING, + Any, + Final, + cast, # noqa: TID251 # SDK file handles have no runtime validation schema +) import httpx +from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream +from openai.types.audio import ( + Transcription, + TranscriptionDiarized, + TranscriptionStreamEvent, + TranscriptionVerbose, +) +from openai.types.audio.transcription_create_params import ( + TranscriptionCreateParamsBase, + TranscriptionCreateParamsStreaming, +) +from pydantic import BaseModel, TypeAdapter +from litellm.litellm_core_utils.audio_utils.transcription_streaming import transcription_stream_requested +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.types.llms.openai import ( AllMessageValues, @@ -11,6 +32,31 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import FileTypes, ModelResponse, TranscriptionResponse +_TRANSCRIPTION_REQUEST: Final = TypeAdapter(Mapping[str, object]) +_TRANSCRIPTION_TEXT: Final = TypeAdapter(str) +_TRANSCRIPTION_BODY: Final = TypeAdapter(dict[str, object]) +_TRANSCRIPTION_TYPES_BY_FORMAT: Final = MappingProxyType( + {"verbose_json": TranscriptionVerbose, "diarized_json": TranscriptionDiarized} +) + + +def sdk_compatible_transcription_request_data(data: object) -> Mapping[str, object]: + request_data: Final = _TRANSCRIPTION_REQUEST.validate_python(data) + extension_keys: Final = ("keywords", "languages") + extension_body: Final = {key: request_data[key] for key in extension_keys if key in request_data} + if not extension_body: + return request_data + + existing_extra_body: Final = request_data.get("extra_body") + return { + **{key: value for key, value in request_data.items() if key not in extension_keys}, + "extra_body": { + **(as_str_mapping(existing_extra_body) or {}), + **extension_body, + }, + } + + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer @@ -20,6 +66,77 @@ else: LiteLLMLoggingObj = Any +def sdk_compatible_transcription_response(response: object) -> BaseModel | str: + if isinstance(response, BaseModel): + return response + return _TRANSCRIPTION_TEXT.validate_python(response, strict=True) + + +def _content_type(http_response: httpx.Response) -> str: + return next(iter(http_response.headers.get_list("content-type")), "") + + +def _is_event_stream(http_response: httpx.Response) -> bool: + return _content_type(http_response).startswith("text/event-stream") + + +def _unstreamed_transcription_response(http_response: httpx.Response, response_format: object) -> BaseModel | str: + if not _content_type(http_response).startswith("application/json"): + return http_response.text + transcription_type: Final = _TRANSCRIPTION_TYPES_BY_FORMAT.get(str(response_format), Transcription) + return transcription_type.model_construct(None, **_TRANSCRIPTION_BODY.validate_json(http_response.content)) + + +def _non_streaming_params(data: Mapping[str, object]) -> TranscriptionCreateParamsBase: + return cast( # cast-ok: SDK file handles have no runtime validation schema + TranscriptionCreateParamsBase, {key: value for key, value in data.items() if key != "stream"} + ) + + +def _streaming_params(data: Mapping[str, object]) -> TranscriptionCreateParamsStreaming: + return cast( # cast-ok: SDK file handles have no runtime validation schema + TranscriptionCreateParamsStreaming, data + ) + + +def create_sdk_transcription( + client: OpenAI, data: Mapping[str, object], timeout: float | httpx.Timeout +) -> tuple[dict[str, str], Stream[TranscriptionStreamEvent] | BaseModel | str]: + if not transcription_stream_requested(data.get("stream")): + raw_response: Final = client.audio.transcriptions.with_raw_response.create( + **_non_streaming_params(data), timeout=timeout + ) + return dict(raw_response.headers), sdk_compatible_transcription_response(raw_response.parse()) + raw_stream: Final = client.audio.transcriptions.with_raw_response.create( + **_streaming_params({**data, "stream": True}), timeout=timeout + ) + if _is_event_stream(raw_stream.http_response): + return dict(raw_stream.headers), raw_stream.parse(to=Stream[TranscriptionStreamEvent]) + raw_stream.http_response.read() + return dict(raw_stream.headers), _unstreamed_transcription_response( + raw_stream.http_response, data.get("response_format") + ) + + +async def acreate_sdk_transcription( + client: AsyncOpenAI, data: Mapping[str, object], timeout: float | httpx.Timeout +) -> tuple[dict[str, str], AsyncStream[TranscriptionStreamEvent] | BaseModel | str]: + if not transcription_stream_requested(data.get("stream")): + raw_response: Final = await client.audio.transcriptions.with_raw_response.create( + **_non_streaming_params(data), timeout=timeout + ) + return dict(raw_response.headers), sdk_compatible_transcription_response(raw_response.parse()) + raw_stream: Final = await client.audio.transcriptions.with_raw_response.create( + **_streaming_params({**data, "stream": True}), timeout=timeout + ) + if _is_event_stream(raw_stream.http_response): + return dict(raw_stream.headers), raw_stream.parse(to=AsyncStream[TranscriptionStreamEvent]) + await raw_stream.http_response.aread() + return dict(raw_stream.headers), _unstreamed_transcription_response( + raw_stream.http_response, data.get("response_format") + ) + + @dataclass class AudioTranscriptionRequestData: """ diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index c562c8d0f04..254783edf39 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -431,6 +431,10 @@ class DecisionsResponseValidationError(BaseLLMException): pass +def _lacks_confidence(answer: DecisionsIRAnswer) -> bool: + return isinstance(answer, (DecisionsIRChoiceAnswer, DecisionsIRScoreAnswer)) and answer.confidence is None + + def parse_systemone_response(payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse: return systemone_response_to_ir( _SYSTEMONE_RESPONSE_ADAPTER.validate_python(_normalize_systemone_response(payload)), request @@ -534,6 +538,8 @@ class BaseDecisionsConfig(ABC): max_total_image_bytes: int | None = None image_media_types: frozenset[str] | None = None inlines_remote_images: bool = False + openai_evidence_fields: frozenset[str] = frozenset() + confidence_optional: bool = False health_check_questions: Mapping[str, Mapping[str, object]] = MappingProxyType( {"reachable": MappingProxyType({"type": "noul", "instructions": "Is the service reachable?"})} ) @@ -682,13 +688,20 @@ class BaseDecisionsConfig(ABC): ) -> DecisionsIRResponse: try: payload: Final[object] = _PAYLOAD_ADAPTER.validate_json(raw_response.content) - return self.parse_response(payload, request) + response: Final = self.parse_response(payload, request) except ValidationError as error: raise DecisionsResponseValidationError( status_code=500, message=f"Decisions provider '{custom_llm_provider}' returned an unexpected response: " f"{error.title} {_validation_summary(error)}", ) from error + if not self.confidence_optional and any(_lacks_confidence(answer) for answer in response.answers): + raise DecisionsResponseValidationError( + status_code=500, + message=f"Decisions provider '{custom_llm_provider}' returned an unexpected response: " + "a choice or score answer has no confidence", + ) + return response def get_error_class( self, diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 232d9d8eb42..0c7054326b7 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -164,7 +164,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): tools: ToolConfigDict | None, generate_content_config_dict: dict, system_instruction: object | None = None, - ) -> dict: + ) -> dict[str, object]: """ Transform the request parameters for the generate content API. diff --git a/litellm/llms/base_llm/realtime/http_transformation.py b/litellm/llms/base_llm/realtime/http_transformation.py index a02054ec8cb..c1c943c2cd6 100644 --- a/litellm/llms/base_llm/realtime/http_transformation.py +++ b/litellm/llms/base_llm/realtime/http_transformation.py @@ -7,10 +7,13 @@ These are HTTP (not WebSocket) endpoints used by the WebRTC flow: """ from abc import ABC, abstractmethod +from collections.abc import Mapping +from typing import Final import httpx from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.types.router import GenericLiteLLMParams class BaseRealtimeHTTPConfig(ABC): @@ -50,6 +53,14 @@ class BaseRealtimeHTTPConfig(ABC): explicit api_key → litellm.api_key → env var → "" """ + def get_extra_headers( + self, + litellm_params: GenericLiteLLMParams, + api_key: str, + extra_headers: Mapping[str, object] | None, + ) -> Mapping[str, object] | None: + return extra_headers + # ------------------------------------------------------------------ # # client_secrets endpoint # # ------------------------------------------------------------------ # @@ -64,6 +75,12 @@ class BaseRealtimeHTTPConfig(ABC): base = base.removesuffix("/v1") return join_api_base(base, "v1/realtime/transcription_sessions") + def get_translation_client_secret_url( + self, api_base: str | None, model: str, api_version: str | None = None + ) -> str: + base: Final = (api_base or "").rstrip("/") + return f"{base}/v1/realtime/translations/client_secrets" + @abstractmethod def validate_environment( self, @@ -86,6 +103,10 @@ class BaseRealtimeHTTPConfig(ABC): """Return the full URL for POST /realtime/calls (SDP exchange).""" return join_api_base(api_base or "", "v1/realtime/calls") + def get_translation_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: + base: Final = (api_base or "").rstrip("/") + return f"{base}/v1/realtime/translations/calls" + def get_realtime_calls_headers(self, ephemeral_key: str) -> dict: """ Build headers for the realtime_calls POST. diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 21af97e523e..056f34f1bf4 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -23,6 +23,11 @@ from litellm.constants import ( BEDROCK_MIN_THINKING_BUDGET_TOKENS, RESPONSE_FORMAT_TOOL_NAME, ) +from litellm.integrations.anthropic_cache_control_hook import ( + MAX_CACHE_CONTROL_BLOCKS, + AnthropicCacheControlHook, + configured_injection_points, +) from litellm.litellm_core_utils.core_helpers import ( as_str_mapping, drops_unsupported_params, @@ -52,6 +57,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import parts_of, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, @@ -63,6 +69,7 @@ from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.base_llm.chat.transformation import REQUEST_PARAMS, BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.tool_result_images import place_tool_result_images from litellm.llms.bedrock.common_utils import ( + bedrock_chat_opts_into_native_structured_output, bedrock_chat_uses_native_structured_output, bedrock_cleaned_output_schema, bedrock_model_is_unknown, @@ -129,6 +136,7 @@ from ..common_utils import ( bedrock_capability_model, bedrock_converse_supports_parallel_tool_use_config, bedrock_model_accepts_cache_points, + bedrock_model_known_to_accept_cache_points, bedrock_reasoning_effort_disabled, bedrock_rejects_stop_sequences, get_anthropic_beta_from_headers, @@ -202,6 +210,11 @@ def _has_native_text_format(output_config: object) -> bool: return False +def _cache_mark_kept(block: Mapping[str, object], keeps_cache_marks: bool) -> Mapping[str, object]: + cache_control: Final = block.get("cache_control") + return {"cache_control": cache_control} if keeps_cache_marks and cache_control is not None else {} + + class AmazonConverseConfig(BaseConfig): """ Reference - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html @@ -241,7 +254,7 @@ class AmazonConverseConfig(BaseConfig): @staticmethod def _convert_consecutive_user_messages_to_guarded_text( - messages: list[AllMessageValues], optional_params: dict + messages: list[AllMessageValues], optional_params: dict, keeps_cache_marks: bool = False ) -> list[AllMessageValues]: """ Convert consecutive user messages at the end to guarded_text type if guardrailConfig is present @@ -279,7 +292,11 @@ class AmazonConverseConfig(BaseConfig): new_content = [] for item in content: if isinstance(item, dict) and item.get("type") == "text": - new_item = {"type": "guarded_text", "text": item["text"]} + new_item = { + "type": "guarded_text", + "text": item["text"], + **_cache_mark_kept(_JSON_OBJECT.validate_python(item), keeps_cache_marks), + } new_content.append(new_item) else: new_content.append(item) @@ -287,7 +304,13 @@ class AmazonConverseConfig(BaseConfig): messages_copy[user_message_index]["content"] = new_content elif isinstance(content, str): # If content is a string, convert it to guarded_text - messages_copy[user_message_index]["content"] = [{"type": "guarded_text", "text": content}] + messages_copy[user_message_index]["content"] = [ + { + "type": "guarded_text", + "text": content, + **_cache_mark_kept(_JSON_OBJECT.validate_python(user_message), keeps_cache_marks), + } + ] return messages_copy @@ -933,8 +956,9 @@ class AmazonConverseConfig(BaseConfig): """Check if the Bedrock model takes native structured outputs (outputConfig.textFormat) for this request. ``supports_native_structured_output`` in ``litellm.model_cost`` turns it on for every request, and - ``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) only when the - request asks for thinking, which the json tool path would drop. + ``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) also when the + request asks for thinking, which the json tool path would drop. Without thinking, a model that cannot + force the json tool gets it unless that key is false. Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html """ return bedrock_chat_uses_native_structured_output( @@ -1068,20 +1092,23 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool, model_info: Mapping[str, object] | None = None, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, model=model, route="bedrock", drop_params=drop_params + ) + is_thinking_enabled: Final = self.is_thinking_enabled(request_params) base_model: Final = BedrockModelInfo.get_base_model(model) drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model) drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param( model=model, base_model=base_model ) - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param == "response_format" and isinstance(value, dict): optional_params = self._translate_response_format_param( value=value, model=model, optional_params=optional_params, - non_default_params=non_default_params, + non_default_params=request_params, is_thinking_enabled=is_thinking_enabled, model_info=model_info, ) @@ -1135,7 +1162,7 @@ class AmazonConverseConfig(BaseConfig): ), ), model=model, - non_default_params=non_default_params, + non_default_params=request_params, optional_params=optional_params, ) if param == "tool_choice": @@ -1160,7 +1187,7 @@ class AmazonConverseConfig(BaseConfig): and value.get("type") == "adaptive" and not AnthropicConfig.is_adaptive_thinking_model(model, "bedrock") ): - max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens") + max_tokens = request_params.get("max_completion_tokens") or request_params.get("max_tokens") legacy_thinking = AnthropicConfig.map_reasoning_effort( reasoning_effort="medium", model=model, @@ -1230,7 +1257,7 @@ class AmazonConverseConfig(BaseConfig): # Nova 2 handles token budgeting differently through reasoningConfig if "gpt-oss" not in model and not self._is_nova_2_model(model): self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=optional_params + non_default_params=request_params, optional_params=optional_params ) final_is_thinking_enabled: Final = self.is_thinking_enabled(optional_params) @@ -1312,19 +1339,15 @@ class AmazonConverseConfig(BaseConfig): if "type" in value and value["type"] == "text": return optional_params + thinking_requested: Final = self.is_thinking_requested(REQUEST_PARAMS.validate_python(non_default_params)) native_schema: Final = ( - self._native_output_schema( - model, - json_schema, - self.is_thinking_requested(REQUEST_PARAMS.validate_python(non_default_params)), - model_info, - ) + self._native_output_schema(model, json_schema, thinking_requested, model_info) if json_schema is not None else None ) if native_schema is not None: # Use Bedrock's native structured outputs API (outputConfig.textFormat) - # No synthetic tool injection, no fake_stream needed. + # No synthetic tool injection. # Requires an explicit schema — json_object with no schema falls through # to the tool-call path below. output_config: Final = self._create_output_config_for_response_format( @@ -1350,8 +1373,6 @@ class AmazonConverseConfig(BaseConfig): optional_params["tool_choice"] = ToolChoiceValuesBlock( tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME) ) - if non_default_params.get("stream", False) is True: - optional_params["fake_stream"] = True # else: response_format=json_object with no schema. # Don't inject the synthetic json_tool_call tool here. When no # schema is given, _create_json_tool_call_for_response_format @@ -1359,6 +1380,13 @@ class AmazonConverseConfig(BaseConfig): # returns {} instead of the requested JSON. The model already # returns JSON when the prompt asks for it. + # A request that goes native only because the model can't force the json tool keeps the json tool's faked + # stream, so it still needs only bedrock:InvokeModel. + streams_natively: Final = native_schema is not None and bedrock_chat_opts_into_native_structured_output( + model, thinking_requested, self.custom_llm_provider or "bedrock_converse", model_info + ) + if json_schema is not None and not streams_natively and non_default_params.get("stream", False) is True: + optional_params["fake_stream"] = True optional_params["json_mode"] = True return optional_params @@ -1413,6 +1441,7 @@ class AmazonConverseConfig(BaseConfig): block_type: Literal["system"], model: str | None = None, deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> SystemContentBlock | None: pass @@ -1426,6 +1455,7 @@ class AmazonConverseConfig(BaseConfig): block_type: Literal["content_block"], model: str | None = None, deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> ContentBlock | None: pass @@ -1438,9 +1468,10 @@ class AmazonConverseConfig(BaseConfig): block_type: Literal["system", "content_block"], model: str | None = None, deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> SystemContentBlock | ContentBlock | None: cache_control: Final = message_block.get("cache_control", None) - if cache_control is None or not bedrock_model_accepts_cache_points(model, deployment_model): + if cache_control is None or not bedrock_model_accepts_cache_points(model, deployment_model, model_info): return None cache_point: Final = self._build_cache_point_block(cache_control, model, deployment_model) @@ -1585,8 +1616,80 @@ class AmazonConverseConfig(BaseConfig): return ChatCompletionTextObject(type="text", text=text) return ChatCompletionTextObject(type="text", text=text, cache_control=cache_control) + def _messages_with_cache_marks( + self, + messages: Sequence[AllMessageValues], + model: str, + deployment_model: str, + drop_params: bool, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + model_info: Mapping[str, object] | None, + ) -> Sequence[AllMessageValues]: + if not bedrock_model_known_to_accept_cache_points(model, deployment_model, model_info): + return messages + markable: Final = ( + AnthropicCacheControlHook.prompt_cache_breakpoints_dropped_from_tool_results(messages, drop_params, model) + if "nova" in model.lower() + else messages + ) + return AnthropicCacheControlHook.prompt_cache_breakpoints_as_capped_cache_control( + markable, + drop_params, + AnthropicCacheControlHook.count_converse_cache_points, + lambda: self._tool_cache_points_on_the_wire( + model, deployment_model, optional_params, litellm_params, model_info + ), + ) + + def _tool_cache_points_on_the_wire( + self, + model: str, + deployment_model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + model_info: Mapping[str, object] | None, + ) -> int: + tools: Final = optional_params.get("tools") + bedrock_tools, _ = self._process_tools_and_beta( + tools if isinstance(tools, list) else [], model, None, {}, litellm_params, deployment_model, model_info + ) + injects_tool_config_point: Final = bool(bedrock_tools) and any( + point.get("location") == "tool_config" + for point in configured_injection_points(optional_params.get("cache_control_injection_points")) + ) + return sum(1 for tool in bedrock_tools if "cachePoint" in tool) + int(injects_tool_config_point) + + def _guarded_text_keeps_cache_marks( + self, + messages: Sequence[AllMessageValues], + system_content_blocks: Sequence[SystemContentBlock], + model: str, + deployment_model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + model_info: Mapping[str, object] | None, + ) -> bool: + if "guardrailConfig" not in optional_params or not bedrock_model_known_to_accept_cache_points( + model, deployment_model, model_info + ): + return False + message_cache_points: Final = sum(map(AnthropicCacheControlHook.count_converse_cache_points, messages)) + if message_cache_points == 0: + return False + cache_points: Final = ( + sum(1 for block in system_content_blocks if "cachePoint" in block) + + message_cache_points + + self._tool_cache_points_on_the_wire(model, deployment_model, optional_params, litellm_params, model_info) + ) + return cache_points <= MAX_CACHE_CONTROL_BLOCKS + def _transform_system_message( - self, messages: list[AllMessageValues], model: str | None = None, deployment_model: str | None = None + self, + messages: list[AllMessageValues], + model: str | None = None, + deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> tuple[list[AllMessageValues], list[SystemContentBlock]]: leading_count: Final = next( (i for i, m in enumerate(messages) if not is_system_message(m)), @@ -1602,7 +1705,7 @@ class AmazonConverseConfig(BaseConfig): if isinstance(content, str) and content: system_content_blocks.append(SystemContentBlock(text=content)) cache_block = self.get_cache_point_block( - message, block_type="system", model=model, deployment_model=deployment_model + message, block_type="system", model=model, deployment_model=deployment_model, model_info=model_info ) if cache_block: system_content_blocks.append(cache_block) @@ -1611,7 +1714,11 @@ class AmazonConverseConfig(BaseConfig): if m.get("type") == "text" and m.get("text"): system_content_blocks.append(SystemContentBlock(text=m["text"])) cache_block = self.get_cache_point_block( - m, block_type="system", model=model, deployment_model=deployment_model + m, + block_type="system", + model=model, + deployment_model=deployment_model, + model_info=model_info, ) if cache_block: system_content_blocks.append(cache_block) @@ -1809,6 +1916,7 @@ class AmazonConverseConfig(BaseConfig): additional_request_params: dict, litellm_params: Mapping[str, object] | None = None, deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> tuple[list[ToolBlock], list]: """Process tools and collect anthropic_beta values.""" bedrock_tools: list[ToolBlock] = [] @@ -1870,6 +1978,7 @@ class AmazonConverseConfig(BaseConfig): _tools_the_model_accepts(regular_tools, model, litellm_params), model=model, deployment_model=deployment_model, + model_info=model_info, ) # Add computer use tools and anthropic_beta if needed (only when computer use tools are present) @@ -1939,6 +2048,7 @@ class AmazonConverseConfig(BaseConfig): _tools_the_model_accepts(_TOOLS_AS_SENT.validate_python(filtered_tools), model, litellm_params), model=model, deployment_model=deployment_model, + model_info=model_info, ) # Append pre-formatted tools (systemTool etc.) after transformation @@ -2014,6 +2124,7 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool = False, litellm_params: Mapping[str, object] | None = None, deployment_model: str | None = None, + model_info: Mapping[str, object] | None = None, ) -> CommonRequestObject: # Drop thinking param if thinking is enabled but thinking_blocks are missing # This prevents the error: "Expected thinking or redacted_thinking, but found tool_use" @@ -2052,7 +2163,7 @@ class AmazonConverseConfig(BaseConfig): # Process tools and collect beta values bedrock_tools, anthropic_beta_list = self._process_tools_and_beta( - original_tools, model, headers, additional_request_params, litellm_params, deployment_model + original_tools, model, headers, additional_request_params, litellm_params, deployment_model, model_info ) requested_thinking: Final = additional_request_params.get("thinking") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # _prepare_request_params hands back a bare dict body_betas: Final = additional_request_params.get("anthropic_beta") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # same bare dict @@ -2070,7 +2181,7 @@ class AmazonConverseConfig(BaseConfig): if ( cache_injection_points and len(bedrock_tools) > 0 - and bedrock_model_accepts_cache_points(model, deployment_model) + and bedrock_model_accepts_cache_points(model, deployment_model, model_info) ): for point in cache_injection_points: if point.get("location") == "tool_config": @@ -2130,22 +2241,27 @@ class AmazonConverseConfig(BaseConfig): return data - @staticmethod def _claude_messages_within_cache_control_cap( + self, messages: Sequence[AllMessageValues], capability_model: str, + deployment_model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object], + model_info: Mapping[str, object] | None, ) -> Sequence[AllMessageValues]: if not drops_unsupported_params(litellm_params.get("drop_params")) or "claude" not in capability_model.lower(): return messages from litellm.integrations.anthropic_cache_control_hook import ( - messages_within_cache_control_cap, + messages_within_reserved_cache_control_cap, tool_message_with_single_cache_mark, ) - return messages_within_cache_control_cap( - tuple(map(tool_message_with_single_cache_mark, messages)), optional_params.get("tools") + return messages_within_reserved_cache_control_cap( + tuple(map(tool_message_with_single_cache_mark, messages)), + self._tool_cache_points_on_the_wire( + capability_model, deployment_model, optional_params, litellm_params, model_info + ), ) async def _async_transform_request( @@ -2156,27 +2272,53 @@ class AmazonConverseConfig(BaseConfig): litellm_params: dict[str, object], headers: dict | None = None, ) -> RequestObject: + from litellm.integrations.anthropic_cache_control_hook import routed_model_info + base_model: Final = litellm_params.get("base_model") capability_model: Final = bedrock_capability_model(model, base_model) + model_info: Final = routed_model_info(litellm_params.get("model_info")) + json_optional_params: Final = _JSON_OBJECT.validate_python(optional_params) + marked_messages: Final = self._messages_with_cache_marks( + messages, + capability_model, + model, + drops_unsupported_params(litellm_params.get("drop_params")), + json_optional_params, + litellm_params, + model_info, + ) messages, system_content_blocks = self._transform_system_message( list( self._claude_messages_within_cache_control_cap( - messages, capability_model, _JSON_OBJECT.validate_python(optional_params), litellm_params + marked_messages, capability_model, model, json_optional_params, litellm_params, model_info ) ), model=capability_model, deployment_model=model, + model_info=model_info, ) # Convert last user message to guarded_text if guardrailConfig is present + guard_input: Final = self._handle_orphaned_tool_blocks(messages, optional_params) messages = self._convert_consecutive_user_messages_to_guarded_text( - self._handle_orphaned_tool_blocks(messages, optional_params), optional_params + guard_input, + optional_params, + keeps_cache_marks=self._guarded_text_keeps_cache_marks( + guard_input, + system_content_blocks, + capability_model, + model, + json_optional_params, + litellm_params, + model_info, + ), ) ## TRANSFORMATION ## _data: Final[CommonRequestObject] = self._transform_request_helper( model=capability_model, deployment_model=model, + model_info=model_info, system_content_blocks=system_content_blocks, optional_params=optional_params, messages=messages, @@ -2191,6 +2333,7 @@ class AmazonConverseConfig(BaseConfig): messages=messages, model=capability_model, deployment_model=model, + model_info=model_info, llm_provider="bedrock_converse", user_continue_message=litellm_params.pop("user_continue_message", None), base_model=self._deployment_base_model(_DEPLOYMENT_PARAMS.validate_python(litellm_params)), @@ -2243,26 +2386,52 @@ class AmazonConverseConfig(BaseConfig): litellm_params: dict[str, object], headers: dict | None = None, ) -> RequestObject: + from litellm.integrations.anthropic_cache_control_hook import routed_model_info + base_model: Final = litellm_params.get("base_model") capability_model: Final = bedrock_capability_model(model, base_model) + model_info: Final = routed_model_info(litellm_params.get("model_info")) + json_optional_params: Final = _JSON_OBJECT.validate_python(optional_params) + marked_messages: Final = self._messages_with_cache_marks( + messages, + capability_model, + model, + drops_unsupported_params(litellm_params.get("drop_params")), + json_optional_params, + litellm_params, + model_info, + ) messages, system_content_blocks = self._transform_system_message( list( self._claude_messages_within_cache_control_cap( - messages, capability_model, _JSON_OBJECT.validate_python(optional_params), litellm_params + marked_messages, capability_model, model, json_optional_params, litellm_params, model_info ) ), model=capability_model, deployment_model=model, + model_info=model_info, ) # Convert last user message to guarded_text if guardrailConfig is present + guard_input: Final = self._handle_orphaned_tool_blocks(messages, optional_params) messages = self._convert_consecutive_user_messages_to_guarded_text( - self._handle_orphaned_tool_blocks(messages, optional_params), optional_params + guard_input, + optional_params, + keeps_cache_marks=self._guarded_text_keeps_cache_marks( + guard_input, + system_content_blocks, + capability_model, + model, + json_optional_params, + litellm_params, + model_info, + ), ) _data: Final[CommonRequestObject] = self._transform_request_helper( model=capability_model, deployment_model=model, + model_info=model_info, system_content_blocks=system_content_blocks, optional_params=optional_params, messages=messages, @@ -2278,6 +2447,7 @@ class AmazonConverseConfig(BaseConfig): messages=messages, model=capability_model, deployment_model=model, + model_info=model_info, llm_provider="bedrock_converse", user_continue_message=litellm_params.pop("user_continue_message", None), base_model=self._deployment_base_model(_DEPLOYMENT_PARAMS.validate_python(litellm_params)), diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py index e32df0bbdbb..37acb5423ff 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py @@ -1,18 +1,22 @@ -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import TypeAdapter +import litellm from litellm.anthropic_beta_gated_fields import strip_beta_gated_request_body from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( + RemoteMedia, async_inline_remote_media, convert_url_to_base64, + inline_every_remote_url, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import ( @@ -25,7 +29,9 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation ) from litellm.llms.bedrock.common_utils import ( BEDROCK_INVOKE_TRACE_FIELD, + BedrockModelInfo, apply_bedrock_invoke_structured_output, + bedrock_capability_model, bedrock_chat_uses_native_structured_output, bedrock_cleaned_output_schema, bedrock_supports_tool_search, @@ -35,6 +41,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, requests_opt_in_native_output, + strip_bedrock_routing_prefix, strip_unsupported_bedrock_invoke_output_config_keys, tools_without_eager_input_streaming, ) @@ -59,6 +66,19 @@ _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) _LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) +def _anthropic_config_max_tokens() -> int | None: + return litellm.AnthropicConfig.max_tokens + + +def _mapped_max_tokens(model: str) -> int | None: + row: Final = object_mapping(litellm.model_cost.get(model)) or {} # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # cost map rows are untyped JSON + match row.get("max_output_tokens"): + case int() as max_tokens: + return max_tokens + case _: + return None + + class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): """ Reference: @@ -79,9 +99,18 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_version: str = "bedrock-2023-05-31" - def __init__(self, model_info: Mapping[str, object] | None = None) -> None: + def __init__( + self, + model_info: Mapping[str, object] | None = None, + read_configured_max_tokens: Callable[[], int | None] = _anthropic_config_max_tokens, + ) -> None: super().__init__() self._deployment_model_info: Final = model_info + self._read_configured_max_tokens: Final = read_configured_max_tokens + + @classmethod + def get_config(cls, *, model: str | None = None) -> dict[str, object]: # mutable-ok: overrides BaseConfig's dict + return {key: value for key, value in super().get_config(model=model).items() if key != "max_tokens"} @property def custom_llm_provider(self) -> str | None: @@ -90,6 +119,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): def should_strip_billing_metadata(self) -> bool: return True + def inlines_remote_media(self, media: RemoteMedia) -> bool: + return inline_every_remote_url(media) + def get_supported_openai_params(self, model: str) -> list[str]: return AnthropicConfig.get_supported_openai_params(self, model) @@ -102,10 +134,15 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): json_schema: Mapping[str, object] | None, request_params: Mapping[str, object], ) -> bool: - return bedrock_chat_uses_native_structured_output( - model, self.is_thinking_requested(request_params), model_info=self._deployment_model_info + return self.chat_uses_native_structured_output( + model, self.is_thinking_requested(request_params), self._deployment_model_info ) and (json_schema is None or bedrock_cleaned_output_schema(model, json_schema) is not None) + def chat_uses_native_structured_output( + self, model: str, thinking_requested: bool, model_info: Mapping[str, object] | None + ) -> bool: + return bedrock_chat_uses_native_structured_output(model, thinking_requested, model_info=model_info) + def uses_opt_in_native_output( self, model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object] ) -> bool: @@ -241,6 +278,36 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): headers=headers, ) + def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: + return ( + BedrockModelInfo.get_base_model( + bedrock_capability_model(strip_bedrock_routing_prefix(model), litellm_params.get("base_model")) + ), + ) + + def _default_max_tokens(self, model: str, litellm_params: Mapping[str, object]) -> int: + configured_max_tokens: Final = self._read_configured_max_tokens() + match deployment_model_info(litellm_params).get("max_output_tokens"): + case int() as deployment_max_tokens: + return ( + deployment_max_tokens + if configured_max_tokens is None + else min(deployment_max_tokens, configured_max_tokens) + ) + case _ if configured_max_tokens is not None: + return configured_max_tokens + case _: + return next( + ( + max_tokens + for max_tokens in map( + _mapped_max_tokens, self._default_max_tokens_models(model, litellm_params) + ) + if max_tokens is not None + ), + DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + ) + def _build_bedrock_anthropic_request_base( self, model: str, @@ -257,23 +324,35 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model=model, output_config=filtered_params["output_config"], ) - filtered_params = self._normalize_bedrock_tool_search_tools(filtered_params) + normalized_params: Final = self._normalize_bedrock_tool_search_tools(filtered_params) + default_max_tokens: Final = ( + None + if "max_tokens" in normalized_params + else self._default_max_tokens(model, _LITELLM_PARAMS.validate_python(litellm_params)) + ) + request_params: Final = ( + normalized_params if default_max_tokens is None else {**normalized_params, "max_tokens": default_max_tokens} + ) anthropic_request: Final = AnthropicConfig.transform_request( self, model=model, messages=messages, - optional_params=filtered_params, + optional_params=request_params, litellm_params=litellm_params, headers=headers, ) anthropic_request.pop("model", None) anthropic_request.pop("stream", None) + model_info: Final = deployment_model_info(_LITELLM_PARAMS.validate_python(litellm_params)) apply_bedrock_invoke_structured_output( model=model, request_body=anthropic_request, - model_info=deployment_model_info(_LITELLM_PARAMS.validate_python(litellm_params)), + model_info=model_info, + chat_uses_native_output=lambda: self.chat_uses_native_structured_output( + model, self.is_thinking_requested(_NON_DEFAULT_PARAMS.validate_python(filtered_params)), model_info + ), ) strip_unsupported_bedrock_invoke_output_config_keys( model=model, @@ -377,7 +456,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict: + def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict[str, object]: """ Convert tool search entries to the format supported by the Bedrock Invoke API. """ diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index b5a352bab48..952ca7a2cd4 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -15,7 +15,7 @@ from litellm.llms.anthropic.chat.transformation import JsonModeRequest from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) -from litellm.llms.bedrock.common_utils import build_mantle_messages_url +from litellm.llms.bedrock.common_utils import bedrock_chat_opts_into_native_structured_output, build_mantle_messages_url from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -54,6 +54,14 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): region=region, ) + def chat_uses_native_structured_output( + self, model: str, thinking_requested: bool, model_info: Mapping[str, object] | None + ) -> bool: + return bedrock_chat_opts_into_native_structured_output(model, thinking_requested, model_info=model_info) + + def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: + return (f"bedrock_mantle/{model}", *super()._default_max_tokens_models(model, litellm_params)) + def validate_environment( self, headers: dict, diff --git a/litellm/llms/bedrock/claude_platform/transformation.py b/litellm/llms/bedrock/claude_platform/transformation.py index a5ecc37d394..6efc4b506e8 100644 --- a/litellm/llms/bedrock/claude_platform/transformation.py +++ b/litellm/llms/bedrock/claude_platform/transformation.py @@ -30,6 +30,9 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig): def should_strip_billing_metadata(self) -> bool: return True + def runs_hosted_mcp(self) -> bool: + return True + def validate_environment( self, headers: dict, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index d6f5d5b6bd7..c1b3b5616b8 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -10,7 +10,7 @@ import json import os import re import warnings -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias @@ -33,9 +33,11 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.litellm_core_utils.core_helpers import lax_bool from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.common_utils import schema_property_count from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -43,6 +45,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned from litellm.secret_managers.main import get_secret, get_secret_str +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams if TYPE_CHECKING: @@ -50,6 +53,12 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues BEDROCK_INVOKE_TRACE_FIELD: Final = "amazon-bedrock-trace" +BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS: Final[Mapping[str, str]] = MappingProxyType( + { + "compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, + "clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + } +) _ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs" @@ -308,7 +317,33 @@ def bedrock_supports_native_structured_output( ) or _bedrock_native_output_opt_in(model, custom_llm_provider, model_info) -def bedrock_chat_uses_native_structured_output( +def _bedrock_map_native_output_flag(model: str) -> object: + routed: Final = strip_bedrock_routing_prefix(model) + flags: Final = ( + (entry or {}).get("supports_bedrock_native_structured_output") + for entry in ( + _model_cost_entry(f"bedrock/{routed}"), + _model_cost_entry(routed), + _model_cost_entry(split_bedrock_region_path(model)[1]), + ) + ) + return next((flag for flag in flags if isinstance(flag, bool)), None) + + +def _bedrock_native_output_opted_out(model: str, model_info: Mapping[str, object] | None) -> bool: + deployment_flag: Final = (model_info or {}).get("supports_bedrock_native_structured_output") + if isinstance(deployment_flag, bool): + return not deployment_flag + return _bedrock_map_native_output_flag(model) is False + + +def _bedrock_native_output_replaces_unforceable_tool(model: str, model_info: Mapping[str, object] | None) -> bool: + return AnthropicModelInfo.forced_tool_use_unsupported(model) and not _bedrock_native_output_opted_out( + model, model_info + ) + + +def bedrock_chat_opts_into_native_structured_output( model: str, thinking_requested: bool, custom_llm_provider: str = "bedrock", @@ -319,7 +354,18 @@ def bedrock_chat_uses_native_structured_output( ) -def bedrock_native_output_is_opt_in( +def bedrock_chat_uses_native_structured_output( + model: str, + thinking_requested: bool, + custom_llm_provider: str = "bedrock", + model_info: Mapping[str, object] | None = None, +) -> bool: + return bedrock_chat_opts_into_native_structured_output( + model, thinking_requested, custom_llm_provider, model_info + ) or (not thinking_requested and _bedrock_native_output_replaces_unforceable_tool(model, model_info)) + + +def bedrock_native_output_is_deployment_opt_in( model: str, custom_llm_provider: str = "bedrock", model_info: Mapping[str, object] | None = None ) -> bool: return not _bedrock_model_supports( @@ -327,6 +373,15 @@ def bedrock_native_output_is_opt_in( ) and _bedrock_native_output_opt_in(model, custom_llm_provider, model_info) +def bedrock_native_output_is_opt_in( + model: str, custom_llm_provider: str = "bedrock", model_info: Mapping[str, object] | None = None +) -> bool: + return not _bedrock_model_supports(model, "supports_native_structured_output", custom_llm_provider) and ( + _bedrock_native_output_opt_in(model, custom_llm_provider, model_info) + or _bedrock_native_output_replaces_unforceable_tool(model, model_info) + ) + + def without_bedrock_native_output_opt_in(litellm_params: Mapping[str, object]) -> Mapping[str, object]: return { **litellm_params, @@ -417,13 +472,18 @@ def requests_output_config_format(optional_params: Mapping[str, object]) -> bool return False +def requests_native_output(anthropic_optional_params: Mapping[str, object]) -> bool: + return isinstance(anthropic_optional_params.get("output_format"), Mapping) or requests_output_config_format( + anthropic_optional_params + ) + + def requests_opt_in_native_output( model: str, anthropic_optional_params: Mapping[str, object], model_info: Mapping[str, object] | None ) -> bool: - native_output_requested: Final = isinstance( - anthropic_optional_params.get("output_format"), Mapping - ) or requests_output_config_format(anthropic_optional_params) - return native_output_requested and bedrock_native_output_is_opt_in(model, model_info=model_info) + return requests_native_output(anthropic_optional_params) and bedrock_native_output_is_opt_in( + model, model_info=model_info + ) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) @@ -484,6 +544,7 @@ def apply_bedrock_invoke_structured_output( model: str, request_body: dict[str, object], # mutable-ok: edited in place like siblings model_info: Mapping[str, object] | None = None, + chat_uses_native_output: Callable[[], bool] = lambda: False, ) -> None: """ Route Anthropic structured-output params to what the Bedrock model supports. @@ -495,7 +556,8 @@ def apply_bedrock_invoke_structured_output( forwarded as ``output_config.format``, which Bedrock relays to the model for enforced structured output. For every other model the schema is inlined into the last user message as best-effort text, with a warning because nothing - enforces it. + enforces it. ``chat_uses_native_output`` is asked only when a schema is + present, and a true answer sends it natively after the same cleanup. """ legacy_output_format: Final = request_body.pop("output_format", None) output_config_format: Final = pop_bedrock_invoke_output_config_format(request_body) @@ -504,7 +566,15 @@ def apply_bedrock_invoke_structured_output( if schema_format is None: return - native_format: Final = bedrock_native_output_format(model, schema_format, model_info=model_info) + native_format: Final = bedrock_native_output_format( + model, + schema_format, + model_info=( + {**(model_info or {}), "supports_bedrock_native_structured_output": True} + if chat_uses_native_output() + else model_info + ), + ) if native_format is not None: existing_output_config: Final = request_body.get("output_config") if isinstance(existing_output_config, dict): @@ -1446,7 +1516,9 @@ def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: ) -def bedrock_model_accepts_cache_points(model: str | None, deployment_model: str | None = None) -> bool: +def bedrock_model_accepts_cache_points( + model: str | None, deployment_model: str | None = None, model_info: Mapping[str, object] | None = None +) -> bool: """ Whether Converse ``cachePoint`` blocks may be sent to this model. @@ -1466,43 +1538,102 @@ def bedrock_model_accepts_cache_points(model: str | None, deployment_model: str never silently degrade. ``litellm.utils.supports_prompt_caching`` is not reusable here: it returns False for unmapped models, the opposite polarity. ``deployment_model`` is the model as configured when ``model`` is its ``base_model``: its own declaration wins. + A routed call passes its deployment's ``model_info``, whose breakpoint flag wins outright and + whose ``supports_prompt_caching`` wins over the catalog's for a known model. Other flags come + from the catalog as loaded, since the registered keys are shared by every deployment of the + same model string and carry whichever sibling registered last. """ if model is None: return True deployment_declared: Final = ( - None if deployment_model is None else bedrock_declared_cache_point_support(deployment_model) + None if deployment_model is None else bedrock_declared_cache_point_support(deployment_model, model_info) ) declared: Final = ( - deployment_declared if deployment_declared is not None else bedrock_declared_cache_point_support(model) + deployment_declared + if deployment_declared is not None + else bedrock_declared_cache_point_support(model, model_info) ) return declared is not False -def bedrock_declared_cache_point_support(model: str) -> bool | None: +def bedrock_declared_cache_point_support(model: str, model_info: Mapping[str, object] | None = None) -> bool | None: if _OPENAI_FAMILY_MODEL_RE.search(model): return False - map_keys: Final = (model, get_bedrock_base_model(model)) + explicit_marker_support: Final = _explicit_cache_point_support(model, model_info) + if explicit_marker_support is not None: + return explicit_marker_support + return _mapped_prompt_caching_support(model, model_info) + + +def bedrock_model_known_to_accept_cache_points( + model: str, deployment_model: str | None = None, model_info: Mapping[str, object] | None = None +) -> bool: + """Like ``bedrock_model_accepts_cache_points``, but only for the Claude and Nova families Bedrock documents + as taking cache checkpoints, unless the entry sets ``supports_prompt_cache_breakpoint``: an OpenAI + ``prompt_cache_breakpoint`` reaches every model of a mixed fleet, not only the ones a client marked for + Bedrock caching, and other families price implicit cache reads yet may reject the cachePoint field.""" + deployment_declared: Final = ( + None if deployment_model is None else _declared_marker_support(deployment_model, model_info) + ) + declared: Final = ( + deployment_declared if deployment_declared is not None else _declared_marker_support(model, model_info) + ) + return declared is True + + +def _declared_marker_support(model: str, model_info: Mapping[str, object] | None) -> bool | None: + if _OPENAI_FAMILY_MODEL_RE.search(model): + return False + explicit_marker_support: Final = _explicit_cache_point_support(model, model_info) + if explicit_marker_support is not None: + return explicit_marker_support + mapped: Final = _mapped_prompt_caching_support(model, model_info) + if not _documents_cache_checkpoints(model): + return None if mapped is None else False + return mapped is not False + + +def _documents_cache_checkpoints(model: str) -> bool: + return _is_claude(model) or get_bedrock_base_model(model).startswith("amazon.nova") + + +def _explicit_cache_point_support(model: str, model_info: Mapping[str, object] | None) -> bool | None: + deployment_flag: Final = None if model_info is None else model_info.get("supports_prompt_cache_breakpoint") + if deployment_flag is not None: + return deployment_flag is True + catalog_entry: Final = _catalog_entry_reader(model_info) registered_keys: Final = tuple(f"bedrock/{route}{model}" for route in ("", "converse/", "converse_like/")) - explicit_marker_support: Final = next( + return next( ( entry.get("supports_prompt_cache_breakpoint") is True - for key in (*registered_keys, *map_keys) - if (entry := litellm.model_cost.get(key)) is not None - and entry.get("supports_prompt_cache_breakpoint") is not None + for key in (*registered_keys, model, get_bedrock_base_model(model)) + if (entry := catalog_entry(key)) is not None and entry.get("supports_prompt_cache_breakpoint") is not None ), None, ) - if explicit_marker_support is not None: - return explicit_marker_support - if not any(key in litellm.model_cost for key in map_keys): + + +def _mapped_prompt_caching_support(model: str, model_info: Mapping[str, object] | None) -> bool | None: + catalog_entry: Final = _catalog_entry_reader(model_info) + map_keys: Final = (model, get_bedrock_base_model(model)) + if not any(catalog_entry(key) is not None for key in map_keys): return None + deployment_caching: Final = None if model_info is None else lax_bool(model_info.get("supports_prompt_caching")) + if deployment_caching is not None: + return deployment_caching return any( - entry.get("supports_prompt_caching") is True - for key in map_keys - if (entry := litellm.model_cost.get(key)) is not None + entry.get("supports_prompt_caching") is True for key in map_keys if (entry := catalog_entry(key)) is not None ) +def _catalog_entry_reader(model_info: Mapping[str, object] | None) -> Callable[[str], Mapping[str, object] | None]: + from litellm.utils import catalog_with_runtime_registrations + + if model_info is not None: + return catalog_with_runtime_registrations().get + return _model_cost_entry + + def bedrock_supports_tool_search(model: str) -> bool: """ Whether Bedrock InvokeModel admits the ``tool_search_tool_*`` tool types on ``model``. diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 26bc5904e06..dec099f37c0 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -13,9 +13,14 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.common_utils import context_management_betas, context_management_with_edits from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model +from litellm.llms.bedrock.common_utils import ( + BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS, + BedrockError, + get_bedrock_base_model, +) from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model @@ -74,6 +79,16 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: ) +def _mantle_retry_outcome( + runtime: CountTokensOutcome, mantle: CountTokensOutcome, carried_edits: bool +) -> CountTokensOutcome: + match runtime, mantle: + case CountTokensFailure(status_code=400), CountTokensFailure(status_code=404) if carried_edits: + return runtime + case _: + return mantle + + def _request_data( model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object ) -> Mapping[str, object]: @@ -85,6 +100,33 @@ def _request_data( } +def _invoke_model_message(message: Mapping[str, object]) -> Mapping[str, object]: + content: Final = message.get("content") + return ( + {**message, "content": [{"type": "text", "text": content}]} if isinstance(content, str) and content else message + ) + + +def _invoke_supports_edit(edit: JsonValue) -> bool: + edit_type: Final = edit.get("type") if isinstance(edit, dict) else None + return isinstance(edit_type, str) and edit_type in BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS + + +def _invoke_context_management(context_management: CountContextManagement | None) -> Mapping[str, JsonValue]: + spec: Final = context_management_with_edits(None if context_management is None else context_management.spec) + edits: Final = spec.get("edits") if isinstance(spec, dict) else None + supported_edits: Final = tuple(filter(_invoke_supports_edit, edits if isinstance(edits, list) else ())) + if not isinstance(spec, dict) or not supported_edits: + return MappingProxyType({}) + supported_spec: Final[JsonValue] = {**spec, "edits": list(supported_edits)} + return MappingProxyType( + { + "context_management": supported_spec, + "anthropic_beta": list(context_management_betas(supported_spec)), + } + ) + + class BedrockTokenCounter(BaseTokenCounter): """Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model""" @@ -167,11 +209,19 @@ class BedrockTokenCounter(BaseTokenCounter): ) if isinstance(translated, ChatCountTranslationError): return translated.response(request_model, model_to_use, primary_tokenizer) - runtime_request: Final = ( - _request_data(model_to_use, messages, tools, system) - if translated is None - else _request_data(model_to_use, translated.messages, translated.tools, translated.system) + native_context_management: Final = _invoke_context_management( + context_management if is_mantle_claude_model(resolved_model) else None ) + counted_messages: Final = messages if translated is None else translated.messages + runtime_request: Final = { + **_request_data( + model_to_use, + tuple(map(_invoke_model_message, counted_messages)) if native_context_management else counted_messages, + tools if translated is None else translated.tools, + system if translated is None else translated.system, + ), + **native_context_management, + } mantle_request: Final = runtime_request if translated is None else {**runtime_request, **translated.options} primary: Final = ( @@ -184,8 +234,12 @@ class BedrockTokenCounter(BaseTokenCounter): ) ) outcome: Final = ( - await self._count_with( - self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model + _mantle_retry_outcome( + primary, + await self._count_with( + self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model + ), + carried_edits=bool(native_context_management), ) if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model) else primary diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 5cd26235ba2..5e5c5516b94 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -8,7 +8,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.anthropic.common_utils import tools_without_custom_field +from litellm.llms.anthropic.common_utils import context_management_betas, tools_without_custom_field from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url @@ -22,6 +22,7 @@ class MantleCountTokensRequest(TypedDict): messages: ReadOnly[list[dict[str, JsonValue]]] system: ReadOnly[NotRequired[JsonValue]] tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]] + context_management: ReadOnly[NotRequired[JsonValue]] _COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest) @@ -74,13 +75,19 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): optional_params=_COUNT_OPTIONS.validate_python( {name: value for name, value in request_data.items() if name in COUNT_TOKEN_OPTION_NAMES} ), + context_management=request.get("context_management"), ) + betas: Final = context_management_betas(request.get("context_management")) verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url) api_key: Final = litellm_params.get("api_key") signed_headers, signed_body = await run_aws_signing( self._sign_request, service_name="bedrock", - headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION}, + headers={ + "Content-Type": "application/json", + "anthropic-version": MANTLE_ANTHROPIC_VERSION, + **({"anthropic-beta": ",".join(betas)} if betas else {}), + }, optional_params=litellm_params, request_data=body, api_base=endpoint_url, diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py index e5d6a168d57..8b4e2e5b5c5 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py @@ -37,8 +37,10 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation bedrock_invoke_guardrail_headers, ) from litellm.llms.bedrock.common_utils import ( + BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS, BedrockError, apply_bedrock_invoke_structured_output, + bedrock_native_output_is_deployment_opt_in, bedrock_supports_tool_search, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, @@ -46,7 +48,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, - requests_opt_in_native_output, + requests_native_output, requests_output_config_format, strip_unsupported_bedrock_invoke_output_config_keys, supports_bedrock_prompt_cache_ttl, @@ -130,7 +132,10 @@ class AmazonAnthropicClaudeMessagesConfig( litellm_params: GenericLiteLLMParams, ) -> GenericLiteLLMParams | None: model_info: Final = _MODEL_INFO.validate_python(litellm_params.model_info) - if not requests_opt_in_native_output(model, anthropic_messages_optional_request_params, model_info): + if not ( + requests_native_output(anthropic_messages_optional_request_params) + and bedrock_native_output_is_deployment_opt_in(model, model_info=model_info) + ): return None return litellm_params.model_copy( update={"model_info": {**(model_info or {}), "supports_bedrock_native_structured_output": False}} @@ -499,12 +504,7 @@ class AmazonAnthropicClaudeMessagesConfig( # Bedrock InvokeModel DOES support ``clear_tool_uses_20250919`` under the # ``context-management-2025-06-27`` beta. AWS docs: # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-tool-use.md - _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType( - { - "compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, - "clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, - } - ) + _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS @classmethod def _filter_context_management_for_bedrock_invoke( diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 00622a1e287..49939399ec8 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -9,7 +9,7 @@ import contextlib import importlib.metadata import json from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from types import MappingProxyType from typing import Final, NoReturn, Protocol, runtime_checkable @@ -34,12 +34,15 @@ from litellm.llms.base_llm.realtime.transformation import RealtimeSessionUpdateR from litellm.types.llms.bedrock import AwsAuthParams from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput +from litellm.types.utils import StandardLoggingPayloadErrorInformation from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig, requested_modalities _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_REQUEST_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) +_SPEND_LOG_METADATA_KEYS: Final = ("litellm_metadata", "metadata") _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -61,6 +64,10 @@ _BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType( ) +def _stream_ended_before_session_closed() -> BedrockError: + return BedrockError(status_code=503, message="Bedrock realtime stream ended before the session closed") + + def _as_bedrock_error(error: BaseException) -> BaseException: status_code: Final = _BEDROCK_STREAM_ERROR_STATUS.get(type(error).__name__) if status_code is None: @@ -165,6 +172,12 @@ class _BridgeOutcome: client_disconnected: bool +@dataclass(frozen=True, slots=True) +class SessionEndSignals: + input_closing: asyncio.Event = field(default_factory=asyncio.Event) + output_failed: asyncio.Event = field(default_factory=asyncio.Event) + + async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: str | None) -> AsyncIterator[str]: if initial_message is not None: yield initial_message @@ -188,8 +201,8 @@ def raise_if_session_committed(scope: Mapping[str, object] | None) -> None: raise BedrockError( status_code=400, message=( - "Bedrock realtime session already committed to a provider stream; it cannot be replayed" - + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") + (f"{committed_failure}. " if committed_failure else "") + + "Bedrock realtime session already committed to a provider stream; it cannot be replayed" ), ) @@ -201,11 +214,48 @@ def _pending_opening_message(scope: Mapping[str, object]) -> str | None: return pending if isinstance(pending, str) else None +def _spend_log_metadata(litellm_params: Mapping[str, object]) -> tuple[str, Mapping[str, object]] | None: + return next( + ( + (key, metadata) + for key in _SPEND_LOG_METADATA_KEYS + if (metadata := _REQUEST_MAPPING_ADAPTER.validate_python(litellm_params.get(key))) + ), + None, + ) + + +def _record_failure_on_spend_log(logging_obj: LiteLLMLogging, error: BaseException) -> None: + """The spend row of a session the provider ended reads failure while it still bills the forwarded events.""" + litellm_params: Final = _REQUEST_MAPPING_ADAPTER.validate_python( + logging_obj.model_call_details.get("litellm_params") + ) + if litellm_params is None: + return + spend_log_metadata: Final = _spend_log_metadata(litellm_params) + if spend_log_metadata is None: + return + status_code: Final[object] = getattr(error, "status_code", None) + error_information: Final[StandardLoggingPayloadErrorInformation] = { + "error_code": "" if status_code is None else str(status_code), + "error_class": type(error).__name__, + "error_message": redact_string(str(error)), + "llm_provider": "bedrock", + } + metadata_key, metadata = spend_log_metadata + logging_obj.model_call_details["litellm_params"] = { # rebind-ok: the spend row reads the logging object + **litellm_params, + metadata_key: {**metadata, "status": "failure", "error_information": error_information}, + } + + def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn: error: Final = _as_bedrock_error(failure) verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", redact_string(str(error))) if _is_session_committed(scope): scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = redact_string(str(error)) + if error is failure: + raise failure raise error from failure @@ -267,6 +317,25 @@ async def _reject_session_update( return True +async def _close_bedrock_input( + bedrock_stream: BedrockBidirectionalStream, + close_messages: Sequence[str], + send_to_bedrock: Callable[[str], Awaitable[None]], + timeout: float, +) -> None: + async def flush_and_close() -> None: + for close_message in close_messages: + with contextlib.suppress(Exception): + await send_to_bedrock(close_message) + with contextlib.suppress(Exception): + await bedrock_stream.input_stream.close() + + try: + await asyncio.wait_for(flush_and_close(), timeout) + except (asyncio.TimeoutError, TimeoutError): + verbose_proxy_logger.debug("Bedrock Realtime: timed out closing Bedrock input stream after %.1fs", timeout) + + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" @@ -278,6 +347,18 @@ class BedrockRealtime(BaseAWSLLM): super().__init__() self._sdk_version_lookup: Final = sdk_version_lookup self._close_timeout: Final = close_timeout + self._failed_input_closes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs to running closes + + def _close_failed_input_in_background( + self, + bedrock_stream: BedrockBidirectionalStream, + send_to_bedrock: Callable[[str], Awaitable[None]], + ) -> None: + close_task: Final = asyncio.create_task( + _close_bedrock_input(bedrock_stream, (), send_to_bedrock, self._close_timeout) + ) + self._failed_input_closes.add(close_task) + close_task.add_done_callback(self._failed_input_closes.discard) async def async_realtime( self, @@ -433,23 +514,26 @@ class BedrockRealtime(BaseAWSLLM): if _should_log_event(leftover_event) ), ) + if outcome.provider_failure is not None and outcome.client_disconnected: + verbose_proxy_logger.debug( + "Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure + ) + provider_failure: Final = None if outcome.client_disconnected else outcome.provider_failure if logged_events: + if provider_failure is not None: + _record_failure_on_spend_log(logging_obj, _as_bedrock_error(provider_failure)) + usage_logging_obj: Final = logging_obj if provider_failure is None else logging_obj.attempt_scoped_copy() GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - logging_obj.dispatch_success_handlers( + usage_logging_obj.dispatch_success_handlers( list(logged_events), prefer_async_handlers=True, ) ) logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True - if outcome.provider_failure is None: + if provider_failure is None: return - if outcome.client_disconnected: - verbose_proxy_logger.debug( - "Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure - ) - return - _raise_provider_failure(websocket.scope, outcome.provider_failure) + _raise_provider_failure(websocket.scope, provider_failure) async def _bridge( self, @@ -463,16 +547,34 @@ class BedrockRealtime(BaseAWSLLM): ) -> _BridgeOutcome: """Run both forwarding directions until the client leaves or either side fails.""" logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend + signals: Final = SessionEndSignals() async def collect_logged_events() -> None: - async for event in self._forward_bedrock_to_client( - bedrock_stream, websocket, transformation_config, model, logging_obj, session_state - ): - logged.append(event) + try: + async for event in self._forward_bedrock_to_client( + bedrock_stream, + websocket, + transformation_config, + model, + logging_obj, + session_state, + signals=signals, + ): + logged.append(event) + except Exception: + signals.output_failed.set() + raise client_task: Final = asyncio.create_task( self._forward_client_to_bedrock( - websocket, bedrock_stream, transformation_config, model, session_state, logging_obj, initial_message + websocket, + bedrock_stream, + transformation_config, + model, + session_state, + logging_obj, + initial_message, + signals=signals, ) ) bedrock_task: Final = asyncio.create_task(collect_logged_events()) @@ -506,6 +608,8 @@ class BedrockRealtime(BaseAWSLLM): session_state: RealtimeResponseTransformInput, logging_obj: LiteLLMLogging | None = None, initial_message: str | None = None, + *, + signals: SessionEndSignals, ) -> None: """Forward messages from client WebSocket to Bedrock stream. @@ -521,7 +625,10 @@ class BedrockRealtime(BaseAWSLLM): async def send_to_bedrock(bedrock_message: str) -> None: event: Final = build_input_chunk(bedrock_message.encode("utf-8")) - await bedrock_stream.input_stream.send(event) + try: + await bedrock_stream.input_stream.send(event) + except Exception as e: + raise _stream_ended_before_session_closed() from e verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) try: @@ -553,20 +660,13 @@ class BedrockRealtime(BaseAWSLLM): client_ws, transformation_config, model, logging_obj, parsed_client_message ): break + signals.input_closing.set() finally: - - async def close_bedrock_input() -> None: - for close_message in transformation_config.session_close_messages(): - with contextlib.suppress(Exception): - await send_to_bedrock(close_message) - with contextlib.suppress(Exception): - await bedrock_stream.input_stream.close() - - try: - await asyncio.wait_for(close_bedrock_input(), self._close_timeout) - except (asyncio.TimeoutError, TimeoutError): - verbose_proxy_logger.debug( - "Bedrock Realtime: timed out closing Bedrock input stream after %.1fs", self._close_timeout + if signals.output_failed.is_set(): + self._close_failed_input_in_background(bedrock_stream, send_to_bedrock) + else: + await _close_bedrock_input( + bedrock_stream, transformation_config.session_close_messages(), send_to_bedrock, self._close_timeout ) async def _forward_bedrock_to_client( @@ -577,6 +677,8 @@ class BedrockRealtime(BaseAWSLLM): model: str, logging_obj: LiteLLMLogging, session_state: RealtimeResponseTransformInput, + *, + signals: SessionEndSignals, ) -> AsyncIterator[OpenAIRealtimeEvents]: """Forward messages from Bedrock to the client, yielding the ones to record for spend logging. @@ -597,6 +699,8 @@ class BedrockRealtime(BaseAWSLLM): result = await output[1].receive() if result is None: + if not signals.input_closing.is_set(): + raise _stream_ended_before_session_closed() verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended") with contextlib.suppress(Exception): await client_ws.close() diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 05855c307e9..1f7b1712d68 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -15,7 +15,8 @@ role / access key / profile / web identity), signed via the shared BaseAWSLLM._sign_request after the request body is finalized. """ -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence +from functools import partial from typing import Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict import httpx @@ -23,6 +24,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.prompt_templates.common_utils import chat_completion_prompt_cache_breakpoint_for_wire from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items @@ -62,6 +64,10 @@ _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) _BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"}) +def _prompt_cache_breakpoint_as_sent_or_none(marker: object) -> object: + return marker if chat_completion_prompt_cache_breakpoint_for_wire(marker, drop_params=True) == marker else None + + class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig): def __init__( self, @@ -113,6 +119,13 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI headers["OpenAI-Project"] = litellm_params.aws_bedrock_project_id return headers + def _prompt_cache_breakpoint_to_wire( + self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None + ) -> Callable[[object], object]: + if reads_breakpoints or known_support is True: + return partial(chat_completion_prompt_cache_breakpoint_for_wire, drop_params=True) + return _prompt_cache_breakpoint_as_sent_or_none + def supports_native_file_search(self) -> bool: return False diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index cf033e10b41..141a174e558 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx import litellm +from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.cohere import CohereV2ChatResponse @@ -248,6 +249,9 @@ class CohereV2ChatConfig(OpenAIGPTConfig): prompt_tokens: Final = token_usage.get("input_tokens", 0) completion_tokens: Final = token_usage.get("output_tokens", 0) + model_response.choices[0].finish_reason = map_finish_reason( # rebind-ok: ModelResponse is filled in place + cohere_v2_chat_response.get("finish_reason", "COMPLETE") + ) model_response.created = int(time.time()) model_response.model = model usage: Final = Usage( diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 2f60d94c08a..34fde444880 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/litellm/llms/cohere/embed/handler.py @@ -82,7 +82,7 @@ async def async_embedding( ) try: - response: Final = await client.post(api_base, headers=headers, data=json.dumps(data)) + response: Final = await client.post(api_base, headers=headers, data=json.dumps(data), timeout=timeout) except httpx.HTTPStatusError as e: ## LOGGING logging_obj.post_call( @@ -166,7 +166,7 @@ def embedding( if client is None or not isinstance(client, HTTPHandler): client = HTTPHandler(concurrent_limit=1) - response: Final = client.post(embed_url, headers=headers, data=json.dumps(data)) + response: Final = client.post(embed_url, headers=headers, data=json.dumps(data), timeout=timeout) return CohereEmbeddingConfig().transform_response( response=response, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8a094cc6424..2f17dc6dd10 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -37,8 +37,13 @@ import litellm import litellm.litellm_core_utils import litellm.types import litellm.types.utils +from litellm._internal_context import internal_sub_call from litellm._logging import redact_string, verbose_logger -from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, update_headers_with_filtered_beta +from litellm.anthropic_beta_headers_manager import ( + aload_beta_headers_config, + resolve_beta_provider, + update_headers_with_filtered_beta, +) from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs @@ -455,6 +460,20 @@ def _google_genai_streaming_hidden_params( } +def _log_generate_content_request( + logging_obj: LiteLLMLoggingObj, + contents: object, + api_base: str, + headers: Mapping[str, object], + body: Mapping[str, object], +) -> None: + logging_obj.pre_call( + input=contents, + api_key="", + additional_args={"complete_input_dict": body, "api_base": api_base, "headers": headers}, + ) + + @lru_cache(maxsize=None) def _responses_api_optional_request_param_names() -> frozenset[str]: return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys()) @@ -2581,14 +2600,16 @@ class BaseLLMHTTPHandler: litellm_params.get("additional_drop_params") or (), anthropic_messages_optional_request_params, ) - if ( - request_params.get("safeguards") is not None - and anthropic_messages_provider_config.should_filter_anthropic_beta_headers() - ): - await aload_beta_headers_config() affinity_api_base: Final = ( AnthropicModelInfo.get_api_base(api_base) if custom_llm_provider == "anthropic" else api_base ) + takes_another_beta_row: Final = ( + resolve_beta_provider(custom_llm_provider, affinity_api_base) != custom_llm_provider + ) + provider_filters_betas: Final = anthropic_messages_provider_config.should_filter_anthropic_beta_headers() + filters_beta_headers: Final = provider_filters_betas or takes_another_beta_row + if request_params.get("safeguards") is not None and filters_beta_headers: + await aload_beta_headers_config() default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) ( headers, @@ -2607,8 +2628,14 @@ class BaseLLMHTTPHandler: api_base=api_base, ) - if anthropic_messages_provider_config.should_filter_anthropic_beta_headers(): - headers = update_headers_with_filtered_beta(headers=headers, provider=custom_llm_provider) + if filters_beta_headers and headers.get("anthropic-beta") is not None: + await aload_beta_headers_config() + headers = update_headers_with_filtered_beta( + headers=headers, + provider=custom_llm_provider, + api_base=affinity_api_base, + unlisted="drop" if provider_filters_betas else "forward", + ) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -6226,12 +6253,13 @@ class BaseLLMHTTPHandler: ) try: - response: ResponsesAPIResponse | BaseResponsesAPIStreamingIterator = await litellm.aresponses( - model=patch.model or model, - input=patch.messages, # pyright: ignore[reportArgumentType] # pre-existing mismatch surfaced by the Router import; patch messages are valid response input at runtime - **optional_params, - **kwargs_for_followup, - ) + with internal_sub_call(logging_obj.litellm_call_id): + response: ResponsesAPIResponse | BaseResponsesAPIStreamingIterator = await litellm.aresponses( + model=patch.model or model, + input=patch.messages, # pyright: ignore[reportArgumentType] # pre-existing mismatch surfaced by the Router import; patch messages are valid response input at runtime + **optional_params, + **kwargs_for_followup, + ) if callback is not None: try: @@ -6320,6 +6348,7 @@ class BaseLLMHTTPHandler: max_loops: int, fingerprints: list[str], fingerprint: str, + logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> ModelResponse | CustomStreamWrapper: patch: Final = plan.request_patch or AgenticLoopRequestPatch() if patch.messages is None: @@ -6359,12 +6388,14 @@ class BaseLLMHTTPHandler: fingerprint=fingerprint, ) - return await litellm.acompletion( - model=full_model_name, - messages=patch.messages, - **optional_params_for_followup, - **kwargs_for_followup, - ) + with internal_sub_call(logging_obj.litellm_call_id if logging_obj is not None else None): + response: Final = await litellm.acompletion( + model=full_model_name, + messages=patch.messages, + **optional_params_for_followup, + **kwargs_for_followup, + ) + return response def _maybe_wrap_in_fake_stream( self, @@ -6713,6 +6744,7 @@ class BaseLLMHTTPHandler: messages=messages, optional_params=optional_params, kwargs=kwargs_with_provider, + logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, depth=depth, max_loops=max_loops, @@ -7065,7 +7097,7 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout, provider_config: BaseRealtimeHTTPConfig | None = None, model: str | None = None, - extra_headers: dict[str, object] | None = None, + extra_headers: Mapping[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, api_version: str | None = None, ) -> httpx.Response: @@ -7098,7 +7130,7 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout, provider_config: BaseRealtimeHTTPConfig | None = None, model: str | None = None, - extra_headers: dict[str, object] | None = None, + extra_headers: Mapping[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, api_version: str | None = None, ) -> httpx.Response: @@ -7117,9 +7149,8 @@ class BaseLLMHTTPHandler: api_version=api_version, ) - async def _async_realtime_session_post( + async def async_realtime_translation_client_secret_handler( self, - endpoint: Literal["client_secrets", "transcription_sessions"], api_base: str, api_key: str, request_data: dict[str, object], @@ -7127,7 +7158,35 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout, provider_config: BaseRealtimeHTTPConfig | None = None, model: str | None = None, - extra_headers: dict[str, object] | None = None, + extra_headers: Mapping[str, object] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + api_version: str | None = None, + ) -> httpx.Response: + return await self._async_realtime_session_post( + endpoint="translation_client_secrets", + api_base=api_base, + api_key=api_key, + request_data=request_data, + logging_obj=logging_obj, + timeout=timeout, + provider_config=provider_config, + model=model, + extra_headers=extra_headers, + client=client, + api_version=api_version, + ) + + async def _async_realtime_session_post( + self, + endpoint: Literal["client_secrets", "transcription_sessions", "translation_client_secrets"], + api_base: str, + api_key: str, + request_data: dict[str, object], + logging_obj: LiteLLMLoggingObj, + timeout: float | httpx.Timeout, + provider_config: BaseRealtimeHTTPConfig | None = None, + model: str | None = None, + extra_headers: Mapping[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, api_version: str | None = None, ) -> httpx.Response: @@ -7150,13 +7209,20 @@ class BaseLLMHTTPHandler: url = provider_config.get_transcription_session_url( api_base=api_base, model=model or "", api_version=api_version ) + elif endpoint == "translation_client_secrets": + url = provider_config.get_translation_client_secret_url( + api_base=api_base, model=model or "", api_version=api_version + ) else: url = provider_config.get_complete_url(api_base=api_base, model=model or "", api_version=api_version) headers: dict[str, object] = provider_config.validate_environment( headers={}, model=model or "", api_key=api_key ) else: - url = join_api_base(api_base, f"v1/realtime/{endpoint}") + endpoint_path: Final = ( + "translations/client_secrets" if endpoint == "translation_client_secrets" else endpoint + ) + url = join_api_base(api_base, f"v1/realtime/{endpoint_path}") headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", @@ -7190,6 +7256,12 @@ class BaseLLMHTTPHandler: ) raise + @staticmethod + def _get_realtime_async_http_client(client: object | None) -> AsyncHTTPHandler: + if isinstance(client, AsyncHTTPHandler): + return client + return get_async_httpx_client(llm_provider=litellm.LlmProviders.OPENAI) + async def async_realtime_calls_handler( self, api_base: str, @@ -7199,10 +7271,11 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout, provider_config: BaseRealtimeHTTPConfig | None = None, model: str | None = None, - session_config: dict[str, object] | None = None, - extra_headers: dict[str, object] | None = None, + session_config: Mapping[str, object] | None = None, + extra_headers: Mapping[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, api_version: str | None = None, + translation: bool = False, ) -> httpx.Response: """ Forward POST /v1/realtime/calls (SDP exchange) to upstream provider. @@ -7214,18 +7287,30 @@ class BaseLLMHTTPHandler: - sdp: the SDP offer (text) - session: JSON string with {"type": "realtime", "model": "...", ...} """ - if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.OPENAI, - ) - else: - async_httpx_client = client + session_data: Final[dict[str, object]] = { # mutable-ok: model and session type are resolved locally + **(session_config or {}) + } + if "type" not in session_data: + session_data["type"] = "translation" if translation else "realtime" + if "model" not in session_data and model: + session_data["model"] = model + + sdp_text: Final = sdp_body.decode("utf-8") if isinstance(sdp_body, bytes) else sdp_body + + async_httpx_client: Final = self._get_realtime_async_http_client(client) if provider_config is not None: - url = provider_config.get_realtime_calls_url(api_base=api_base, model=model or "", api_version=api_version) + url = ( + provider_config.get_translation_calls_url(api_base=api_base, model=model or "", api_version=api_version) + if translation + else provider_config.get_realtime_calls_url( + api_base=api_base, model=model or "", api_version=api_version + ) + ) headers: dict[str, object] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) else: - url = join_api_base(api_base, "v1/realtime/calls") + path: Final = "translations/calls" if translation else "calls" + url = join_api_base(api_base, f"v1/realtime/{path}") headers = { "Authorization": f"Bearer {openai_ephemeral_key}", } @@ -7233,14 +7318,8 @@ class BaseLLMHTTPHandler: if extra_headers: headers.update(extra_headers) - # Build multipart form data: sdp + session JSON - session_data: Final = session_config or {} - if "type" not in session_data: - session_data["type"] = "realtime" - if "model" not in session_data and model: - session_data["model"] = model - - sdp_text: Final = sdp_body.decode("utf-8") if isinstance(sdp_body, bytes) else sdp_body + if translation: + headers["Content-Type"] = "application/sdp" files: Final = { "sdp": (None, sdp_text, "text/plain"), @@ -7258,12 +7337,14 @@ class BaseLLMHTTPHandler: ) try: - return await async_httpx_client.post( - url=url, - headers=headers, - files=files, - timeout=timeout, - ) + if translation: + return await async_httpx_client.post( + url=url, + headers=headers, + content=sdp_text, + timeout=timeout, + ) + return await async_httpx_client.post(url=url, headers=headers, files=files, timeout=timeout) except Exception as e: if provider_config is not None: raise self._handle_error( @@ -8073,7 +8154,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8269,7 +8350,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8450,7 +8531,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8614,7 +8695,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8755,7 +8836,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8928,7 +9009,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9115,7 +9196,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9248,7 +9329,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9330,7 +9411,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9515,7 +9596,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -12443,6 +12524,10 @@ class BaseLLMHTTPHandler: from litellm.google_genai.streaming_iterator import ( GoogleGenAIGenerateContentStreamingIterator, ) + from litellm.llms.vertex_ai.gemini.rejected_thought_signature_retry import ( + post_with_accepted_thought_signatures_sync, + with_placeholder_native_thought_signatures, + ) if _is_async: return self.async_generate_content_handler( @@ -12494,27 +12579,19 @@ class BaseLLMHTTPHandler: if extra_body: data.update(extra_body) - ## LOGGING - logging_obj.pre_call( - input=contents, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": headers, - }, - ) + log_request: Final = partial(_log_generate_content_request, logging_obj, contents, api_base, headers) + log_request(data) + + def post(body: Mapping[str, object]) -> httpx.Response: + return sync_httpx_client.post( + url=api_base, headers=headers, json=dict(body), timeout=timeout, stream=stream + ) try: + sent_body, response = post_with_accepted_thought_signatures_sync( + data, with_placeholder_native_thought_signatures, None, post, log_request + ) if stream: - response = sync_httpx_client.post( - url=api_base, - headers=headers, - json=data, - timeout=timeout, - stream=True, - ) - # Return streaming iterator return GoogleGenAIGenerateContentStreamingIterator( response=response, model=model, @@ -12522,7 +12599,7 @@ class BaseLLMHTTPHandler: generate_content_provider_config=generate_content_provider_config, litellm_metadata=litellm_metadata or {}, custom_llm_provider=custom_llm_provider, - request_body=data, + request_body=dict(sent_body), hidden_params=_google_genai_streaming_hidden_params( api_base=api_base, litellm_params=litellm_params, @@ -12530,13 +12607,6 @@ class BaseLLMHTTPHandler: response_headers=response.headers, ), ) - else: - response = sync_httpx_client.post( - url=api_base, - headers=headers, - json=data, - timeout=timeout, - ) except Exception as e: raise self._handle_error( e=e, @@ -12574,6 +12644,10 @@ class BaseLLMHTTPHandler: from litellm.google_genai.streaming_iterator import ( AsyncGoogleGenAIGenerateContentStreamingIterator, ) + from litellm.llms.vertex_ai.gemini.rejected_thought_signature_retry import ( + post_with_accepted_thought_signatures_async, + with_placeholder_native_thought_signatures, + ) if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -12609,27 +12683,19 @@ class BaseLLMHTTPHandler: if extra_body: data.update(extra_body) - ## LOGGING - logging_obj.pre_call( - input=contents, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": headers, - }, - ) + log_request: Final = partial(_log_generate_content_request, logging_obj, contents, api_base, headers) + log_request(data) + + async def post(body: Mapping[str, object]) -> httpx.Response: + return await async_httpx_client.post( + url=api_base, headers=headers, json=dict(body), timeout=timeout, stream=stream + ) try: + sent_body, response = await post_with_accepted_thought_signatures_async( + data, with_placeholder_native_thought_signatures, None, post, log_request + ) if stream: - response = await async_httpx_client.post( - url=api_base, - headers=headers, - json=data, - timeout=timeout, - stream=True, - ) - # Return async streaming iterator return AsyncGoogleGenAIGenerateContentStreamingIterator( response=response, model=model, @@ -12637,7 +12703,7 @@ class BaseLLMHTTPHandler: generate_content_provider_config=generate_content_provider_config, litellm_metadata=litellm_metadata or {}, custom_llm_provider=custom_llm_provider, - request_body=data, + request_body=dict(sent_body), hidden_params=_google_genai_streaming_hidden_params( api_base=api_base, litellm_params=litellm_params, @@ -12645,13 +12711,6 @@ class BaseLLMHTTPHandler: response_headers=response.headers, ), ) - else: - response = await async_httpx_client.post( - url=api_base, - headers=headers, - json=data, - timeout=timeout, - ) except Exception as e: raise self._handle_error( e=e, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 60081c38073..b5393fa1603 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -4,6 +4,7 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion import os from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from itertools import chain, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload @@ -12,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.anthropic_beta_gated_fields import log_beta_gated_strips, strip_beta_gated_thinking from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( handle_invalid_parallel_tool_calls, should_convert_tool_call_to_json_mode, @@ -23,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_name_from_message, strip_prefill_prefix_field, ) +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import ( @@ -65,7 +68,7 @@ from ...anthropic.chat.transformation import ( tools_without_unsupported_tool_search, ) from ...openai_like.chat.transformation import OpenAILikeChatConfig -from ..common_utils import DatabricksBase, DatabricksException +from ..common_utils import DATABRICKS_HOSTED_MODEL_PREFIXES, DatabricksBase, DatabricksException _RESPONSE_FORMAT_ADAPTER: ( # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object Final[TypeAdapter[dict[str, object] | None]] @@ -81,15 +84,139 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: ) +_CONTENT_BLOCKS: Final = TypeAdapter(tuple[object, ...]) + + +def _content_blocks(message: object) -> tuple[object, ...]: + match message.model_dump(exclude_none=True) if isinstance(message, BaseModel) else message: + case {"content": list() as content}: + return _CONTENT_BLOCKS.validate_python(content) + case _: + return () + + +def _plain_text(block: object) -> str | None: + match block: + case {"type": "text", "text": str() as text} if block == {"type": "text", "text": text}: + return text + case _: + return None + + +def _plain_texts(message: Mapping[str, object]) -> tuple[str, ...]: + blocks: Final = _content_blocks(message) + texts: Final = tuple(text for block in blocks if (text := _plain_text(block)) is not None) + return texts if len(texts) == len(blocks) else () + + def _collapse_single_text_block(message: Mapping[str, object]) -> Mapping[str, object]: - content: Final = message.get("content") - match content: - case [{"type": "text", "text": str() as text}] if content == [{"type": "text", "text": text}]: + match _plain_texts(message): + case (text,): return {**message, "content": text} case _: return message +def _without_text_cache_control(block: object) -> object: + match block: + case {"type": "text", "text": str() as text, "cache_control": object() as cache_control} if block == { + "type": "text", + "text": text, + "cache_control": cache_control, + }: + return {"type": "text", "text": text} + case _: + return block + + +def _without_cache_control(message: AllMessageValues) -> AllMessageValues: + uncached: Final = {key: value for key, value in message.items() if key != "cache_control"} + match content := uncached.get("content"): + case list(): + blocks: Final = [_without_text_cache_control(block) for block in _CONTENT_BLOCKS.validate_python(content)] + return cast(AllMessageValues, {**uncached, "content": blocks}) # cast-ok: spread keeps the shape + case _: + return cast(AllMessageValues, uncached) # cast-ok: only cache_control was left out + + +def _has_string_content(message: Mapping[str, object]) -> bool: + return isinstance(message.get("content"), str) + + +def _is_system(message: Mapping[str, object]) -> bool: + return message.get("role") == "system" + + +def _json_mode_message(message: AllMessageValues, drop_params: bool) -> AllMessageValues: + candidate: Final = _without_cache_control(message) if drop_params else message + texts: Final = _plain_texts(candidate) + if "cache_control" in candidate or not (texts or _has_string_content(candidate)): + return message + if sum(1 for text in texts if text.strip()) < 2: + return candidate + return cast(AllMessageValues, {**candidate, "content": "".join(texts)}) # cast-ok: spread keeps the shape + + +def _is_dropped_on_send(message: Mapping[str, object]) -> bool: + sanitized: Final = {**message} + _sanitize_empty_content(sanitized) + return _is_bare_assistant_message(sanitized) + + +def _is_system_or_dropped(message: Mapping[str, object]) -> bool: + return _is_system(message) or _is_dropped_on_send(message) + + +def _json_checked_start(messages: Sequence[AllMessageValues]) -> int: + trailing: Final = tuple(takewhile(_is_system_or_dropped, reversed(messages))) + last_sent_offset: Final = 0 if any(_is_system(message) for message in trailing) else 1 + return max(len(messages) - len(trailing) - last_sent_offset, 0) + + +def _json_object_ready_messages( + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited transform_request contract + response_format: object, + drop_params: bool, +) -> list[AllMessageValues]: # mutable-ok: inherited transform_request contract + match response_format: + case {"type": "json_object"} if "claude" not in model and model.startswith(DATABRICKS_HOSTED_MODEL_PREFIXES): + checked_start: Final = _json_checked_start(messages) + return [ + *messages[:checked_start], + *(_json_mode_message(message, drop_params) for message in messages[checked_start:]), + ] + case _: + return messages + + +def _is_empty_text_block(block: object) -> bool: + match block: + case {"type": "text", "text": str() as text}: + return not text.strip() + case {"type": "text", "text": object() as text}: + return not text + case {"type": "text"}: + return True + case _: + return False + + +def _non_string_text_type(block: object) -> str | None: + match block: + case {"type": "text", "text": str()}: + return None + case {"type": "text", "text": object() as text} if text: + return type(text).__name__ + case _: + return None + + +def _first_non_string_text_type(messages: Sequence[AllMessageValues]) -> str | None: + blocks: Final = chain.from_iterable(map(_content_blocks, messages)) + return next((type_name for block in blocks if (type_name := _non_string_text_type(block))), None) + + def _sanitize_empty_content(message_dict: dict[str, object]) -> None: """ Remove or filter content so empty text blocks are not sent. @@ -108,9 +235,7 @@ def _sanitize_empty_content(message_dict: dict[str, object]) -> None: message_dict.pop("content") return filtered: Final = [ - block - for block in content - if not (isinstance(block, dict) and block.get("type") == "text" and not (block.get("text") or "").strip()) + block for block in _CONTENT_BLOCKS.validate_python(content) if not _is_empty_text_block(block) ] if not filtered: message_dict.pop("content") @@ -343,7 +468,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): use_ai_gateway: Final = self.uses_ai_gateway(model) return join_api_base(self._get_api_base(api_base, use_ai_gateway=use_ai_gateway), "chat/completions") - def get_supported_openai_params(self, model: str | None = None) -> list: + def get_supported_openai_params(self, model: str | None = None) -> list[str]: return [ "stream", "stop", @@ -449,8 +574,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): drop_params: bool, replace_max_completion_tokens_with_max_tokens: bool = True, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) - mapped_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) + request_params: Final = ( + tool_params_without_hosted_mcp(non_default_params, model=model, route="databricks", drop_params=drop_params) + if "claude" in model + else non_default_params + ) + is_thinking_enabled: Final = self.is_thinking_enabled(request_params) + mapped_params: Final = super().map_openai_params(request_params, optional_params, model, drop_params) if "claude" in model: AnthropicConfig.translate_legacy_thinking_for_adaptive_model( model=model, optional_params=mapped_params, custom_llm_provider="databricks", drop_params=drop_params @@ -459,16 +589,16 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): mapped_params["tools"] = self._map_openai_to_dbrx_tool( model=model, tools=mapped_params["tools"], drop_params=drop_params ) - if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens: - mapped_params["max_tokens"] = non_default_params[ + if "max_completion_tokens" in request_params and replace_max_completion_tokens_with_max_tokens: + mapped_params["max_tokens"] = request_params[ "max_completion_tokens" ] # most openai-compatible providers support 'max_tokens' not 'max_completion_tokens' mapped_params.pop("max_completion_tokens", None) - if "response_format" in non_default_params and "claude" in model: + if "response_format" in request_params and "claude" in model: _tool: Final = self.map_response_format_to_databricks_tool( model, - non_default_params["response_format"], + request_params["response_format"], mapped_params, is_thinking_enabled, ) @@ -486,8 +616,8 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): "response_format", None ) # unsupported for claude models - if json_schema -> convert to tool call - if "reasoning_effort" in non_default_params and self._uses_anthropic_thinking_param(model): - reasoning_effort_value: Final = non_default_params.get("reasoning_effort") + if "reasoning_effort" in request_params and self._uses_anthropic_thinking_param(model): + reasoning_effort_value: Final = request_params.get("reasoning_effort") mapped_thinking: Final = AnthropicConfig._map_reasoning_effort( reasoning_effort=reasoning_effort_value, model=model, @@ -513,7 +643,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): optional_params.pop("reasoning_effort", None) ## handle thinking tokens self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=mapped_params + non_default_params=request_params, optional_params=mapped_params ) gated_thinking: Final = strip_beta_gated_thinking( mapped_params.get("thinking"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # mapped_params is the base config's bare dict @@ -535,6 +665,35 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): return False + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + headers: dict[str, object], # mutable-ok: inherited contract + ) -> dict[str, object]: # mutable-ok: inherited contract + match _first_non_string_text_type(messages): + case str() as type_name: + raise DatabricksException( + status_code=400, + message=f"A text content block needs a string 'text', got {type_name}", + ) + case None: + pass + return super().transform_request( + model=model, + messages=_json_object_ready_messages( + model, + messages, + optional_params.get("response_format"), + drop_params=drops_unsupported_params(litellm_params.get("drop_params")), + ), + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + @overload def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: Literal[True] diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index b1562bef6b8..e8313d9a003 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -56,6 +56,7 @@ DATABRICKS_WORKSPACE_ROUTE_PREFIXES: Final = ( DATABRICKS_AI_GATEWAY_MESSAGES_PATH, ) DATABRICKS_NATIVE_MESSAGES_MODEL_PREFIXES: Final = ("databricks-claude-", "system.ai.claude-") +DATABRICKS_HOSTED_MODEL_PREFIXES: Final = ("databricks-", "system.ai.") DATABRICKS_ENDPOINT_PATHS: Final = MappingProxyType( { "chat_completions": "/chat/completions", diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index 5b7c2a2c410..9c222f03c8c 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_search_results_text, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -33,17 +34,15 @@ def _warn_on_dropped_budget(thinking: object, model: str) -> None: class DeepSeekChatConfig(OpenAIGPTConfig): - def get_supported_openai_params(self, model: str) -> list: + def get_supported_openai_params(self, model: str) -> list[str]: """ DeepSeek reasoner models support thinking parameter. """ - params: Final = super().get_supported_openai_params(model) - params.extend(["thinking", "reasoning_effort"]) - return params + return [*super().get_supported_openai_params(model), "thinking", "reasoning_effort"] def map_openai_params( self, - non_default_params: dict, + non_default_params: dict[str, object], optional_params: dict, model: str, drop_params: bool, @@ -59,7 +58,12 @@ class DeepSeekChatConfig(OpenAIGPTConfig): Reference: https://api-docs.deepseek.com/guides/thinking_mode """ # Let parent handle standard params first - optional_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) + optional_params = super().map_openai_params( + tool_params_without_hosted_mcp(non_default_params, model=model, route="deepseek", drop_params=drop_params), + optional_params, + model, + drop_params, + ) # Pop thinking/reasoning_effort from optional_params first (parent may have added them) # Then re-add only if valid for DeepSeek @@ -265,6 +269,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): tools: Final = optional_params.get("tools") if not isinstance(tools, list) or not tools: return optional_params + sent_tools: Final[tuple[object, ...]] = tuple(tools) def _is_function_tool(tool: object) -> bool: return isinstance(tool, dict) and tool.get("type") == "function" @@ -289,14 +294,14 @@ class DeepSeekChatConfig(OpenAIGPTConfig): name: Final = function.get("name") return isinstance(name, str) and name in function_tool_names - function_tools: Final = [tool for tool in tools if _is_function_tool(tool)] - if len(function_tools) == len(tools): + function_tools: Final = [tool for tool in sent_tools if _is_function_tool(tool)] + if len(function_tools) == len(sent_tools): return optional_params dropped_types: Final = sorted( { str(tool.get("type")) if isinstance(tool, dict) else type(tool).__name__ - for tool in tools + for tool in sent_tools if not _is_function_tool(tool) } ) diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index 463b9cbfc9a..d599f0b65f7 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -29,6 +29,7 @@ from litellm.utils import ( declared_mid_conversation_system_support, deployment_declared_mid_conversation_system_support, deployment_model_info, + get_bundled_model_cost_map, ) if TYPE_CHECKING: @@ -54,6 +55,15 @@ def _forwarded_user_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, ) +def _bundled_cost_map_declares_mid_conversation_system(model: str) -> bool: + bundled_cost_map: Final[Mapping[str, object]] = get_bundled_model_cost_map() + names: Final = (model.removeprefix("fireworks_ai/"), resolve_fireworks_resource_name(model)) + return any( + deployment_declared_mid_conversation_system_support(bundled_cost_map.get(f"fireworks_ai/{name}")) is True + for name in names + ) + + class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): @property def custom_llm_provider(self) -> LlmProviders: @@ -90,9 +100,13 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): validated_input: Final = self._validate_input_param(input) model_info: Final = deployment_model_info(litellm_params) declared: Final = declared_mid_conversation_system_support(model, self.custom_llm_provider.value, model_info) + keeps_developer_items_in_place: Final = declared is True and ( + deployment_declared_mid_conversation_system_support(model_info) is True + or _bundled_cost_map_declares_mid_conversation_system(model) + ) folded: Final = ( None - if deployment_declared_mid_conversation_system_support(model_info) is True + if keeps_developer_items_in_place else instruction_items_folded( validated_input, response_api_optional_request_params, "user" if declared is False else "system" ) diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 936919af83d..1539f13455f 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -1,5 +1,5 @@ from collections.abc import Mapping -from typing import Final, cast +from typing import ClassVar, Final, cast import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -49,6 +49,8 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): Note: Please make sure to modify the default parameters as required for your use case. """ + _route: ClassVar[str] = "gemini" + temperature: float | None = None max_output_tokens: int | None = None top_p: float | None = None diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 88c6494e3f7..74cd6f22704 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -537,13 +537,16 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.llms.anthropic.messages.context_management.inference_view import ( + InferenceRequestViewError, + count_context_edits, + inference_request_view, + ) from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, - count_context_edits, count_context_error_response, gemini_count_contents, - inference_request_view, malformed_tool_result_text_count_error, provider_error_response, unconvertible_messages_error, @@ -556,7 +559,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): messages_format, count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use), ) - if isinstance(request_view, VertexGeminiCountContextError): + if isinstance(request_view, InferenceRequestViewError): return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools, model=model_to_use) if isinstance(count_context, VertexGeminiCountContextError): diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 81a186eadf4..7bc783b960f 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -17,6 +17,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.gigachat.utils import convert_usage, get_api_base @@ -169,7 +170,10 @@ class GigaChatConfig(BaseConfig): drop_params: bool, ) -> dict: # mutable-ok: base class contract returns dict """Map OpenAI parameters to GigaChat parameters.""" - for param, value in non_default_params.items(): + request_params: Final = tool_params_without_hosted_mcp( + dict(non_default_params), model=model, route="gigachat", drop_params=drop_params + ) + for param, value in request_params.items(): if param == "stream": optional_params["stream"] = value elif param == "temperature": diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index a1d1a3f6a19..5d997babc6e 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -38,7 +38,9 @@ class _SupportsTokenEncode(Protocol): def encode(self, text: str) -> list[int]: ... -def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object: +def get_hf_task_embedding_for_model( + model: str, task_type: str | None, api_base: str, timeout: float | httpx.Timeout | None = None +) -> object: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type @@ -46,7 +48,7 @@ def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: raise Exception(f"Invalid task_type={task_type}. Expected one of={hf_tasks_embeddings}") http_client: Final = HTTPHandler(concurrent_limit=1) - model_info: Final = http_client.get(url=f"{api_base}/api/models/{model}") + model_info: Final = http_client.get(url=f"{api_base}/api/models/{model}", timeout=timeout) model_info_dict: Final = decoded_json_object(_DECODED_JSON.validate_python(model_info.json())) @@ -148,6 +150,7 @@ class HuggingFaceEmbedding(BaseLLM): call_type: Literal["sync", "async"], optional_params: dict, embed_url: str, + timeout: float | httpx.Timeout | None = None, ) -> dict: data: dict = {} @@ -165,7 +168,9 @@ class HuggingFaceEmbedding(BaseLLM): task_type: Final = optional_params.pop("input_type", None) if call_type == "sync": - hf_task: Final = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) + hf_task: Final = get_hf_task_embedding_for_model( + model=model, task_type=task_type, api_base=HF_HUB_URL, timeout=timeout + ) elif call_type == "async": return self._async_transform_input(model=model, task_type=task_type, embed_url=embed_url, input=input) @@ -253,6 +258,7 @@ class HuggingFaceEmbedding(BaseLLM): call_type="sync", optional_params=optional_params, embed_url=api_base, + timeout=timeout, ) ## LOGGING @@ -269,9 +275,10 @@ class HuggingFaceEmbedding(BaseLLM): if client is None: client = get_async_httpx_client( llm_provider=litellm.LlmProviders.HUGGINGFACE, + params={"timeout": timeout}, ) - response: Final = await client.post(api_base, headers=headers, data=json.dumps(data)) + response: Final = await client.post(api_base, headers=headers, data=json.dumps(data), timeout=timeout) ## LOGGING logging_obj.post_call( @@ -321,7 +328,9 @@ class HuggingFaceEmbedding(BaseLLM): litellm_params=litellm_params, ) task_type: Final = optional_params.get("input_type", None) - task: Final = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) + task: Final = get_hf_task_embedding_for_model( + model=model, task_type=task_type, api_base=HF_HUB_URL, timeout=timeout + ) # print_verbose(f"{model}, {task}") embed_url = "" if model.startswith(("http://", "https://")): @@ -359,6 +368,7 @@ class HuggingFaceEmbedding(BaseLLM): call_type="sync", optional_params=optional_params, embed_url=embed_url, + timeout=timeout, ) ## LOGGING @@ -374,7 +384,7 @@ class HuggingFaceEmbedding(BaseLLM): ## COMPLETION CALL if client is None or not isinstance(client, HTTPHandler): client = HTTPHandler(concurrent_limit=1) - response: Final = client.post(embed_url, headers=headers, data=json.dumps(data)) + response: Final = client.post(embed_url, headers=headers, data=json.dumps(data), timeout=timeout) ## LOGGING logging_obj.post_call( diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 168b562b358..59020b058a2 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, cast from httpx._models import Headers, Response -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm._uuid import uuid @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_reasoning_content, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -40,6 +41,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_REQUEST_PARAMS: Final = TypeAdapter(dict[str, object]) + class OllamaChatConfig(BaseConfig): """ @@ -152,7 +155,13 @@ class OllamaChatConfig(BaseConfig): model: str, drop_params: bool, ) -> dict: - for param, value in non_default_params.items(): + request_params: Final = tool_params_without_hosted_mcp( + _REQUEST_PARAMS.validate_python(non_default_params), + model=model, + route="ollama_chat", + drop_params=drop_params, + ) + for param, value in request_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["num_predict"] = value if param == "stream": diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 9d551f6f859..0479ac5c36c 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + chat_completion_prompt_cache_breakpoints_for_wire, drop_non_python_regex_patterns, drop_tool_reference_parts_from_tool_messages, flatten_combinators_and_drop_non_python_regex_patterns, @@ -555,26 +556,38 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ) -> tuple[Sequence[AllMessageValues], Sequence["ChatCompletionToolParam"] | None]: """ A deployment that reads OpenAI's ``prompt_cache_breakpoint`` (gpt-5.6+ on - api.openai.com, or one opted in through ``prompt_cache_options`` or the cost - map) gets each Anthropic ``cache_control`` mark rewritten in place. Any other - endpoint keeps or strips the marks as before, and under drop_params also loses - every ``prompt_cache_breakpoint`` it cannot read. + api.openai.com, or one opted in through ``prompt_cache_options``, its own + ``model_info``, or the cost map) gets each + Anthropic ``cache_control`` mark rewritten in place, and under drop_params each + content part's ``prompt_cache_breakpoint`` cut down to what chat completions + accepts. Any other endpoint keeps or strips the marks as before, and under + drop_params also loses every ``prompt_cache_breakpoint`` it cannot read. """ - from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook + from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook, routed_model_info custom_llm_provider: Final = _litellm_params_str(litellm_params, "custom_llm_provider") api_base: Final = _litellm_params_str(litellm_params, "api_base") reads_breakpoints: Final = AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( - model, custom_llm_provider, api_base, prompt_cache_options + model, + custom_llm_provider, + api_base, + prompt_cache_options, + routed_model_info(litellm_params.get("model_info")), ) if not reads_breakpoints and self._should_preserve_cache_control_for_endpoint(custom_llm_provider, api_base): return ( _claude_messages_within_cache_control_cap(model, messages, tools, litellm_params, request_params), tools, ) - drops_breakpoints: Final = not reads_breakpoints and _drops_params(litellm_params) + drop_params: Final = _drops_params(litellm_params) + drops_breakpoints: Final = not reads_breakpoints and drop_params renamed: Final = ( - [AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(message) for message in messages] + [ + AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints( + chat_completion_prompt_cache_breakpoints_for_wire(message, drop_params) + ) + for message in messages + ] if reads_breakpoints else [without_prompt_cache_breakpoints(message) for message in messages] if drops_breakpoints diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 369b6959411..5aa72579d6a 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -307,10 +307,13 @@ def _openai_answer(question: DecisionsIRQuestion, answer: DecisionsIRAnswer) -> def ir_to_openai_response( - response: DecisionsIRResponse, request: DecisionsIRRequest, requested_model: str + response: DecisionsIRResponse, + request: DecisionsIRRequest, + requested_model: str, + evidence_fields: frozenset[str] = frozenset(), ) -> OpenAIDecisionResponse: usage: Final = response.usage or DecisionsIRUsage() - return OpenAIDecisionResponse( + formatted: Final = OpenAIDecisionResponse( model=response.model or requested_model, answers=tuple( _openai_answer(question, answer) @@ -327,6 +330,7 @@ def ir_to_openai_response( total_tokens=usage.input_tokens + usage.output_tokens, ), ) + return formatted.model_copy(update={name: response.extra[name] for name in evidence_fields & response.extra.keys()}) class OpenAIDecisionsConfig(BaseDecisionsConfig): diff --git a/litellm/llms/openai/live/handler.py b/litellm/llms/openai/live/handler.py index 297a147a288..5b899ac4162 100644 --- a/litellm/llms/openai/live/handler.py +++ b/litellm/llms/openai/live/handler.py @@ -18,6 +18,7 @@ from litellm._logging import redact_secrets from litellm.constants import ( OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS, OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, ) @@ -316,12 +317,15 @@ async def _forward_live_backend_frames( backend_websocket: LiveBackendWebSocket, client_gone: asyncio.Event, event_retention: _LiveEventRetention, + upstream_ended_abnormally: asyncio.Event, ) -> None: from starlette.websockets import WebSocketDisconnect from websockets.exceptions import ConnectionClosed event_loop: Final = asyncio.get_running_loop() - async for message, close_deadline in _live_backend_messages(websocket, backend_websocket): + async for message, close_deadline in _live_backend_messages( + websocket, backend_websocket, upstream_ended_abnormally + ): if (event := _metering_event(message)) is not None: event_retention.retain(event) if client_gone.is_set(): @@ -341,6 +345,7 @@ async def _forward_live_backend_frames( async def _live_backend_messages( websocket: LiveClientWebSocket, backend_websocket: LiveBackendWebSocket, + upstream_ended_abnormally: asyncio.Event, ) -> AsyncIterator[tuple[str, float | None]]: from starlette.websockets import WebSocketDisconnect from websockets.exceptions import ConnectionClosed @@ -350,6 +355,8 @@ async def _live_backend_messages( yield message, None except ConnectionClosed as error: upstream_close: Final = backend_close_from(error) + if not upstream_close.ended_cleanly: + upstream_ended_abnormally.set() with suppress(WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): await websocket.close(code=upstream_close.code, reason=upstream_close.reason) return @@ -434,6 +441,7 @@ class OpenAILiveSessions(OpenAIRealtime): "complete_input_dict": {"session_start": rewritten_start}, } event_retention: Final = _LiveEventRetention() + session_ended_abnormally: Final = asyncio.Event() try: logging_obj.pre_call(input=None, api_key=api_key, additional_args=pre_call_args) async with self._websocket_connector( @@ -448,7 +456,11 @@ class OpenAILiveSessions(OpenAIRealtime): websocket=websocket, backend_websocket=backend_websocket, event_retention=event_retention, + upstream_ended_abnormally=session_ended_abnormally, ) + except Exception: + session_ended_abnormally.set() + raise finally: session_end_time: Final = self._clock() retained_events: Final = event_retention.results() @@ -463,6 +475,10 @@ class OpenAILiveSessions(OpenAIRealtime): logging_obj.model_call_details[ # rebind-ok: flags the shared logging object as already logged REALTIME_SESSION_SUCCESS_LOGGED_KEY ] = True + if session_ended_abnormally.is_set(): + logging_obj.model_call_details[ # rebind-ok: flags the shared logging object for the router + REALTIME_SESSION_ENDED_ABNORMALLY_KEY + ] = True except InvalidStatus as error: await close_after_upstream_handshake_refusal(websocket, error.response.status_code) except Exception as error: # noqa: BLE001 # any relay failure must still close the client socket with a redacted reason @@ -484,6 +500,7 @@ class OpenAILiveSessions(OpenAIRealtime): websocket: LiveClientWebSocket, backend_websocket: LiveBackendWebSocket, event_retention: _LiveEventRetention, + upstream_ended_abnormally: asyncio.Event, ) -> None: from websockets.exceptions import ConnectionClosed @@ -503,6 +520,7 @@ class OpenAILiveSessions(OpenAIRealtime): backend_websocket=backend_websocket, client_gone=client_gone, event_retention=event_retention, + upstream_ended_abnormally=upstream_ended_abnormally, ) ) client_gone_task: Final = asyncio.create_task(client_gone.wait()) diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 5be488f0f6e..40ead28d952 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -141,7 +141,7 @@ class OpenAIRealtime(OpenAIChatCompletion): return ssl_config - def _construct_url(self, api_base: str, query_params: RealtimeQueryParams) -> str: + def _construct_url(self, api_base: str, query_params: RealtimeQueryParams, realtime_mode: str = "realtime") -> str: """ Construct the backend websocket URL with the client's query parameters. @@ -158,7 +158,8 @@ class OpenAIRealtime(OpenAIChatCompletion): api_base = api_base.replace("http://", "ws://") url = URL(api_base) # Set the correct path - url = url.copy_with(path="/v1/realtime") + path: Final = "/v1/realtime/translations" if realtime_mode == "translation" else "/v1/realtime" + url = url.copy_with(path=path) upstream_params: Final = tuple( (key, value) for key, value in query_params.items() if not (drops_model and key == "model") ) @@ -197,6 +198,8 @@ class OpenAIRealtime(OpenAIChatCompletion): query_params: RealtimeQueryParams | None = None, user_api_key_dict: object | None = None, litellm_metadata: dict | None = None, + realtime_mode: str = "realtime", + drop_params: bool = False, **kwargs: object, ): import websockets @@ -210,13 +213,15 @@ class OpenAIRealtime(OpenAIChatCompletion): # Use all query params if provided, else fallback to just model if query_params is None: query_params = {"model": model} - url: Final = self._construct_url(api_base, query_params) + url: Final = self._construct_url(api_base, query_params, realtime_mode=realtime_mode) try: # Get provider-specific SSL configuration ssl_config: Final = self._get_ssl_config(url) - openai_beta_realtime: Final = client_sent_openai_beta_realtime_header(websocket) + openai_beta_realtime: Final = realtime_mode != "translation" and client_sent_openai_beta_realtime_header( + websocket + ) if not openai_beta_realtime: verbose_logger.debug( "OpenAI Realtime: connecting with GA protocol (no OpenAI-Beta header). " @@ -253,6 +258,8 @@ class OpenAIRealtime(OpenAIChatCompletion): ), event_normalizer=self._make_event_normalizer(), usage_meter=self._make_usage_meter(), + translation_session=realtime_mode == "translation", + drop_params=drop_params, ) await realtime_streaming.bidirectional_forward() diff --git a/litellm/llms/openai/realtime/http_transformation.py b/litellm/llms/openai/realtime/http_transformation.py index 49d33582c6f..ab32ae6085e 100644 --- a/litellm/llms/openai/realtime/http_transformation.py +++ b/litellm/llms/openai/realtime/http_transformation.py @@ -28,6 +28,18 @@ class OpenAIRealtimeHTTPConfig(BaseRealtimeHTTPConfig): base = base.removesuffix("/v1") return join_api_base(base, "v1/realtime/transcription_sessions") + def get_translation_client_secret_url( + self, api_base: str | None, model: str, api_version: str | None = None + ) -> str: + base = self.get_api_base(api_base).rstrip("/") + base = base.removesuffix("/v1") + return f"{base}/v1/realtime/translations/client_secrets" + + def get_translation_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: + base = self.get_api_base(api_base).rstrip("/") + base = base.removesuffix("/v1") + return f"{base}/v1/realtime/translations/calls" + def validate_environment( self, headers: dict, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a905bff3521..a18cbe13925 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1988,7 +1988,7 @@ class OpenAIResponsesHandler(BaseTranslation): "annotations": (), "logprobs": None, } - return ( + events: Final = ( *_open_item_closing_events(responses_so_far), OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -2033,6 +2033,10 @@ class OpenAIResponsesHandler(BaseTranslation): response=_blocked_response(exc, response_id=response_id, model=model, output_item=item), ), ) + return tuple( + event.model_copy(update={"sequence_number": number}) + for number, event in enumerate(events, start=_next_stream_sequence_number(responses_so_far)) + ) class _BlockedContentPart(TypedDict): diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 7cca31c5d3f..6bbf3598bd8 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,5 +1,5 @@ from collections.abc import Callable, Collection, Iterable, Mapping, Sequence -from functools import lru_cache +from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints @@ -12,10 +12,11 @@ import litellm from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, - openai_prompt_cache_breakpoint_support, + deployment_model_info, + deployment_prompt_cache_breakpoint_support, targets_openai_api, ) -from litellm.litellm_core_utils.core_helpers import normalize_drop_params, process_response_headers +from litellm.litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params, process_response_headers from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( safe_convert_created_field, @@ -25,6 +26,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_non_python_regex_patterns, flatten_combinators_and_drop_non_python_regex_patterns, prompt_cache_breakpoint_for_wire, + whole_prompt_cache_breakpoint_or_none, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -44,6 +46,7 @@ from litellm.responses.litellm_completion_transformation.tool_search import ( tools_with_anthropic_tool_search_mapped, tools_without_anthropic_tool_search, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.secret_managers.main import get_secret_str from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import * @@ -103,13 +106,6 @@ _INPUT_ITEM: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, _CONTENT_PARTS: Final[TypeAdapter[Sequence[Mapping[str, object]]]] = TypeAdapter(Sequence[Mapping[str, object]]) -class _DeclaredModel(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) - - base_model: object = None - model_info: object = None - - class _DeploymentHost(BaseModel): """`responses()` builds the params from raw kwargs, where callers may spell the host `base_url`.""" @@ -241,12 +237,12 @@ def _json_array_or_empty(value: object) -> Sequence[object]: return value if isinstance(value, (list, tuple)) else () -def _item_with_wire_safe_prompt_cache_breakpoints(item: object, keeps_markers: bool) -> object: +def _item_with_wire_safe_prompt_cache_breakpoints(item: object, to_wire: Callable[[object], object] | None) -> object: fields: Final = _json_object_or_none(item) if fields is None or not _carries_prompt_cache_breakpoint(fields): return item wire_safe_fields: Final = { - key: _parts_with_wire_safe_prompt_cache_breakpoints(key, value, keeps_markers) + key: _parts_with_wire_safe_prompt_cache_breakpoints(key, value, to_wire) for key, value in fields.items() if key != "prompt_cache_breakpoint" } @@ -270,21 +266,23 @@ def _parts_carry_prompt_cache_breakpoint(parts: object) -> bool: return any(isinstance(part, Mapping) and "prompt_cache_breakpoint" in part for part in _json_array_or_empty(parts)) -def _parts_with_wire_safe_prompt_cache_breakpoints(key: str, value: object, keeps_markers: bool) -> object: +def _parts_with_wire_safe_prompt_cache_breakpoints( + key: str, value: object, to_wire: Callable[[object], object] | None +) -> object: if key not in _INPUT_PART_LIST_KEYS or not _parts_carry_prompt_cache_breakpoint(value): return value parts: Final = _json_array_or_empty(value) - wire_safe_parts: Final = tuple(_part_with_wire_safe_prompt_cache_breakpoint(part, keeps_markers) for part in parts) + wire_safe_parts: Final = tuple(_part_with_wire_safe_prompt_cache_breakpoint(part, to_wire) for part in parts) return value if _same_objects(wire_safe_parts, parts) else list(wire_safe_parts) -def _part_with_wire_safe_prompt_cache_breakpoint(part: object, keeps_markers: bool) -> object: +def _part_with_wire_safe_prompt_cache_breakpoint(part: object, to_wire: Callable[[object], object] | None) -> object: fields: Final = _json_object_or_none(part) if fields is None or "prompt_cache_breakpoint" not in fields: return part marker: Final = fields["prompt_cache_breakpoint"] - wire_marker: Final = prompt_cache_breakpoint_for_wire(marker, drop_params=True) if keeps_markers else None - if keeps_markers and wire_marker == marker: + wire_marker: Final = None if to_wire is None else to_wire(marker) + if to_wire is not None and wire_marker == marker: return part unmarked: Final = {key: value for key, value in fields.items() if key != "prompt_cache_breakpoint"} return with_prompt_cache_breakpoint(unmarked, wire_marker) @@ -623,11 +621,45 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): "model": model, "input": replay_safe_input, **response_api_optional_request_params, + **self._prompt_cache_options_with_original_comparison_id(request_params), }, ) return final_request_params + def transform_extra_body( + self, + extra_body: Mapping[str, object], + request: Mapping[str, object], + model: str, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object]: + return {**extra_body, **self._prompt_cache_options_with_original_comparison_id(extra_body)} + + def _prompt_cache_options_with_original_comparison_id(self, params: Mapping[str, object]) -> Mapping[str, object]: + options: Final = as_str_mapping(params.get("prompt_cache_options")) + comparison_response_id: Final = None if options is None else options.get("comparison_response_id") + if ( + self.custom_llm_provider != LlmProviders.OPENAI + or options is None + or not isinstance(comparison_response_id, str) + or not comparison_response_id.startswith("resp_") + ): + return {} + decoded_id: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(comparison_response_id) + decoded_provider: Final = decoded_id.get("custom_llm_provider") + original_response_id: Final = decoded_id.get("response_id", comparison_response_id) + rebuilt_envelope: Final = ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + {"id": original_response_id}, decoded_provider, {"model_info": {"id": decoded_id.get("model_id")}} + ) + if ( + decoded_provider != LlmProviders.OPENAI + or original_response_id == comparison_response_id + or rebuilt_envelope.get("id") != comparison_response_id + ): + return {} + return {"prompt_cache_options": {**options, "comparison_response_id": original_response_id}} + def _prepared_input_and_tools( self, model: str, @@ -681,24 +713,31 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) -> str | ResponseInputParam: if not isinstance(input, list) or not drops_params(litellm_params): return input - keeps_markers: Final = reads_breakpoints or self._keeps_prompt_cache_breakpoints( - model=model, litellm_params=litellm_params - ) - wire_safe_items: Final = tuple( - _item_with_wire_safe_prompt_cache_breakpoints(item, keeps_markers) for item in input + known_support: Final = self._prompt_cache_breakpoint_support(model=model, litellm_params=litellm_params) + to_wire: Final = ( + self._prompt_cache_breakpoint_to_wire(litellm_params, reads_breakpoints, known_support) + if reads_breakpoints or known_support is not False + else None ) + wire_safe_items: Final = tuple(_item_with_wire_safe_prompt_cache_breakpoints(item, to_wire) for item in input) if _same_objects(wire_safe_items, input): return input wire_safe_input: Final = list(wire_safe_items) return cast("ResponseInputParam", wire_safe_input) # cast-ok: items keep their shape, minus a rejected marker - def _keeps_prompt_cache_breakpoints(self, model: str, litellm_params: GenericLiteLLMParams) -> bool: - provider: Final = _provider_name(self.custom_llm_provider) - base_model: Final = self._declared_base_model(model=model, litellm_params=litellm_params) - known_support: Final = ( - openai_prompt_cache_breakpoint_support(name, provider) for name in dict.fromkeys((base_model, model)) + def _prompt_cache_breakpoint_to_wire( + self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None + ) -> Callable[[object], object]: + opted_in_openai_host: Final = reads_breakpoints and self.custom_llm_provider == LlmProviders.OPENAI + known_openai_reader: Final = known_support is True and self._targets_openai_api(litellm_params) + if opted_in_openai_host or known_openai_reader: + return partial(prompt_cache_breakpoint_for_wire, drop_params=True) + return whole_prompt_cache_breakpoint_or_none + + def _prompt_cache_breakpoint_support(self, model: str, litellm_params: GenericLiteLLMParams) -> bool | None: + return deployment_prompt_cache_breakpoint_support( + model, _provider_name(self.custom_llm_provider), litellm_params ) - return next((support for support in known_support if support is not None), None) is not False def _targets_openai_api(self, litellm_params: GenericLiteLLMParams) -> bool: host: Final = _DeploymentHost.model_validate(litellm_params) @@ -729,6 +768,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): custom_llm_provider=_provider_name(self.custom_llm_provider), api_base=host.api_base or host.base_url, prompt_cache_options=prompt_cache_options, + model_info=deployment_model_info(litellm_params), ) @staticmethod @@ -872,13 +912,6 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): base_model: Final[object] = model_info.get("base_model") if isinstance(model_info, dict) else None return base_model if isinstance(base_model, str) and base_model else model - @staticmethod - def _declared_base_model(model: str, litellm_params: GenericLiteLLMParams) -> str: - declared: Final = _DeclaredModel.model_validate(litellm_params) - model_info: Final = _json_object_or_none(declared.model_info) - candidates: Final = (declared.base_model, None if model_info is None else model_info.get("base_model")) - return next((name for name in candidates if isinstance(name, str) and name), model) - @staticmethod def _sanitized_tool_entry( entry: Mapping[str, object], diff --git a/litellm/llms/openai/transcriptions/gpt_transformation.py b/litellm/llms/openai/transcriptions/gpt_transformation.py index e7f8dd04516..b9eaa204567 100644 --- a/litellm/llms/openai/transcriptions/gpt_transformation.py +++ b/litellm/llms/openai/transcriptions/gpt_transformation.py @@ -26,8 +26,8 @@ class OpenAIGPTAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig): self, model: str, audio_file: FileTypes, - optional_params: dict, - litellm_params: dict, + optional_params: dict[str, object], + litellm_params: dict[str, object], ) -> AudioTranscriptionRequestData: """ Transform the audio transcription request @@ -37,3 +37,14 @@ class OpenAIGPTAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig): return AudioTranscriptionRequestData( data=data, ) + + +class OpenAIGPTTranscribeAudioTranscriptionConfig(OpenAIGPTAudioTranscriptionConfig): + def get_supported_openai_params(self, model: str) -> list[OpenAIAudioTranscriptionOptionalParams]: + return [ + "prompt", + "response_format", + "keywords", + "languages", + "stream", + ] diff --git a/litellm/llms/openai/transcriptions/handler.py b/litellm/llms/openai/transcriptions/handler.py index 48a44575989..23c5e70a3cc 100644 --- a/litellm/llms/openai/transcriptions/handler.py +++ b/litellm/llms/openai/transcriptions/handler.py @@ -1,8 +1,9 @@ -from collections.abc import Mapping +from collections.abc import Coroutine, Mapping from typing import TYPE_CHECKING, Final, Optional, cast import httpx -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream +from openai.types.audio import TranscriptionStreamEvent from pydantic import BaseModel if TYPE_CHECKING: @@ -12,6 +13,9 @@ from litellm.litellm_core_utils.core_helpers import set_provider_response_header from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, + acreate_sdk_transcription, + create_sdk_transcription, + sdk_compatible_transcription_request_data, ) from litellm.types.utils import FileTypes from litellm.utils import ( @@ -28,31 +32,18 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): async def make_openai_audio_transcriptions_request( self, openai_aclient: AsyncOpenAI, - data: dict, + data: Mapping[str, object], timeout: float | httpx.Timeout, ): - try: - raw_response = await openai_aclient.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) - headers: Final = dict(raw_response.headers) - response: Final = raw_response.parse() - - return headers, response - except Exception as e: - raise e + return await acreate_sdk_transcription(openai_aclient, sdk_compatible_transcription_request_data(data), timeout) def make_sync_openai_audio_transcriptions_request( self, openai_client: OpenAI, - data: dict, + data: Mapping[str, object], timeout: float | httpx.Timeout, ): - try: - raw_response: Final = openai_client.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) - headers: Final = dict(raw_response.headers) - response: Final = raw_response.parse() - return headers, response - except Exception as e: - raise e + return create_sdk_transcription(openai_client, sdk_compatible_transcription_request_data(data), timeout) def audio_transcriptions( self, @@ -70,7 +61,11 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): atranscription: bool = False, provider_config: BaseAudioTranscriptionConfig | None = None, shared_session: Optional["ClientSession"] = None, - ) -> TranscriptionResponse: + ) -> ( + TranscriptionResponse + | Stream[TranscriptionStreamEvent] + | Coroutine[object, object, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] + ): """ Handle audio transcription request """ @@ -128,6 +123,9 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): ) logging_obj.model_call_details["response_headers"] = headers + if isinstance(response, Stream): + return response + if isinstance(response, BaseModel): stringified_response = response.model_dump() else: @@ -153,7 +151,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): async def async_audio_transcriptions( self, audio_file: FileTypes, - data: dict, + data: Mapping[str, object], model_response: TranscriptionResponse, timeout: float, logging_obj: LiteLLMLoggingObj, @@ -192,6 +190,8 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): timeout=timeout, ) logging_obj.model_call_details["response_headers"] = headers + if isinstance(response, AsyncStream): + return response if isinstance(response, BaseModel): stringified_response = response.model_dump() else: diff --git a/litellm/llms/opencode/chat/transformation.py b/litellm/llms/opencode/chat/transformation.py index dcb30b5835c..c6d288b6d84 100644 --- a/litellm/llms/opencode/chat/transformation.py +++ b/litellm/llms/opencode/chat/transformation.py @@ -1,5 +1,5 @@ from collections.abc import AsyncIterator, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, ClassVar, Final import httpx from pydantic import TypeAdapter @@ -112,6 +112,9 @@ class OpenCodeGoMessagesChatConfig(_GoReportedCostMixin, AnthropicConfig): def custom_llm_provider(self) -> str: return "opencode_go" + def translates_prompt_cache_breakpoint(self, model: str) -> bool: + return "claude" in model.lower() + def get_complete_url( self, api_base: str | None, @@ -159,6 +162,7 @@ class OpenCodeGoMessagesChatConfig(_GoReportedCostMixin, AnthropicConfig): class OpenCodeGoGeminiChatConfig(_GoReportedCostMixin, GoogleAIStudioGeminiConfig): go_usage_key: str = "usageMetadata" + _route: ClassVar[str] = "opencode_go" @property def supports_stream_param_in_request_body(self) -> bool: diff --git a/litellm/llms/parallel_ai/decisions/__init__.py b/litellm/llms/parallel_ai/decisions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/parallel_ai/decisions/transformation.py b/litellm/llms/parallel_ai/decisions/transformation.py new file mode 100644 index 00000000000..34d03047301 --- /dev/null +++ b/litellm/llms/parallel_ai/decisions/transformation.py @@ -0,0 +1,89 @@ +from collections.abc import Mapping +from typing import Final + +import httpx +from pydantic import BaseModel, ConfigDict + +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, ir_to_systemone_request +from litellm.types.decisions import ( + DecisionsIRRequest, + DecisionsIRResponse, + UnsupportedDecisionsRequest, +) + + +class _ParallelResponse(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + mode: str + answers: Mapping[str, Mapping[str, object]] + + +class ParallelAIDecisionsConfig(BaseDecisionsConfig): + path = "/v1/decisions" + api_key_env = ("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY") + api_base_env = ("PARALLEL_AI_API_BASE",) + confidence_optional = True + openai_evidence_fields = frozenset({"citations", "decision_id", "mode"}) + + def get_default_api_base(self) -> str: + return "https://api.parallel.ai" + + def resolve_connection(self, api_base: str | None, api_key: str | None) -> tuple[str | None, str | None]: + configured_base: Final = self.resolve_api_base(None) + resolved_base: Final = api_base or configured_base + destination: Final = self._decisions_url(resolved_base) + if api_key: + return resolved_base, api_key + if destination not in {self._decisions_url(None), self._decisions_url(configured_base)}: + raise ValueError( + "Parallel Decisions requires an explicit api_key for a custom api_base unless " + "the destination is configured in PARALLEL_AI_API_BASE" + ) + return resolved_base, self.resolve_api_key(None) + + def get_complete_url(self, api_base: str, model: str) -> str: + return super().get_complete_url(api_base.rstrip("/").removesuffix("/v1/search").removesuffix(self.path), model) + + def _decisions_url(self, api_base: str | None) -> httpx.URL: + return httpx.URL(self.get_complete_url(api_base or self.get_default_api_base(), "decision")) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> Mapping[str, str]: + return { + **{ + name: value + for name, value in headers.items() + if name.lower() not in {"authorization", "x-api-key", "content-type"} + }, + **({"x-api-key": api_key} if api_key is not None else {}), + "Content-Type": "application/json", + } + + def transform_decisions_request( + self, model: str, request: DecisionsIRRequest, custom_llm_provider: str + ) -> Mapping[str, object] | UnsupportedDecisionsRequest: + if model != "decision": + return UnsupportedDecisionsRequest(reason="Parallel Decisions supports the decision model") + if len(request.questions) > 100: + return UnsupportedDecisionsRequest(reason="Parallel Decisions supports at most 100 questions") + body: Final = ir_to_systemone_request(model, request) + if isinstance(body, UnsupportedDecisionsRequest): + return body + return {"mode": "fast", **{key: value for key, value in body.items() if key != "model"}} + + def parse_response(self, payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse: + response: Final = _ParallelResponse.model_validate(payload) + return super().parse_response( + { + **response.model_dump(), + "model": "decision", + "answers": {key: answer for key, answer in response.answers.items() if answer.get("type") != "refusal"}, + }, + request, + ) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index ad64d852ef9..52345e1ff40 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1458,6 +1458,7 @@ class VertexAITokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.llms.anthropic.common_utils import context_management_with_edits from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIError as PartnerVertexAIError, @@ -1539,6 +1540,9 @@ class VertexAITokenCounter(BaseTokenCounter): system=system, tools=tools, count_options=count_options, + context_management=( + None if context_management is None else context_management_with_edits(context_management.spec) + ), ) except (PartnerVertexAIError, httpx.HTTPStatusError) as e: status_code: Final = e.response.status_code @@ -1565,12 +1569,15 @@ class VertexAITokenCounter(BaseTokenCounter): original_response=result, ) else: + from litellm.llms.anthropic.messages.context_management.inference_view import ( + InferenceRequestViewError, + count_context_edits, + inference_request_view, + ) from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, - count_context_edits, count_context_error_response, gemini_count_contents, - inference_request_view, malformed_tool_result_text_count_error, provider_error_response, unconvertible_messages_error, @@ -1584,7 +1591,7 @@ class VertexAITokenCounter(BaseTokenCounter): messages_format, count_context_edits(context_management, tools, count_tokens_params_request, model_to_use), ) - if isinstance(request_view, VertexGeminiCountContextError): + if isinstance(request_view, InferenceRequestViewError): return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context( system=request_view.system, tools=tools, model=model_to_use diff --git a/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py index 43e4ee8dc84..7334fb0520a 100644 --- a/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py +++ b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py @@ -1,4 +1,4 @@ -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from functools import partial from typing import Final, TypeVar @@ -17,16 +17,27 @@ Built: Final = TypeVar("Built") Sent: Final = TypeVar("Sent") -def cache_rebuild( +def _build_from_copies( + build: Callable[..., Built], messages: Sequence[AllMessageValues], optional_params: Mapping[str, object] +) -> Built: + return build(messages=deepcopy(list(messages)), optional_params=independent_snapshot(dict(optional_params))) + + +def cache_rebuilds( build: Callable[..., Built], messages: list[AllMessageValues], optional_params: dict -) -> Callable[[], Built] | None: - """Copies the inputs before the first transform, which pops keys from ``optional_params`` and ``messages``.""" +) -> tuple[Callable[[], Built], Callable[[], Built]] | tuple[None, None]: + """Returns one build with a fresh cache and one with the history inline, both over a snapshot of the inputs + taken before the first transform, which pops keys from ``optional_params`` and ``messages``. Each call + copies the snapshot again, since every transform consumes the inputs it is given.""" if not any(map(is_cached_message, messages)): - return None - return partial(build, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params)) + return None, None + from_snapshot: Final = partial( + _build_from_copies, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params) + ) + return partial(from_snapshot, build), partial(from_snapshot, partial(build, context_caching=False)) -def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]: +def status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]: match error: case httpx.HTTPStatusError(): return error.response.status_code, error.response.text @@ -35,7 +46,7 @@ def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tupl def _cache_rejected(body: RequestBody, error: httpx.HTTPStatusError | BaseLLMException) -> bool: - return body.get("cachedContent") is not None and names_rejected_cache(*_status_and_message(error)) + return body.get("cachedContent") is not None and names_rejected_cache(*status_and_message(error)) async def post_with_fresh_cache_async( diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 884c82fe331..4603f612926 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -425,6 +425,7 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header: str | None, extra_headers: dict | None = None, cached_content: str | None = None, + context_caching: bool = True, ) -> tuple[list[AllMessageValues], dict, str | None]: """ Receives @@ -436,7 +437,7 @@ class ContextCachingEndpoints(VertexBase): Follows - https://ai.google.dev/api/caching#request-body """ - if cached_content is not None: + if cached_content is not None or not context_caching: return messages, optional_params, cached_content cached_messages, non_cached_messages = separate_cached_messages(messages=messages) @@ -573,6 +574,7 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header: str | None, extra_headers: dict | None = None, cached_content: str | None = None, + context_caching: bool = True, ) -> tuple[list[AllMessageValues], dict, str | None]: """ Receives @@ -584,7 +586,7 @@ class ContextCachingEndpoints(VertexBase): Follows - https://ai.google.dev/api/caching#request-body """ - if cached_content is not None: + if cached_content is not None or not context_caching: return messages, optional_params, cached_content cached_messages, non_cached_messages = separate_cached_messages(messages=messages) diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index fc7d6545624..c29f2ecbcf8 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -11,10 +11,11 @@ from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_ from litellm.types.llms.anthropic import AnthropicMessagesTextParam from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ContentType -from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse if TYPE_CHECKING: from litellm.exceptions import APIError + from litellm.llms.anthropic.messages.context_management.inference_view import InferenceRequestViewError @dataclass(frozen=True, slots=True) @@ -152,7 +153,7 @@ def malformed_tool_result_text_count_error( def count_context_error_response( - error: VertexGeminiCountContextError, request_model: str, model_used: str + error: "VertexGeminiCountContextError | InferenceRequestViewError", request_model: str, model_used: str ) -> TokenCountResponse: return TokenCountResponse( total_tokens=0, @@ -178,130 +179,6 @@ def provider_error_response(error: "APIError", request_model: str, model_used: s ) -@dataclass(frozen=True, slots=True) -class InferenceRequestView: - messages: Sequence[Mapping[str, object]] - system: object | None - - -@dataclass(frozen=True, slots=True) -class CountContextEdits: - spec: JsonValue - model: str - tools: tuple[Mapping[str, object], ...] | None - drop_invalid_edits: bool = False - - -_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter( - Mapping[str, object] | tuple[Mapping[str, object], ...] -) - - -def count_context_edits( - context_management: CountContextManagement | None, - tools: Sequence[Mapping[str, object]] | None, - litellm_params: Mapping[str, object], - fallback_model: str, -) -> CountContextEdits: - model: Final = litellm_params.get("model") - return CountContextEdits( - spec=None if context_management is None else context_management.spec, - model=model if isinstance(model, str) else fallback_model, - tools=None if tools is None else tuple(tools), - drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits, - ) - - -def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None: - try: - return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec) - except ValidationError: - return None - - -def _typed_system( - system: object | None, -) -> str | tuple[AnthropicMessagesTextParam, ...] | None | VertexGeminiCountContextError: - if system is None or isinstance(system, str): - return system - try: - return _SYSTEM_BLOCKS.validate_python(system) - except ValidationError: - return _INVALID_SYSTEM - - -def _view_without_edits( - messages: Sequence[Mapping[str, object]], - system: object | None, - typed_system: str | Sequence[Mapping[str, object]] | None, -) -> InferenceRequestView: - from litellm.llms.anthropic.messages.context_management.editors.compact import ( - apply_client_compaction_block_history, - ) - - history: Final = apply_client_compaction_block_history( - messages=[dict(message) for message in messages], - system=typed_system - if typed_system is None or isinstance(typed_system, str) - else [dict(block) for block in typed_system], - ) - if history is None: - return InferenceRequestView(messages=messages, system=system) - return InferenceRequestView(messages=history.messages, system=history.system) - - -async def _anthropic_inference_view( - messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits -) -> InferenceRequestView | VertexGeminiCountContextError: - from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request - from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError - from litellm.llms.anthropic.messages.native.utils import normalize_native_messages - - typed_system: Final = _typed_system(system) - if isinstance(typed_system, VertexGeminiCountContextError): - return typed_system - normalized: Final = normalize_native_messages(messages) - spec: Final = _typed_spec(edits.spec) - try: - prepared: Final = await prepare_context_managed_request( - model=edits.model, - messages=normalized, - tools=None if edits.tools is None else [dict(tool) for tool in edits.tools], - system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system, - context_management_spec=( - None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec] - ), - litellm_metadata=None, - additional_drop_params=None, - llm_router=None, - count_only=True, - drop_invalid_edits=edits.drop_invalid_edits, - ) - except AnthropicContextManagementError as error: - if error.status_code >= 500: - return _view_without_edits(normalized, system, typed_system) - return VertexGeminiCountContextError(error_message=error.message, status_code=error.status_code) - if prepared is None: - return InferenceRequestView(messages=normalized, system=system) - return InferenceRequestView(messages=prepared.messages, system=prepared.system) - - -async def inference_request_view( - messages: Sequence[Mapping[str, object]], - system: object | None, - messages_format: TokenCountMessagesFormat, - edits: CountContextEdits, -) -> InferenceRequestView | VertexGeminiCountContextError: - """Apply the history rewrites and context edits /v1/messages runs before it converts a body for Gemini.""" - match messages_format: - case "openai": - return InferenceRequestView(messages=messages, system=system) - case "anthropic": - return await _anthropic_inference_view(messages, system, edits) - case _: - assert_never(messages_format) - - def _chat_messages( messages: Sequence[Mapping[str, object]], messages_format: TokenCountMessagesFormat, diff --git a/litellm/llms/vertex_ai/gemini/rejected_thought_signature_retry.py b/litellm/llms/vertex_ai/gemini/rejected_thought_signature_retry.py new file mode 100644 index 00000000000..32893471c11 --- /dev/null +++ b/litellm/llms/vertex_ai/gemini/rejected_thought_signature_retry.py @@ -0,0 +1,235 @@ +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from typing import Final, TypeVar + +import httpx +from pydantic import TypeAdapter, ValidationError + +from litellm import verbose_logger +from litellm.litellm_core_utils.prompt_templates.factory import get_dummy_thought_signature +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.vertex_ai import ContentType, PartType, RequestBody + +from ..context_caching.rejected_cache_retry import status_and_message + +Body: Final = TypeVar("Body", bound=Mapping[str, object]) +Sent: Final = TypeVar("Sent") + +_REJECTED_SIGNATURE_MARKER: Final = "thought signature" +_SIGNATURE_KEYS: Final = frozenset({"thoughtSignature", "thought_signature"}) +_ENTRIES: Final = TypeAdapter(list[dict[str, object]]) + + +def names_rejected_thought_signature(status_code: int, message: str) -> bool: + """A signature only validates on the backend that minted it. AI Studio answers a replayed Vertex signature + with ``400 Corrupted thought signature.`` and Vertex answers an AI Studio one with ``400 Invalid thought + signature.``""" + return status_code == 400 and _REJECTED_SIGNATURE_MARKER in message.lower() + + +def _is_rejected_signature(error: httpx.HTTPStatusError | BaseLLMException) -> bool: + return names_rejected_thought_signature(*status_and_message(error)) + + +def _entries(value: object) -> Sequence[Mapping[str, object]]: + try: + return _ENTRIES.validate_python(value) + except ValidationError: + return [] + + +def _foreign_signatures(body: Mapping[str, object], placeholder: str) -> Iterator[object]: + for content in _entries(body.get("contents")): + for part in _entries(content.get("parts")): + yield from ( + value for key, value in part.items() if key in _SIGNATURE_KEYS and value and value != placeholder + ) + + +def _foreign_signature_count(body: Mapping[str, object]) -> int: + return sum(1 for _ in _foreign_signatures(body, get_dummy_thought_signature())) + + +def _part_with_placeholder(part: PartType, placeholder: str) -> PartType: + if not part.get("thoughtSignature"): + return part + return {**part, "thoughtSignature": placeholder} + + +def _content_with_placeholders(content: ContentType, placeholder: str) -> ContentType: + return {**content, "parts": [_part_with_placeholder(part, placeholder) for part in content["parts"]]} + + +def with_placeholder_thought_signatures(body: RequestBody) -> RequestBody | None: + """Swaps every thought signature in the history for Google's documented skip-validator placeholder, or + returns ``None`` when the body carries nothing the backend could have rejected.""" + placeholder: Final = get_dummy_thought_signature() + if next(_foreign_signatures(body, placeholder), None) is None: + return None + return {**body, "contents": [_content_with_placeholders(content, placeholder) for content in body["contents"]]} + + +def _native_part_with_placeholder(part: Mapping[str, object], placeholder: str) -> Mapping[str, object]: + return {key: placeholder if key in _SIGNATURE_KEYS and value else value for key, value in part.items()} + + +def _native_content_with_placeholders(content: Mapping[str, object], placeholder: str) -> Mapping[str, object]: + parts: Final = _entries(content.get("parts")) + if not parts: + return content + return {**content, "parts": [_native_part_with_placeholder(part, placeholder) for part in parts]} + + +def with_placeholder_native_thought_signatures(body: Mapping[str, object]) -> Mapping[str, object] | None: + """The same swap on a ``generateContent`` body as the client sent it, whose parts may spell the key + ``thoughtSignature`` or ``thought_signature``.""" + placeholder: Final = get_dummy_thought_signature() + if next(_foreign_signatures(body, placeholder), None) is None: + return None + return { + **body, + "contents": [ + _native_content_with_placeholders(content, placeholder) for content in _entries(body.get("contents")) + ], + } + + +async def _inline_history_async(body: Body, uncached: Callable[[], Awaitable[Body]] | None) -> Body | None: + if uncached is None or body.get("cachedContent") is None: + return None + inline: Final = await uncached() + return inline if _foreign_signature_count(inline) > _foreign_signature_count(body) else None + + +def _inline_history_sync(body: Body, uncached: Callable[[], Body] | None) -> Body | None: + if uncached is None or body.get("cachedContent") is None: + return None + inline: Final = uncached() + return inline if _foreign_signature_count(inline) > _foreign_signature_count(body) else None + + +def _resend_body(signed: Body, placeholders: Callable[[Body], Body | None]) -> Body | None: + resend_body: Final = placeholders(signed) + if resend_body is not None: + verbose_logger.warning( + "Gemini rejected a replayed thought signature, resending with the skip_thought_signature_validator" + " placeholder in place of the replayed signatures" + ) + return resend_body + + +async def build_with_accepted_thought_signatures_async( + build: Callable[[], Awaitable[Body]], + placeholders: Callable[[Body], Body | None], + uncached: Callable[[], Awaitable[Body]] | None, +) -> Body: + """Vertex validates the signatures in a cache's contents when it creates the cache, so a rejected one there + sends the history inline with the placeholder instead.""" + try: + return await build() + except BaseLLMException as error: + if uncached is None or not _is_rejected_signature(error): + raise + inline_body: Final = _resend_body(await uncached(), placeholders) + if inline_body is None: + raise + return inline_body + + +def build_with_accepted_thought_signatures_sync( + build: Callable[[], Body], + placeholders: Callable[[Body], Body | None], + uncached: Callable[[], Body] | None, +) -> Body: + try: + return build() + except BaseLLMException as error: + if uncached is None or not _is_rejected_signature(error): + raise + inline_body: Final = _resend_body(uncached(), placeholders) + if inline_body is None: + raise + return inline_body + + +async def _post_inline_async( + body: Body, + placeholders: Callable[[Body], Body | None], + uncached: Callable[[], Awaitable[Body]] | None, + post: Callable[[Body], Awaitable[Sent]], + log_request: Callable[[Body], None], + rejection: httpx.HTTPStatusError | BaseLLMException, +) -> tuple[Body, Sent]: + inline: Final = await _inline_history_async(body, uncached) + resend_body: Final = None if inline is None else _resend_body(inline, placeholders) + if resend_body is None: + raise rejection + log_request(resend_body) + return resend_body, await post(resend_body) + + +def _post_inline_sync( + body: Body, + placeholders: Callable[[Body], Body | None], + uncached: Callable[[], Body] | None, + post: Callable[[Body], Sent], + log_request: Callable[[Body], None], + rejection: httpx.HTTPStatusError | BaseLLMException, +) -> tuple[Body, Sent]: + inline: Final = _inline_history_sync(body, uncached) + resend_body: Final = None if inline is None else _resend_body(inline, placeholders) + if resend_body is None: + raise rejection + log_request(resend_body) + return resend_body, post(resend_body) + + +async def post_with_accepted_thought_signatures_async( + body: Body, + placeholders: Callable[[Body], Body | None], + uncached: Callable[[], Awaitable[Body]] | None, + post: Callable[[Body], Awaitable[Sent]], + log_request: Callable[[Body], None], +) -> tuple[Body, Sent]: + """Innermost retry around the raw post, so a cache rebuild (which re-transforms the messages and brings + the rejected signatures back) is itself wrapped by it. The resend keeps the cache and swaps the signatures + the body carries, and sends the history inline only when that is rejected too and the cache holds one. + Returns the body the backend accepted.""" + try: + return body, await post(body) + except (httpx.HTTPStatusError, BaseLLMException) as error: + if not _is_rejected_signature(error): + raise + resend_body: Final = _resend_body(body, placeholders) + if resend_body is None: + return await _post_inline_async(body, placeholders, uncached, post, log_request, error) + log_request(resend_body) + try: + return resend_body, await post(resend_body) + except (httpx.HTTPStatusError, BaseLLMException) as error: + if not _is_rejected_signature(error): + raise + return await _post_inline_async(body, placeholders, uncached, post, log_request, error) + + +def post_with_accepted_thought_signatures_sync( + body: Body, + placeholders: Callable[[Body], Body | None], + uncached: Callable[[], Body] | None, + post: Callable[[Body], Sent], + log_request: Callable[[Body], None], +) -> tuple[Body, Sent]: + try: + return body, post(body) + except (httpx.HTTPStatusError, BaseLLMException) as error: + if not _is_rejected_signature(error): + raise + resend_body: Final = _resend_body(body, placeholders) + if resend_body is None: + return _post_inline_sync(body, placeholders, uncached, post, log_request, error) + log_request(resend_body) + try: + return resend_body, post(resend_body) + except (httpx.HTTPStatusError, BaseLLMException) as error: + if not _is_rejected_signature(error): + raise + return _post_inline_sync(body, placeholders, uncached, post, log_request, error) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index bbd4a96980a..3244e0ae15b 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1256,6 +1256,7 @@ def sync_transform_request_body( vertex_project: str | None, vertex_location: str | None, vertex_auth_header: str | None, + context_caching: bool = True, ) -> RequestBody: from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints @@ -1275,6 +1276,7 @@ def sync_transform_request_body( timeout=timeout, extra_headers=extra_headers, cached_content=optional_params.pop("cached_content", None), + context_caching=context_caching, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, vertex_project=vertex_project, @@ -1390,6 +1392,7 @@ async def async_transform_request_body( vertex_project: str | None, vertex_location: str | None, vertex_auth_header: str | None, + context_caching: bool = True, ) -> RequestBody: from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints @@ -1409,6 +1412,7 @@ async def async_transform_request_body( timeout=timeout, extra_headers=extra_headers, cached_content=optional_params.pop("cached_content", None), + context_caching=context_caching, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, vertex_project=vertex_project, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 69580687048..32d7b39a20e 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -9,7 +9,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from functools import lru_cache, partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, Union, cast, get_args import httpx from pydantic import JsonValue, TypeAdapter @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id from litellm.litellm_core_utils.prompt_templates.factory import encode_tool_call_id_with_signature from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -103,7 +104,7 @@ from ..common_utils import ( ) from ..context_caching.known_cached_contents import known_cached_contents from ..context_caching.rejected_cache_retry import ( - cache_rebuild, + cache_rebuilds, post_with_fresh_cache_async, post_with_fresh_cache_sync, ) @@ -134,6 +135,13 @@ from .long_decode import ( with_initial_prompt_tokens, with_spans_shifted, ) +from .rejected_thought_signature_retry import ( + build_with_accepted_thought_signatures_async, + build_with_accepted_thought_signatures_sync, + post_with_accepted_thought_signatures_async, + post_with_accepted_thought_signatures_sync, + with_placeholder_thought_signatures, +) from .transformation import ( async_transform_request_body, gemini_convert_messages_with_history, @@ -437,6 +445,8 @@ _NAMED_UNKNOWN_VOICE_KEYS: Final = 10 _NAMED_VOICE_KEY_CHARS: Final = 100 +_NAMED_AUDIO_FORMAT_CHARS: Final = 100 + def _gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object] | None: camel_config: Final = as_str_mapping(value.get("speechConfig")) @@ -473,6 +483,31 @@ def reject_unknown_gemini_tts_voice_keys( ) +def _named_audio_format(audio_format: object) -> str: + if isinstance(audio_format, str): + return repr(audio_format[:_NAMED_AUDIO_FORMAT_CHARS]) + return repr(audio_format)[:_NAMED_AUDIO_FORMAT_CHARS] + + +def unsupported_gemini_tts_audio_format( + model: str, audio: Mapping[str, object], drop_params: bool, llm_provider: str | None +) -> UnsupportedParamsError | None: + nested_config: Final = _gemini_speech_config_wrapper(audio) or {} + audio_format: Final = next( + (item for item in (audio.get("format"), nested_config.get("format")) if item not in (None, "pcm16")), None + ) + if audio_format is None or drops_unsupported_params(drop_params): + return None + return UnsupportedParamsError( + message=( + f"{model} does not take the audio format {_named_audio_format(audio_format)}. Set format " + "to 'pcm16' or leave it out, or set drop_params to drop it." + ), + model=model, + llm_provider=llm_provider, + ) + + def normalize_gemini_speech_config( value: Mapping[str, object], drop_unknown_fields: bool = False ) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict @@ -587,6 +622,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Note: Please make sure to modify the default parameters as required for your use case. """ + _route: ClassVar[str] = "vertex_ai" + temperature: float | None = None max_output_tokens: int | None = None top_p: float | None = None @@ -1555,15 +1592,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): languageCode: "en-US", } """ - # Validate audio format - Gemini TTS only supports pcm16 - audio_format: Final = value.get("format") - if audio_format is not None and audio_format != "pcm16": - raise ValueError( - f"Unsupported audio format for Gemini TTS models: {audio_format}. " - f"Gemini TTS models only support 'pcm16' format as they return audio data in L16 PCM format. " - f"Please set audio format to 'pcm16'." - ) - + audio_format_error: Final = unsupported_gemini_tts_audio_format( + "Gemini TTS", value, drop_unknown_fields, self.custom_llm_provider + ) + if audio_format_error is not None: + raise audio_format_error camel_speech_config: Final = value.get("speechConfig") snake_speech_config: Final = value.get("speech_config") speech_config_value: Final = ( @@ -1601,6 +1634,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _map_chat_audio_params( self, model: str, value: Mapping[str, object], drop_params: bool ) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary + audio_format_error: Final = unsupported_gemini_tts_audio_format( + model, value, drop_params, self.custom_llm_provider + ) + if audio_format_error is not None: + raise audio_format_error if _gemini_speech_config_wrapper(value) is not None: reject_unknown_gemini_tts_voice_keys(model, value, drop_params, self.custom_llm_provider) return self._map_audio_params(value, drop_unknown_fields=drops_unsupported_params(drop_params)) @@ -1639,9 +1677,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): drop_params: bool, allow_integer_enums: bool = True, ) -> dict: - self._apply_include_server_side_tool_invocations(non_default_params, optional_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, model=model, route=self._route, drop_params=drop_params + ) + self._apply_include_server_side_tool_invocations(request_params, optional_params) supports_sampling_params: Final = self.supports_sampling_params(model) - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param in ("temperature", "top_p", "top_k") and not supports_sampling_params: continue if param == "temperature": @@ -1705,7 +1746,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value elif param == "parallel_tool_calls": - tools_list = non_default_params.get("tools", non_default_params.get("functions")) + tools_list = request_params.get("tools", request_params.get("functions")) num_tools = len(tools_list) if isinstance(tools_list, list) else 0 # Gemini does not support parallel_tool_calls=False with multiple # tools. Drop the param instead of failing — Responses API clients @@ -3466,23 +3507,17 @@ class VertexLLM(VertexBase): litellm_params=litellm_params, ) - rebuild: Final = cache_rebuild( - partial( - async_transform_request_body, - **data, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_auth_header=auth_header, - ), - messages, - optional_params, - ) - request_body: Final = await async_transform_request_body( + transform: Final = partial( + async_transform_request_body, **data, vertex_project=vertex_project, vertex_location=vertex_location, vertex_auth_header=auth_header, ) + rebuild, uncached = cache_rebuilds(transform, messages, optional_params) + request_body: Final = await build_with_accepted_thought_signatures_async( + transform, with_placeholder_thought_signatures, uncached + ) ## LOGGING log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers) @@ -3511,7 +3546,14 @@ class VertexLLM(VertexBase): async def post(body: RequestBody) -> ModelResponseIterator: return await open_stream(client=client, data=json.dumps(body), request_body=body) - return (await post_with_fresh_cache_async(request_body, rebuild, post, log_request))[1] + accepted: Final = partial( + post_with_accepted_thought_signatures_async, + placeholders=with_placeholder_thought_signatures, + uncached=uncached, + post=post, + log_request=log_request, + ) + return (await post_with_fresh_cache_async(request_body, rebuild, accepted, log_request))[1][1] streaming_response: Final = CustomStreamWrapper( completion_stream=None, @@ -3578,23 +3620,17 @@ class VertexLLM(VertexBase): litellm_params=litellm_params, ) - rebuild: Final = cache_rebuild( - partial( - async_transform_request_body, - **data, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_auth_header=auth_header, - ), - messages, - optional_params, - ) - request_body: Final = await async_transform_request_body( + transform: Final = partial( + async_transform_request_body, **data, vertex_project=vertex_project, vertex_location=vertex_location, vertex_auth_header=auth_header, ) + rebuild, uncached = cache_rebuilds(transform, messages, optional_params) + request_body: Final = await build_with_accepted_thought_signatures_async( + transform, with_placeholder_thought_signatures, uncached + ) _async_client_params: Final = {} if timeout: @@ -3635,7 +3671,14 @@ class VertexLLM(VertexBase): ) return response - sent_body, first_response = await post_with_fresh_cache_async(request_body, rebuild, post, log_request) + accepted: Final = partial( + post_with_accepted_thought_signatures_async, + placeholders=with_placeholder_thought_signatures, + uncached=uncached, + post=post, + log_request=log_request, + ) + _, (sent_body, first_response) = await post_with_fresh_cache_async(request_body, rebuild, accepted, log_request) return _transform_long_decode_hops( model=model, hops=await collect_hops_async(parse_body(first_response), sent_body, post), @@ -3767,23 +3810,17 @@ class VertexLLM(VertexBase): ) ## TRANSFORMATION ## - rebuild: Final = cache_rebuild( - partial( - sync_transform_request_body, - **transform_request_params, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_auth_header=auth_header, - ), - messages, - optional_params, - ) - data: Final = sync_transform_request_body( + transform: Final = partial( + sync_transform_request_body, **transform_request_params, vertex_project=vertex_project, vertex_location=vertex_location, vertex_auth_header=auth_header, ) + rebuild, uncached = cache_rebuilds(transform, messages, optional_params) + data: Final = build_with_accepted_thought_signatures_sync( + transform, with_placeholder_thought_signatures, uncached + ) ## LOGGING log_request: Final = partial(_log_generate_request, logging_obj, messages, url, headers) @@ -3811,7 +3848,14 @@ class VertexLLM(VertexBase): def post(body: RequestBody) -> ModelResponseIterator: return open_stream(client=client, data=json.dumps(body), request_body=body) - return post_with_fresh_cache_sync(data, rebuild, post, log_request)[1] + accepted: Final = partial( + post_with_accepted_thought_signatures_sync, + placeholders=with_placeholder_thought_signatures, + uncached=uncached, + post=post, + log_request=log_request, + ) + return post_with_fresh_cache_sync(data, rebuild, accepted, log_request)[1][1] streaming_response: Final = CustomStreamWrapper( completion_stream=None, @@ -3855,7 +3899,14 @@ class VertexLLM(VertexBase): ) return response - sent_body, first_response = post_with_fresh_cache_sync(data, rebuild, post, log_request) + accepted: Final = partial( + post_with_accepted_thought_signatures_sync, + placeholders=with_placeholder_thought_signatures, + uncached=uncached, + post=post, + log_request=log_request, + ) + _, (sent_body, first_response) = post_with_fresh_cache_sync(data, rebuild, accepted, log_request) return _transform_long_decode_hops( model=model, hops=collect_hops_sync(parse_body(first_response), sent_body, post), diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index e341c8f5f48..8d39ae76456 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -9,12 +9,14 @@ their respective publisher-specific count-tokens endpoints. from collections.abc import Mapping, Sequence from typing import Any, Final -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter +from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, filter_and_transform_beta_headers from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import ( OBJECT_LIST_ADAPTER, + context_management_betas, count_tokens_feature_betas, is_claude_model_target, tools_without_custom_field, @@ -31,6 +33,7 @@ _COUNT_TOKENS_PUBLISHERS: Final = ( _REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object]) +_CONTEXT_MANAGEMENT: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) def count_tokens_request_body( @@ -55,8 +58,14 @@ def count_tokens_publisher(model: str) -> str | None: def count_tokens_headers(access_token: str, request_data: Mapping[str, object]) -> Mapping[str, str]: authorization: Final = {"Authorization": f"Bearer {access_token}"} - betas: Final = count_tokens_feature_betas( - "vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"]) + betas: Final = ( + *count_tokens_feature_betas( + "vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"]) + ), + *filter_and_transform_beta_headers( + list(context_management_betas(_CONTEXT_MANAGEMENT.validate_python(request_data.get("context_management")))), + "vertex_ai", + ), ) return {**authorization, "anthropic-beta": ",".join(betas)} if betas else authorization @@ -193,6 +202,8 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): api_base=litellm_params.get("api_base"), ) + if request_data.get("context_management") is not None: + await aload_beta_headers_config() headers: Final = dict(count_tokens_headers(access_token, request_data)) # Get async HTTP client diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index a3a3427670d..c298ca404e2 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -268,6 +268,7 @@ class VertexAIPartnerModels(VertexBase): system: object | None = None, tools: list[dict[str, object]] | None = None, count_options: Mapping[str, object] = MappingProxyType({}), + context_management: object = None, ): """ Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.) @@ -302,7 +303,15 @@ class VertexAIPartnerModels(VertexBase): "model": model, "messages": messages, **MappingProxyType( - {key: value for key, value in (("system", system), ("tools", tools)) if value is not None} + { + key: value + for key, value in ( + ("system", system), + ("tools", tools), + ("context_management", context_management), + ) + if value is not None + } ), **count_options, } diff --git a/litellm/llms/vllm/common_utils.py b/litellm/llms/vllm/common_utils.py index 014fabd5187..41db96914f5 100644 --- a/litellm/llms/vllm/common_utils.py +++ b/litellm/llms/vllm/common_utils.py @@ -2,10 +2,12 @@ from collections.abc import Mapping from typing import Final import httpx +from pydantic import TypeAdapter -import litellm +from litellm.constants import NON_LLM_CONNECTION_TIMEOUT, PLACEHOLDER_API_KEY from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.custom_httpx.http_handler import get_httpx_client from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.utils import add_path_to_api_base @@ -37,17 +39,31 @@ class VLLMError(BaseLLMException): ) +_EXTRA_HEADERS: Final = TypeAdapter(Mapping[str, str]) + + +def _auth_headers(api_key: str | None) -> Mapping[str, str]: + if api_key is None or api_key == PLACEHOLDER_API_KEY: + return {} + return {"Authorization": f"Bearer {api_key}"} + + +def _str_param(litellm_params: Mapping[str, object], key: str) -> str | None: + value: Final = litellm_params.get(key) + return value if isinstance(value, str) else None + + class VLLMModelInfo(BaseLLMModelInfo): def validate_environment( self, - headers: dict, + headers: dict[str, str], model: str, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: upstream_key: Final = self.upstream_api_key(api_key, api_base) return {**headers, **self.upstream_auth_headers(upstream_key)} if upstream_key else headers @@ -58,7 +74,7 @@ class VLLMModelInfo(BaseLLMModelInfo): return {"Authorization": f"Bearer {api_key}", "x-api-key": api_key} @staticmethod - def get_api_base(api_base: str | None = None) -> str | None: + def get_api_base(api_base: str | None = None) -> str: api_base = api_base or get_secret_str("VLLM_API_BASE") if api_base is None: raise ValueError( @@ -74,25 +90,55 @@ class VLLMModelInfo(BaseLLMModelInfo): def get_base_model(model: str) -> str | None: return model - def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: - api_base = VLLMModelInfo.get_api_base(api_base) - api_key = VLLMModelInfo.get_api_key(api_key) - endpoint: Final = "/v1/models" - if api_base is None or api_key is None: - raise ValueError( - "VLLM_API_BASE or VLLM_API_KEY is not set. Please set the environment variable, to query VLLM's `/models` endpoint." - ) - - url: Final = add_path_to_api_base(api_base, endpoint) - response: Final = litellm.module_level_client.get( - url=url, - ) - - response.raise_for_status() - - models: Final = response.json()["data"] - - return [model["id"] for model in models] + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: mirrors override contract + return [] def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return VLLMError(status_code=status_code, message=error_message, headers=headers) + + +class HostedVLLMModelInfo(VLLMModelInfo): + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE") + if resolved_api_base is None: + raise ValueError("HOSTED_VLLM_API_BASE is not set.") + return resolved_api_base + + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return api_key or get_secret_str("HOSTED_VLLM_API_KEY") + + def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: + return self._list_models(api_key=api_key, api_base=api_base, extra_headers={}, ssl_verify=None) + + def discover_models( + self, litellm_params: Mapping[str, object] | None = None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + params: Final = litellm_params or {} + ssl_verify: Final = params.get("ssl_verify") + return self._list_models( + api_key=_str_param(params, "api_key"), + api_base=_str_param(params, "api_base"), + extra_headers=_EXTRA_HEADERS.validate_python(params.get("extra_headers") or {}), + ssl_verify=ssl_verify if isinstance(ssl_verify, (bool, str)) else None, + ) + + def _list_models( + self, + *, + api_key: str | None, + api_base: str | None, + extra_headers: Mapping[str, str], + ssl_verify: bool | str | None, + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + url: Final = add_path_to_api_base(self.get_api_base(api_base), "/v1/models") + response: Final = get_httpx_client(params={"ssl_verify": ssl_verify}).get( + url=url, + headers={**extra_headers, **_auth_headers(self.get_api_key(api_key))}, + timeout=NON_LLM_CONNECTION_TIMEOUT, + ) + response.raise_for_status() + return [model["id"] for model in response.json()["data"]] diff --git a/litellm/llms/vllm/passthrough/transformation.py b/litellm/llms/vllm/passthrough/transformation.py index 2b253e7fa38..10d740e49e0 100644 --- a/litellm/llms/vllm/passthrough/transformation.py +++ b/litellm/llms/vllm/passthrough/transformation.py @@ -67,9 +67,6 @@ class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig): ) -> tuple["URL", str]: base_target_url: Final = self.get_api_base(api_base) - if base_target_url is None: - raise Exception("VLLM api base not found") - return ( self.format_url(endpoint_under_api_base(endpoint, base_target_url), base_target_url, request_query_params), base_target_url, diff --git a/litellm/main.py b/litellm/main.py index 78767b03eb7..45177f4223e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -37,6 +37,8 @@ if TYPE_CHECKING: import dotenv import httpx import openai +from openai import AsyncStream, Stream +from openai.types.audio import TranscriptionStreamEvent from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import assert_never, overload @@ -70,6 +72,9 @@ from litellm.constants import ( ) from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + raise_if_allowed_tools_tool_choice_unsupported, +) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, @@ -263,7 +268,7 @@ from .llms.codestral.completion.handler import CodestralTextCompletion from .llms.cohere.embed import handler as cohere_embed from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from .llms.custom_llm import CustomLLM, custom_chat_llm_router +from .llms.custom_llm import custom_chat_llm_router from .llms.databricks.common_utils import DatabricksBase from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha @@ -750,6 +755,7 @@ async def acompletion( from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, + routed_model_info, ) from litellm.types.llms.openai import AllMessageValues @@ -761,6 +767,7 @@ async def acompletion( tools=tools, enable_prompt_caching=cast(bool | None, kwargs.get("enable_prompt_caching")), # cast-ok: untyped kwargs api_base=base_url or kwargs.get("api_base"), + model_info=routed_model_info(caller_kwargs.get("model_info")), ) prompt_id: Final = cast(str | None, kwargs.get("prompt_id", None)) # cast-ok: read from untyped kwargs @@ -998,7 +1005,7 @@ def _model_response_from_cached_dict(cached_response_dict: Mapping[str, object]) def _transcription_response_from_cached_dict(cached_response_dict: Mapping[str, object]) -> TranscriptionResponse: - return TranscriptionResponse(**cached_response_dict) + return TranscriptionResponse.model_validate(cached_response_dict) def _text_completion_response_from_cached_dict(cached_response_dict: Mapping[str, object]) -> TextCompletionResponse: @@ -5334,10 +5341,7 @@ def _complete_custom_providers( stream: Final = ctx.stream timeout: Final = ctx.timeout - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) @@ -5825,6 +5829,7 @@ def completion( from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, + routed_model_info, ) from litellm.types.llms.openai import AllMessageValues @@ -5836,6 +5841,7 @@ def completion( tools=tools, enable_prompt_caching=cast(bool | None, kwargs.get("enable_prompt_caching")), # cast-ok: untyped kwargs api_base=base_url or kwargs.get("api_base"), + model_info=routed_model_info(cast(object, model_info)), # cast-ok: untyped kwargs ) if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( @@ -5956,6 +5962,14 @@ def completion( if not is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): tools = _drop_anthropic_only_tool_keys(tools=tools) + raise_if_allowed_tools_tool_choice_unsupported( + tool_choice=cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic + object, tool_choice + ), + custom_llm_provider=custom_llm_provider, + model=model, + ) + if provider_specific_header is not None: headers.update( ProviderSpecificHeaderUtils.get_provider_specific_headers( @@ -7392,6 +7406,7 @@ def embedding( aembedding=aembedding, litellm_params=litellm_params_dict, headers=headers, + timeout=timeout, ) elif custom_llm_provider == "bedrock": if isinstance(input, str): @@ -7900,10 +7915,7 @@ def embedding( litellm_params={}, ) elif custom_llm_provider in litellm._custom_providers: - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) @@ -8495,7 +8507,10 @@ async def amoderation( @client -async def atranscription(*args, **kwargs) -> TranscriptionResponse: +async def atranscription( + *args, # noqa: ANN002 # public SDK wrapper preserves positional call compatibility + **kwargs, # noqa: ANN003 # kwargs-ok: public SDK wrapper preserves keyword call compatibility +) -> TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]: """ Calls openai + azure whisper endpoints. @@ -8539,6 +8554,14 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: else: # Call the synchronous function using run_in_executor response = await loop.run_in_executor(None, func_with_context) + if isinstance(response, AsyncStream): + if file is not None: + calculated_duration = calculate_request_duration(file) + if calculated_duration is not None: + response._litellm_audio_duration = ( # pyright: ignore[reportAttributeAccessIssue] # SDK omits it + calculated_duration + ) + return response if not isinstance(response, TranscriptionResponse): raise ValueError( f"Invalid response from transcription provider, expected TranscriptionResponse, but got {type(response)}" @@ -8551,9 +8574,9 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: if response is not None and not isinstance(response, Coroutine) and file is not None: existing_duration: Final = getattr(response, "duration", None) if existing_duration is None: - calculated_duration: Final = calculate_request_duration(file) - if calculated_duration is not None: - response.hidden_params["audio_transcription_duration"] = calculated_duration + sync_calculated_duration: Final = calculate_request_duration(file) + if sync_calculated_duration is not None: + response.set_audio_transcription_duration(sync_calculated_duration) return response except Exception as e: @@ -8567,6 +8590,42 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: ) +def _validate_gpt_transcription_request( + model: str, + pricing_model: str, + custom_llm_provider: str, + response_format: str | None, + drop_params: object, +) -> str | None: + model_cost_key: Final = next( + ( + key + for key in (f"{custom_llm_provider}/{pricing_model}", pricing_model) + if (as_str_mapping(litellm.model_cost.get(key)) or {}).get("litellm_provider") == custom_llm_provider + ), + None, + ) + model_info: Final = ( + get_model_info(model=pricing_model, custom_llm_provider=custom_llm_provider) + if model_cost_key is not None + else None + ) + provider_specific_entry: Final = model_info.get("provider_specific_entry") if model_info is not None else None + if ( + provider_specific_entry is None + or provider_specific_entry.get("transcription_json_or_text_only") != 1 + or response_format in (None, "json", "text") + ): + return response_format + if litellm.drop_params is True or normalize_drop_params(drop_params) is True: + return None + raise litellm.UnsupportedParamsError( + message=f"{model} only supports response_format 'json' or 'text'", + model=model, + llm_provider=custom_llm_provider, + ) + + @client def transcription( model: str, @@ -8586,8 +8645,15 @@ def transcription( max_retries: int | None = None, custom_llm_provider=None, base_url: str | None = None, + languages: Sequence[str] | None = None, + keywords: Sequence[str] | None = None, + stream: bool | None = None, **kwargs, -) -> TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]: +) -> ( + TranscriptionResponse + | Stream[TranscriptionStreamEvent] + | Coroutine[Any, Any, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] +): """ Calls openai + azure whisper endpoints. @@ -8625,19 +8691,36 @@ def transcription( api_key = dynamic_api_key if dynamic_api_key is not None else api_key + litellm_params_dict: Final = get_litellm_params(**kwargs) + configured_base_model: Final = litellm_params_dict.get("base_model") + pricing_model: Final = ( + configured_base_model.removeprefix(f"{custom_llm_provider}/") + if isinstance(configured_base_model, str) + else model + ) + + supported_response_format: Final = _validate_gpt_transcription_request( + model=model, + pricing_model=pricing_model, + custom_llm_provider=custom_llm_provider, + response_format=response_format, + drop_params=kwargs.get("drop_params"), + ) + optional_params: Final = get_optional_params_transcription( model=model, language=language, + languages=languages, + keywords=keywords, prompt=prompt, - response_format=response_format, + response_format=supported_response_format, timestamp_granularities=timestamp_granularities, temperature=temperature, + stream=stream, custom_llm_provider=custom_llm_provider, **non_default_params, ) - litellm_params_dict: Final = get_litellm_params(**kwargs) - litellm_logging_obj.update_environment_variables( model=model, user=user, @@ -8654,11 +8737,17 @@ def transcription( custom_llm_provider=custom_llm_provider, ) - response: TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] + response: ( + TranscriptionResponse + | Stream[TranscriptionStreamEvent] + | AsyncStream[TranscriptionStreamEvent] + | Coroutine[Any, Any, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] + ) provider_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( model=model, provider=LlmProviders(custom_llm_provider), + base_model=pricing_model, ) uses_openai_transport: Final = custom_llm_provider in OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS and not ( @@ -8669,7 +8758,7 @@ def transcription( # azure configs api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") - api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") + azure_api_version: Final = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") azure_ad_token: Final = kwargs.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") @@ -8688,7 +8777,7 @@ def transcription( logging_obj=litellm_logging_obj, api_base=api_base, api_key=api_key, - api_version=api_version, + api_version=azure_api_version, azure_ad_token=azure_ad_token, max_retries=max_retries, litellm_params=litellm_params_dict, @@ -8825,15 +8914,17 @@ def transcription( # Store duration in _hidden_params for cost calculation without # exposing it in the response body (see sync path comment above). - if not isinstance(response, Coroutine): - existing_duration: Final = getattr(response, "duration", None) - if existing_duration is None: - calculated_duration: Final = calculate_request_duration(file) + if not isinstance(response, Coroutine) and ( + isinstance(response, (Stream, AsyncStream)) or getattr(response, "duration", None) is None + ): + calculated_duration: Final = calculate_request_duration(file) + if isinstance(response, (Stream, AsyncStream)): if calculated_duration is not None: - response_hidden_params: Final = cast( # cast-ok: preserve dynamic mapping behavior - dict[str, object], getattr(response, HIDDEN_PARAMS_ATTR) + response._litellm_audio_duration = ( # pyright: ignore[reportAttributeAccessIssue] # SDK omits it + calculated_duration ) - response_hidden_params["audio_transcription_duration"] = calculated_duration + elif calculated_duration is not None: + response.set_audio_transcription_duration(calculated_duration) return response diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 35136c41dc2..4c94f26c6c1 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -75,9 +75,7 @@ def _public_request( return native_call(legacy, args, _native_kwargs(fields, kwargs)) -def _resolved(request: NativeCall) -> tuple[str, str | None]: - model: Final = str(request.resolved["model"]) - provider: Final = optional_str(request.resolved.get("custom_llm_provider")) +def _resolved_model(model: str, provider: str | None) -> tuple[str, str | None]: try: resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider) except BadRequestError: @@ -85,21 +83,43 @@ def _resolved(request: NativeCall) -> tuple[str, str | None]: return resolved_model, resolved_provider -def _target(request: NativeCall, model: str, provider: str | None) -> Target: - requested_model: Final = str(request.resolved["model"]) - model_info: Final = request.resolved.get("model_info") - api_base: Final = request.resolved.get("api_base") - native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup - model, provider, model_info, api_base if isinstance(api_base, str) else None - ) +def _resolved(request: NativeCall) -> tuple[str, str | None]: + return _resolved_model(str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider"))) + + +def _route_target( + requested_model: str, + model: str, + provider: str | None, + model_info: object, + api_base: str | None, + compacting: bool, +) -> Target: + native_config: Final = main._native_messages_config(model, provider, model_info, api_base) # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup if native_config is not None: return Target.MESSAGES - compacting: Final = request.resolved.get("compaction") is not None if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice return Target.RESPONSES return Target.CHAT_COMPLETIONS +def _target(request: NativeCall, model: str, provider: str | None) -> Target: + api_base: Final = request.resolved.get("api_base") + return _route_target( + str(request.resolved["model"]), + model, + provider, + request.resolved.get("model_info"), + api_base if isinstance(api_base, str) else None, + compacting=request.resolved.get("compaction") is not None, + ) + + +def messages_target(model: str, custom_llm_provider: str | None, model_info: object, api_base: str | None) -> Target: + resolved_model, resolved_provider = _resolved_model(model, custom_llm_provider) + return _route_target(model, resolved_model, resolved_provider, model_info, api_base, compacting=False) + + _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **detectors.COMMON, @@ -119,6 +139,7 @@ _ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType( { **_FEATURES, Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook, + Feature.FOUNDRY_BETA_ROW: route_detectors.foundry_beta_row, Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception, Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted, Feature.STRICT_TOOLS_RETRY: route_detectors.strict_tools_retry, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e1fdfc89a9f..974f1bb0c8b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1,4 +1,168 @@ { + "azure/gpt-realtime-2.1-2026-07-07": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 2.4e-05, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "cache_read_input_image_token_cost": 5e-07, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-realtime-2.1-mini-2026-07-07": { + "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, + "input_cost_per_audio_token": 1e-05, + "input_cost_per_image_token": 8e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "cache_read_input_image_token_cost": 8e-08, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-realtime-translate-2026-05-06": { + "input_cost_per_second": 0.0005666666666666667, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/a-new-chapter-for-realtime-ai-reasoning-translation-and-real-time-transcription/4517124", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_native_streaming": true + }, + "azure/gpt-realtime-translate-2026-05-07": { + "input_cost_per_second": 0.0005666666666666667, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "source": "https://ai.azure.com/catalog/models/gpt-realtime-translate", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_native_streaming": true + }, + "azure/gpt-realtime-whisper-2026-05-06": { + "input_cost_per_second": 0.0002833333333333333, + "litellm_provider": "azure", + "mode": "audio_transcription", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/gpt-realtime-whisper", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "max_input_tokens": 32000, + "max_output_tokens": 4096 + }, + "azure/gpt-realtime-whisper-2026-05-07": { + "input_cost_per_second": 0.0002833333333333333, + "litellm_provider": "azure", + "mode": "audio_transcription", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/gpt-realtime-whisper", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "max_input_tokens": 32000, + "max_output_tokens": 4096 + }, "sample_spec": { "code_interpreter_cost_per_session": 0.0, "computer_use_input_cost_per_1k_tokens": 0.0, @@ -6424,6 +6588,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6458,6 +6625,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6491,6 +6661,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6525,6 +6698,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6547,6 +6723,7 @@ "azure/gpt-realtime-2.1": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_image_token_cost": 5e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-06-25", "input_cost_per_audio_token": 3.2e-05, @@ -6556,6 +6733,9 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", @@ -6575,12 +6755,15 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true }, "azure/gpt-realtime-2.1-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_image_token_cost": 8e-08, "cache_read_input_token_cost": 6e-08, "deprecation_date": "2027-06-25", "input_cost_per_audio_token": 1e-05, @@ -6590,6 +6773,9 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", @@ -6609,6 +6795,8 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true }, @@ -6626,6 +6814,9 @@ "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6659,6 +6850,9 @@ "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6791,6 +6985,9 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "mode": "audio_transcription", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/gpt-realtime-whisper", "supported_endpoints": [ "/v1/realtime", @@ -13478,6 +13675,10 @@ "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", + "token_count_framing": 149, + "token_count_tool_framing": 302, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 100, "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", @@ -13831,7 +14032,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -13845,7 +14046,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/ap-northeast-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14002,7 +14203,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14016,7 +14217,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/ap-south-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14140,12 +14341,12 @@ "input_cost_per_token": 7.4e-07, "input_cost_per_token_flex": 3.7e-07, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, "output_cost_per_token_flex": 1.11e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.html", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md", "supports_audio_input": false, "supports_function_calling": true, "supports_native_structured_output": true, @@ -14250,7 +14451,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14264,7 +14465,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/eu-north-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14786,7 +14987,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14800,7 +15001,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/sa-east-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15035,7 +15236,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15049,7 +15250,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-east-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15156,7 +15357,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15170,7 +15371,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-east-2/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15552,6 +15753,7 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { + "deprecation_date": "2027-03-30", "input_cost_per_token": 2.65e-06, "litellm_provider": "bedrock", "max_input_tokens": 8000, @@ -15562,6 +15764,7 @@ "supports_pdf_input": true }, "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": { + "deprecation_date": "2027-03-30", "input_cost_per_token": 3e-07, "litellm_provider": "bedrock", "max_input_tokens": 8000, @@ -15708,7 +15911,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15722,7 +15925,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-west-2/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -23519,12 +23722,12 @@ ] }, "deepgram/streaming/nova-3": { - "input_cost_per_second": 8e-05, + "input_cost_per_second": 0.00012833, "litellm_provider": "deepgram", "metadata": { - "calculation": "$0.0048/60 seconds = $0.00008000 per second", + "calculation": "$0.0077/60 seconds = $0.00012833 per second", "note": "Nova-3 monolingual streaming, pay as you go", - "original_pricing_per_minute": 0.0048 + "original_pricing_per_minute": 0.0077 }, "mode": "audio_transcription", "output_cost_per_second": 0.0, @@ -23533,13 +23736,28 @@ "/v1/listen" ] }, - "deepgram/streaming/nova-3-multilingual": { - "input_cost_per_second": 9.667e-05, + "deepgram/streaming/nova-3-medical": { + "input_cost_per_second": 0.00012833, "litellm_provider": "deepgram", "metadata": { - "calculation": "$0.0058/60 seconds = $0.00009667 per second", + "calculation": "$0.0077/60 seconds = $0.00012833 per second", + "note": "Nova-3 Medical streaming, launch price from Deepgram's 2025-03-03 announcement. deepgram.com/pricing does not list Nova-3 Medical, and this matches its current Nova-3 monolingual streaming pay as you go rate", + "original_pricing_per_minute": 0.0077 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://www.businesswire.com/news/home/20250303470439/en", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/nova-3-multilingual": { + "input_cost_per_second": 0.00015333, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0092/60 seconds = $0.00015333 per second", "note": "Nova-3 multilingual (language=multi) streaming, pay as you go", - "original_pricing_per_minute": 0.0058 + "original_pricing_per_minute": 0.0092 }, "mode": "audio_transcription", "output_cost_per_second": 0.0, @@ -24519,7 +24737,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock_converse", - "max_input_tokens": 164000, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -24530,7 +24748,7 @@ "supports_reasoning": true, "supports_native_structured_output": true, "supports_tool_choice": true, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -29217,9 +29435,9 @@ }, "gemini-3-pro-image": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 3.6e-07, "deprecation_date": "2027-05-28", @@ -31203,6 +31421,7 @@ "cache_read_input_token_cost_batches": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, + "deprecation_date": "2026-12-01", "input_cost_per_audio_token": 1e-06, "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 3e-07, @@ -31786,6 +32005,7 @@ "output_cost_per_token_flex": 5e-06, "supports_url_context": true, "cache_read_input_token_cost_above_200k_tokens_batches": 2.5e-07, + "deprecation_date": "2026-12-01", "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 }, @@ -31889,7 +32109,7 @@ }, "gemini/gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_batches": 2e-08, + "cache_read_input_token_cost_batches": 1.5e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, @@ -34788,6 +35008,7 @@ "gpt-4.1": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_priority": 8.75e-07, + "deprecation_date": "2027-01-14", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_priority": 3.5e-06, @@ -34873,6 +35094,7 @@ "gpt-4.1-mini": { "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_priority": 1.75e-07, + "deprecation_date": "2027-01-14", "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, "input_cost_per_token_priority": 7e-07, @@ -35032,6 +35254,7 @@ "gpt-4o": { "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_priority": 2.125e-06, + "deprecation_date": "2027-01-14", "input_cost_per_token": 2.5e-06, "input_cost_per_token_batches": 1.25e-06, "input_cost_per_token_priority": 4.25e-06, @@ -35376,6 +35599,7 @@ "gpt-4o-mini": { "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.25e-07, + "deprecation_date": "2027-01-14", "input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_priority": 2.5e-07, @@ -39044,6 +39268,9 @@ "max_input_tokens": 128000, "max_output_tokens": 32000, "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, "source": "https://developers.openai.com/api/docs/pricing", @@ -39099,6 +39326,8 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true }, @@ -39135,6 +39364,8 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true }, @@ -43081,18 +43312,18 @@ "mode": "embedding" }, "mistral/mistral-large-4": { - "cache_read_input_token_cost": 1.4e-07, - "input_cost_per_token": 1.36e-06, + "cache_read_input_token_cost": 7e-08, + "input_cost_per_token": 7e-07, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 4.18e-06, + "output_cost_per_token": 2.1e-06, "reasoning_effort_levels": [ "none", "high" ], - "source": "https://docs.mistral.ai/models/mistral-large-4", + "source": "https://docs.mistral.ai/models/mistral-large-4-0", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -43102,18 +43333,18 @@ "supports_vision": true }, "mistral/mistral-large-4-0": { - "cache_read_input_token_cost": 1.4e-07, - "input_cost_per_token": 1.36e-06, + "cache_read_input_token_cost": 7e-08, + "input_cost_per_token": 7e-07, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 4.18e-06, + "output_cost_per_token": 2.1e-06, "reasoning_effort_levels": [ "none", "high" ], - "source": "https://docs.mistral.ai/models/mistral-large-4", + "source": "https://docs.mistral.ai/models/mistral-large-4-0", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -46387,6 +46618,19 @@ "supports_response_schema": true, "supports_tool_choice": false }, + "parallel_ai/decision": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "parallel_ai", + "max_input_tokens": 32000, + "mode": "evaluation", + "source": "https://docs.parallel.ai/getting-started/pricing#decisions-api", + "supported_endpoints": [ + "/v1/decisions", + "/v1/systemone" + ] + }, "parallel_ai/search": { "input_cost_per_query": 0.005, "litellm_provider": "parallel_ai", @@ -49614,6 +49858,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -49658,6 +49903,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -52547,9 +52793,9 @@ }, "vertex_ai/gemini-3-pro-image": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 3.6e-07, "deprecation_date": "2027-05-28", @@ -53957,23 +54203,23 @@ "supports_vision": true }, "wandb/meta-llama/Llama-3.1-8B-Instruct": { - "max_input_tokens": 131000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "input_cost_per_token": 2.2e-07, "output_cost_per_token": 2.2e-07, "litellm_provider": "wandb", "mode": "chat", - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/deepseek-ai/DeepSeek-V3.1": { "supports_reasoning": true, - "max_input_tokens": 161000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "input_cost_per_token": 5.5e-07, "output_cost_per_token": 1.65e-06, "litellm_provider": "wandb", "mode": "chat", - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/meta-llama/Llama-3.3-70B-Instruct": { "max_input_tokens": 128000, @@ -54457,6 +54703,8 @@ "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", + "token_count_framing": 185, + "token_count_tool_framing": 214, "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", @@ -59875,6 +60123,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html", "use_openai_responses_path": true, "supported_endpoints": [ @@ -59896,7 +60145,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { "input_cost_per_token": 4.4e-06, @@ -59911,6 +60161,7 @@ "max_input_tokens": 1050000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/responses" @@ -59930,7 +60181,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -60336,6 +60588,7 @@ "mode": "responses", "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, "output_cost_per_token_ultrafast": 0.00033, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60359,7 +60612,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-6-sol": { "input_cost_per_token": 2.2e-06, @@ -60374,6 +60628,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60397,7 +60652,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-6-luna": { "input_cost_per_token": 1.1e-07, @@ -60412,6 +60668,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60435,7 +60692,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", + "supports_web_search": true }, "us.openai.gpt-6-astra": { "supports_bedrock_runtime_chat_completions_response_format": true, @@ -62630,7 +62888,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 23, + "token_count_tool_framing": 12 }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, @@ -62736,7 +62996,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 76, + "token_count_tool_framing": 75 }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, @@ -63349,6 +63611,10 @@ "input_cost_per_second": 7.5e-05, "litellm_provider": "openai", "mode": "audio_transcription", + "provider_specific_entry": { + "transcription_json_or_text_only": 1, + "transcription_usage_duration": 1 + }, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/transcriptions", @@ -63361,7 +63627,8 @@ "supported_output_modalities": [ "text" ], - "supports_audio_input": true + "supports_audio_input": true, + "supports_native_streaming": true }, "gpt-live-transcribe": { "input_cost_per_second": 0.000283333333333, @@ -63379,7 +63646,8 @@ "supported_output_modalities": [ "text" ], - "supports_audio_input": true + "supports_audio_input": true, + "supports_native_streaming": true }, "gpt-live-1": { "input_cost_per_second": 0.000833333333333, @@ -63413,7 +63681,8 @@ "audio" ], "supports_audio_input": true, - "supports_audio_output": true + "supports_audio_output": true, + "supports_native_streaming": true }, "claude-mythos-5": { "supports_anthropic_compaction": true, @@ -64022,6 +64291,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", @@ -64029,6 +64299,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64048,6 +64319,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64090,6 +64362,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64120,6 +64393,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", @@ -64127,6 +64401,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64146,6 +64421,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64220,6 +64496,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64240,6 +64517,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64263,6 +64541,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64338,6 +64617,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64374,6 +64654,7 @@ "output_cost_per_token_priority": 6.375e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_vision": true }, @@ -64410,6 +64691,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64480,6 +64762,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64503,6 +64786,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67296,6 +67580,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67314,6 +67599,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67332,6 +67618,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67350,6 +67637,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67365,6 +67653,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67407,6 +67696,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67425,6 +67715,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67443,6 +67734,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67460,6 +67752,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67478,6 +67771,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67496,6 +67790,7 @@ "output_cost_per_token_priority": 4.05e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -68281,6 +68576,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -68326,6 +68622,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -68554,6 +68851,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -68594,6 +68892,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -68633,7 +68932,9 @@ "cache_read_input_token_cost_above_272k_tokens": 5.4e-07, "output_cost_per_token": 1.62e-05, "output_cost_per_token_above_272k_tokens": 2.43e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { "litellm_provider": "bedrock_mantle", @@ -68668,7 +68969,9 @@ "cache_read_input_token_cost_above_272k_tokens": 5.4e-08, "output_cost_per_token": 1.62e-06, "output_cost_per_token_above_272k_tokens": 2.43e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { "litellm_provider": "bedrock_mantle", @@ -68697,7 +69000,9 @@ "input_cost_per_token": 3.375e-06, "cache_read_input_token_cost": 3.375e-07, "output_cost_per_token": 2.025e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/xai.grok-4.3": { "use_openai_responses_path": true, @@ -71248,6 +71553,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/realtime" @@ -71305,7 +71613,11 @@ "supported_output_modalities": [ "text" ], - "supports_audio_input": true + "supports_audio_input": true, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "supports_native_streaming": true }, "azure/gpt-transcribe": { "deprecation_date": "2028-02-01", @@ -71315,7 +71627,8 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/audio/transcriptions", - "/v1/realtime/transcription_sessions" + "/v1/realtime/transcription_sessions", + "/v1/realtime" ], "supported_modalities": [ "text", @@ -71324,7 +71637,13 @@ "supported_output_modalities": [ "text" ], - "supports_audio_input": true + "supports_audio_input": true, + "provider_specific_entry": { + "realtime_ga_only": 1, + "transcription_json_or_text_only": 1, + "transcription_usage_duration": 1 + }, + "supports_native_streaming": true }, "azure/gpt-realtime-translate": { "deprecation_date": "2027-11-09", @@ -71342,7 +71661,11 @@ "audio" ], "supports_audio_input": true, - "supports_audio_output": true + "supports_audio_output": true, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "supports_native_streaming": true }, "aihubmix/agnes-2.5-flash": { "input_cost_per_token": 3e-08, @@ -72781,7 +73104,8 @@ "max_input_tokens": 65536, "mode": "evaluation", "output_cost_per_token": 0.0, - "source": "https://docs.typesafe.ai/models" + "source": "https://docs.typesafe.ai/models", + "supports_box_questions": true }, "typesafe/jev-preview": { "input_cost_per_token": 4.2e-08, @@ -73731,6 +74055,7 @@ "output_cost_per_token_priority": 1.875e-05, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -75312,6 +75637,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_ultrafast": 9.9e-05, "output_cost_per_token_ultrafast": 6.6e-05, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/chat/completions", @@ -75334,6 +75660,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, @@ -76629,23 +76956,19 @@ }, "vertex_ai/gemini-robotics-er-2-preview": { "cache_read_input_token_cost": 1e-07, - "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1e-06, - "input_cost_per_token_batches": 5e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 131072, "max_output_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 5e-06, "output_cost_per_token": 5e-06, - "output_cost_per_token_batches": 2.5e-06, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/v1/completions" ], "supported_modalities": [ "text", @@ -76675,23 +76998,19 @@ }, "gemini-robotics-er-2-preview": { "cache_read_input_token_cost": 1e-07, - "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1e-06, - "input_cost_per_token_batches": 5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, "max_output_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 5e-06, "output_cost_per_token": 5e-06, - "output_cost_per_token_batches": 2.5e-06, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/v1/completions" ], "supported_modalities": [ "text", @@ -76758,6 +77077,7 @@ "supports_audio_input": true, "supports_audio_output": false, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76817,6 +77137,7 @@ ], "supports_audio_input": true, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76874,6 +77195,7 @@ "text" ], "supports_audio_input": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -77231,6 +77553,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -77389,6 +77712,7 @@ "output_cost_per_token_priority": 9e-07, "source": "https://platform.minimax.io/docs/guides/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -79022,6 +79346,7 @@ "mode": "responses", "output_cost_per_token": 5.5e-06, "output_cost_per_token_above_272k_tokens": 8.25e-06, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-terra.html", "supported_endpoints": [ "/v1/chat/completions", @@ -79044,6 +79369,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, @@ -81164,6 +81490,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81180,6 +81507,7 @@ "output_cost_per_token": 3.08e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81196,6 +81524,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81212,6 +81541,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 700a35edb06..40ea9d15ff4 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -3158,7 +3158,7 @@ "batches": false, "messages": false, "interactions": false, - "systemone": false, + "systemone": true, "a2a": false, "generateContent": false } diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index b14f43635b2..2f613396413 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -504,10 +504,13 @@ def _lazy_lock(app: "FastAPI", module_path: str) -> asyncio.Lock: def _register_feature(app: "FastAPI", feat: LazyFeature, module: object, features: tuple[LazyFeature, ...]) -> None: - before: Final = len(app.router.routes) + routes_before: Final = frozenset(id(route) for route in app.router.routes) feat.register_fn(app, module) lazy_routes: Final[Mapping[str, tuple[BaseRoute, ...]]] = MappingProxyType( - {**_lazy_routes(app), feat.module_path: tuple(app.router.routes[before:])} + { + **_lazy_routes(app), + feat.module_path: tuple(route for route in app.router.routes if id(route) not in routes_before), + } ) app.state.lazy_routes = lazy_routes # rebind-ok: the app owns the record of which routes each feature added app.router.routes[:] = _route_table( # rebind-ok: the app owns its route table diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5a80583b207..5cc1c53d396 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -10237,8 +10237,15 @@ "type": "string" }, "confidence": { - "title": "Confidence", - "type": "number" + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Confidence" }, "probabilities": { "additionalProperties": { @@ -10256,7 +10263,6 @@ "required": [ "type", "choice", - "confidence", "probabilities" ], "title": "ChoiceAnswer", @@ -10687,8 +10693,15 @@ "title": "Choice" }, "confidence": { - "title": "Confidence", - "type": "number" + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Confidence" }, "name": { "anyOf": [ @@ -10718,8 +10731,7 @@ "required": [ "name", "choice", - "probabilities", - "confidence" + "probabilities" ], "title": "OpenAIChoiceAnswer", "type": "object" @@ -10924,8 +10936,15 @@ "OpenAIScoreAnswer": { "properties": { "confidence": { - "title": "Confidence", - "type": "number" + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Confidence" }, "name": { "anyOf": [ @@ -10959,8 +10978,7 @@ "required": [ "name", "score", - "probabilities", - "confidence" + "probabilities" ], "title": "OpenAIScoreAnswer", "type": "object" @@ -10992,8 +11010,15 @@ "additionalProperties": true, "properties": { "confidence": { - "title": "Confidence", - "type": "number" + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Confidence" }, "legend": { "additionalProperties": { @@ -11034,7 +11059,6 @@ "required": [ "type", "score", - "confidence", "legend", "probabilities" ], @@ -23542,8 +23566,15 @@ "type": "string" }, "confidence": { - "title": "Confidence", - "type": "number" + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Confidence" }, "probabilities": { "additionalProperties": { @@ -23561,7 +23592,6 @@ "required": [ "type", "choice", - "confidence", "probabilities" ], "title": "ChoiceAnswer", @@ -29400,8 +29430,15 @@ "additionalProperties": true, "properties": { "confidence": { - "title": "Confidence", - "type": "number" + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Confidence" }, "legend": { "additionalProperties": { @@ -29442,7 +29479,6 @@ "required": [ "type", "score", - "confidence", "legend", "probabilities" ], @@ -34331,6 +34367,7 @@ }, "/openai/v1/realtime/calls": { "post": { + "description": "Exchange a WebRTC SDP offer for an answer using the ephemeral token from client_secrets", "operationId": "proxy_realtime_calls_openai_v1_realtime_calls_post", "responses": { "200": { @@ -34350,6 +34387,7 @@ }, "/openai/v1/realtime/client_secrets": { "post": { + "description": "Create a short-lived client token for a Realtime conversation, transcription, or translation session", "operationId": "create_realtime_client_secret_openai_v1_realtime_client_secrets_post", "responses": { "200": { @@ -54536,6 +54574,7 @@ "paths": { "/openai/v1/realtime/calls": { "post": { + "description": "Exchange a WebRTC SDP offer for an answer using the ephemeral token from client_secrets", "operationId": "proxy_realtime_calls_openai_v1_realtime_calls_post_2", "responses": { "200": { @@ -54555,6 +54594,7 @@ }, "/openai/v1/realtime/client_secrets": { "post": { + "description": "Create a short-lived client token for a Realtime conversation, transcription, or translation session", "operationId": "create_realtime_client_secret_openai_v1_realtime_client_secrets_post_2", "responses": { "200": { @@ -54608,6 +54648,7 @@ }, "/realtime/calls": { "post": { + "description": "Exchange a WebRTC SDP offer for an answer using the ephemeral token from client_secrets", "operationId": "proxy_realtime_calls_realtime_calls_post", "responses": { "200": { @@ -54627,6 +54668,7 @@ }, "/realtime/client_secrets": { "post": { + "description": "Create a short-lived client token for a Realtime conversation, transcription, or translation session", "operationId": "create_realtime_client_secret_realtime_client_secrets_post", "responses": { "200": { @@ -54678,8 +54720,489 @@ ] } }, + "/realtime/translations/calls": { + "post": { + "description": "Exchange a WebRTC SDP offer for an answer using the ephemeral token from client_secrets", + "operationId": "proxy_realtime_calls_realtime_translations_calls_post", + "requestBody": { + "content": { + "application/sdp": { + "schema": { + "type": "string" + } + } + }, + "required": true + }, + "responses": { + "201": { + "content": { + "text/plain": { + "schema": { + "type": "string" + } + } + }, + "description": "WebRTC SDP answer" + } + }, + "summary": "Proxy Realtime Calls", + "tags": [ + "realtime" + ] + } + }, + "/realtime/translations/client_secrets": { + "post": { + "description": "Create a short-lived client token for a Realtime conversation, transcription, or translation session", + "operationId": "create_realtime_client_secret_realtime_translations_client_secrets_post", + "requestBody": { + "content": { + "application/json": { + "example": { + "model": "gpt-realtime-translate", + "session": { + "model": "gpt-realtime-translate", + "type": "translation" + } + }, + "schema": { + "description": "Request body for POST /v1/realtime/client_secrets.\n\nLiteLLM also accepts a top-level `model` field for routing when\nsession.model is absent (LiteLLM extension, not forwarded to OpenAI).", + "properties": { + "expires_after": { + "anyOf": [ + { + "description": "Expiration config for a client secret.", + "properties": { + "anchor": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": "created_at", + "title": "Anchor" + }, + "seconds": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Seconds" + } + }, + "title": "RealtimeExpiresAfter", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "model": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + }, + "session": { + "anyOf": [ + { + "additionalProperties": true, + "description": "Session configuration nested inside the client_secrets request body.\n\nMirrors OpenAI's RealtimeSessionCreateRequest (type=realtime) and\nRealtimeTranscriptionSessionCreateRequest (type=transcription).\nExtra/unknown fields are passed through unchanged.", + "properties": { + "audio": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "input": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "format": { + "anyOf": [ + { + "type": "string" + }, + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Format" + }, + "noise_reduction": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Noise Reduction" + }, + "transcription": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "delay": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Delay" + }, + "keywords": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Keywords" + }, + "language": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Language" + }, + "languages": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Languages" + }, + "model": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + }, + "prompt": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Prompt" + } + }, + "title": "RealtimeAudioTranscriptionConfig", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "turn_detection": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Turn Detection" + } + }, + "title": "RealtimeAudioInputConfig", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "output": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "format": { + "anyOf": [ + { + "type": "string" + }, + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Format" + }, + "language": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Language" + }, + "voice": { + "anyOf": [ + { + "type": "string" + }, + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Voice" + } + }, + "title": "RealtimeAudioOutputConfig", + "type": "object" + }, + { + "type": "null" + } + ] + } + }, + "title": "RealtimeSessionAudioConfig", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "include": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Include" + }, + "instructions": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Instructions" + }, + "max_output_tokens": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Max Output Tokens" + }, + "model": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + }, + "output_modalities": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Output Modalities" + }, + "prompt": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Prompt" + }, + "tool_choice": { + "anyOf": [ + {}, + { + "type": "null" + } + ], + "title": "Tool Choice" + }, + "tools": { + "anyOf": [ + { + "items": { + "additionalProperties": true, + "type": "object" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Tools" + }, + "tracing": { + "anyOf": [ + {}, + { + "type": "null" + } + ], + "title": "Tracing" + }, + "truncation": { + "anyOf": [ + {}, + { + "type": "null" + } + ], + "title": "Truncation" + }, + "type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Type" + } + }, + "title": "RealtimeSessionConfig", + "type": "object" + }, + { + "type": "null" + } + ] + } + }, + "title": "RealtimeClientSecretRequest", + "type": "object" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RealtimeClientSecretResponse" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Create Realtime Client Secret", + "tags": [ + "realtime" + ] + } + }, "/v1/realtime/calls": { "post": { + "description": "Exchange a WebRTC SDP offer for an answer using the ephemeral token from client_secrets", "operationId": "proxy_realtime_calls_v1_realtime_calls_post", "responses": { "200": { @@ -54699,6 +55222,7 @@ }, "/v1/realtime/client_secrets": { "post": { + "description": "Create a short-lived client token for a Realtime conversation, transcription, or translation session", "operationId": "create_realtime_client_secret_v1_realtime_client_secrets_post", "responses": { "200": { @@ -54749,6 +55273,486 @@ "realtime" ] } + }, + "/v1/realtime/translations/calls": { + "post": { + "description": "Exchange a WebRTC SDP offer for an answer using the ephemeral token from client_secrets", + "operationId": "proxy_realtime_calls_v1_realtime_translations_calls_post", + "requestBody": { + "content": { + "application/sdp": { + "schema": { + "type": "string" + } + } + }, + "required": true + }, + "responses": { + "201": { + "content": { + "text/plain": { + "schema": { + "type": "string" + } + } + }, + "description": "WebRTC SDP answer" + } + }, + "summary": "Proxy Realtime Calls", + "tags": [ + "realtime" + ] + } + }, + "/v1/realtime/translations/client_secrets": { + "post": { + "description": "Create a short-lived client token for a Realtime conversation, transcription, or translation session", + "operationId": "create_realtime_client_secret_v1_realtime_translations_client_secrets_post", + "requestBody": { + "content": { + "application/json": { + "example": { + "model": "gpt-realtime-translate", + "session": { + "model": "gpt-realtime-translate", + "type": "translation" + } + }, + "schema": { + "description": "Request body for POST /v1/realtime/client_secrets.\n\nLiteLLM also accepts a top-level `model` field for routing when\nsession.model is absent (LiteLLM extension, not forwarded to OpenAI).", + "properties": { + "expires_after": { + "anyOf": [ + { + "description": "Expiration config for a client secret.", + "properties": { + "anchor": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": "created_at", + "title": "Anchor" + }, + "seconds": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Seconds" + } + }, + "title": "RealtimeExpiresAfter", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "model": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + }, + "session": { + "anyOf": [ + { + "additionalProperties": true, + "description": "Session configuration nested inside the client_secrets request body.\n\nMirrors OpenAI's RealtimeSessionCreateRequest (type=realtime) and\nRealtimeTranscriptionSessionCreateRequest (type=transcription).\nExtra/unknown fields are passed through unchanged.", + "properties": { + "audio": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "input": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "format": { + "anyOf": [ + { + "type": "string" + }, + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Format" + }, + "noise_reduction": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Noise Reduction" + }, + "transcription": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "delay": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Delay" + }, + "keywords": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Keywords" + }, + "language": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Language" + }, + "languages": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Languages" + }, + "model": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + }, + "prompt": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Prompt" + } + }, + "title": "RealtimeAudioTranscriptionConfig", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "turn_detection": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Turn Detection" + } + }, + "title": "RealtimeAudioInputConfig", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "output": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "format": { + "anyOf": [ + { + "type": "string" + }, + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Format" + }, + "language": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Language" + }, + "voice": { + "anyOf": [ + { + "type": "string" + }, + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Voice" + } + }, + "title": "RealtimeAudioOutputConfig", + "type": "object" + }, + { + "type": "null" + } + ] + } + }, + "title": "RealtimeSessionAudioConfig", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "include": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Include" + }, + "instructions": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Instructions" + }, + "max_output_tokens": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Max Output Tokens" + }, + "model": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + }, + "output_modalities": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Output Modalities" + }, + "prompt": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Prompt" + }, + "tool_choice": { + "anyOf": [ + {}, + { + "type": "null" + } + ], + "title": "Tool Choice" + }, + "tools": { + "anyOf": [ + { + "items": { + "additionalProperties": true, + "type": "object" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Tools" + }, + "tracing": { + "anyOf": [ + {}, + { + "type": "null" + } + ], + "title": "Tracing" + }, + "truncation": { + "anyOf": [ + {}, + { + "type": "null" + } + ], + "title": "Truncation" + }, + "type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Type" + } + }, + "title": "RealtimeSessionConfig", + "type": "object" + }, + { + "type": "null" + } + ] + } + }, + "title": "RealtimeClientSecretRequest", + "type": "object" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RealtimeClientSecretResponse" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Create Realtime Client Secret", + "tags": [ + "realtime" + ] + } } } }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c9960d794c6..fc1e7128e10 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -332,6 +332,8 @@ class LiteLLMRoutes(enum.Enum): "/realtime?{model}", "/v1/realtime?{model}", "/openai/v1/realtime?{model}", + "/realtime/translations", + "/v1/realtime/translations", # realtime (GA WebRTC HTTP routes) "/realtime/client_secrets", "/v1/realtime/client_secrets", @@ -342,6 +344,10 @@ class LiteLLMRoutes(enum.Enum): "/realtime/transcription_sessions", "/v1/realtime/transcription_sessions", "/openai/v1/realtime/transcription_sessions", + "/realtime/translations/client_secrets", + "/v1/realtime/translations/client_secrets", + "/realtime/translations/calls", + "/v1/realtime/translations/calls", # responses API "/responses", "/v1/responses", diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py index cdbbd6029fa..57e99edecba 100644 --- a/litellm/proxy/agent_endpoints/auth/managed_authorization.py +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -11,7 +11,15 @@ from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failur from litellm.types.agents import AgentResponse from litellm.types.proxy.agent_identity import AgentIdentityFailure, ManagedAgentContext -_MANAGED_REALTIME_ROUTES: Final = frozenset(("/realtime", "/v1/realtime", "/openai/v1/realtime")) +_MANAGED_REALTIME_ROUTES: Final = frozenset( + ( + "/realtime", + "/v1/realtime", + "/openai/v1/realtime", + "/realtime/translations", + "/v1/realtime/translations", + ) +) _MANAGED_LIVE_ROUTES: Final = frozenset(("/live/sessions", "/v1/live/sessions", "/openai/v1/live/sessions")) _MANAGED_MODEL_ROUTES: Final = frozenset( f"{prefix}/{operation}" @@ -112,7 +120,11 @@ def managed_inference_request( if route in _MANAGED_LIVE_ROUTES: return dict(body) if route in _MANAGED_REALTIME_ROUTES: - model: Final = query_model or body.get("model") + from litellm.proxy.common_utils.http_parsing_utils import resolve_realtime_route_model + + model: Final = resolve_realtime_route_model( + query_model if isinstance(query_model, str) else None, None, route.endswith("/realtime/translations") + ) if not isinstance(model, str) or not model: raise_identity_failure( AgentIdentityFailure(message="Managed inference requires an explicit or configured model") diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 0e204cf9211..baa6602b042 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -13,8 +13,6 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm.anthropic_interface.exceptions import ( - AnthropicErrorDetail, - AnthropicErrorResponse, AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException @@ -34,13 +32,14 @@ from litellm.proxy.common_request_processing import ( resolve_counted_tool_fields, resolve_litellm_call_id, ) -from litellm.proxy.common_utils.error_body_call_id import error_body_call_id +from litellm.proxy.common_utils.anthropic_error_payload import ( + anthropic_handler_error_json_response, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, ) from litellm.proxy.common_utils.openai_error_payload import ( - LITELLM_CALL_ID_HEADER, error_status_code, litellm_call_id_headers, openai_error_param, @@ -61,44 +60,8 @@ router: Final = APIRouter() _NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) -def _with_provider_specific_fields(exc: ProxyException, detail: AnthropicErrorDetail) -> AnthropicErrorDetail: - if not exc.provider_specific_fields: - return detail - with_fields: Final[AnthropicErrorDetail] = {**detail, "provider_specific_fields": exc.provider_specific_fields} - return with_fields - - -def _anthropic_error_detail( - exc: ProxyException, detail: AnthropicErrorDetail, call_id: str | None -) -> AnthropicErrorDetail: - if call_id is None: - return _with_provider_specific_fields(exc, detail) - with_call_id: Final[AnthropicErrorDetail] = { - **_with_provider_specific_fields(exc, detail), - "litellm_call_id": call_id, - } - return with_call_id - - def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSONResponse: - from litellm.proxy.proxy_server import ( - _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does - general_settings_view, - ) - - status_code: Final = int(exc.code) if exc.code is not None and exc.code.isdigit() else 500 - _close_dangling_otel_server_span(request, status_code, exc=exc) - envelope: Final = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=exc.message, - request_id=request.headers.get("x-request-id"), - ) - body_call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER)) - content: Final[AnthropicErrorResponse] = { - **envelope, - "error": _anthropic_error_detail(exc, envelope["error"], body_call_id), - } - return JSONResponse(status_code=status_code, content=content, headers=exc.headers) + return anthropic_handler_error_json_response(exc, request) _TOKEN_COUNT_FIELDS: Final = ("model", "messages", "tools", "system") diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e2ae25686da..9d98e2c9973 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -43,7 +43,9 @@ from litellm.constants import ( REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.dd_tracing import tracer +from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.tool_fields import DropPolicy @@ -125,6 +127,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( NO_TEAM_MEMBERSHIP_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, + budget_windows_absent_cache_key, end_user_cache_key, end_user_restricted_registry_cache_key, get_management_object_ttl, @@ -729,7 +732,7 @@ def _is_cost_explicitly_configured(deployments: Sequence[DeploymentTypedDict]) - def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: - """Whether any of the deployments carries an explicit positive per-token price. + """Whether any of the deployments carries an explicit positive per-token or per-query price. The group's price is read from the deployments named after the routed group and from that group's own alias target, while an alias chain is served from the deployments named @@ -740,14 +743,21 @@ def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> b model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") if model_id is None: continue - raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY) - if _is_positive_cost(raw_entry.get("input_cost_per_token")) or _is_positive_cost( - raw_entry.get("output_cost_per_token") + raw_entry = as_str_mapping(litellm.model_cost.get(model_id)) or _EMPTY_COST_ENTRY + if ( + _is_positive_cost(raw_entry.get("input_cost_per_token")) + or _is_positive_cost(raw_entry.get("output_cost_per_token")) + or _is_positive_cost(_deployment_query_price(deployment["litellm_params"], raw_entry)) ): return True return False +def _deployment_query_price(litellm_params: Mapping[str, object], cost_entry: Mapping[str, object]) -> float | None: + provider: Final = litellm_params.get("custom_llm_provider") or str(litellm_params.get("model")).partition("/")[0] + return decisions_request_price({**cost_entry, "litellm_provider": provider}, decisions_pricing_fields(cost_entry)) + + _EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) @@ -3471,11 +3481,22 @@ async def delete_cache_key_object( ) await publish_auth_cache_invalidation(cache_key=key) + await evict_budget_windows_absent_marker( + entity_type=Litellm_EntityType.KEY, entity_id=hashed_token, user_api_key_cache=user_api_key_cache + ) _delete_cache_key_object: Final = delete_cache_key_object +async def evict_budget_windows_absent_marker( + entity_type: Litellm_EntityType, entity_id: str, user_api_key_cache: UserApiKeyCache +) -> None: + marker: Final = budget_windows_absent_cache_key(entity_type.value, entity_id) + evict_local(cache_keys=(marker,), user_api_key_cache=user_api_key_cache) + await publish_auth_cache_invalidation(cache_key=marker) + + async def delete_cache_key_objects( hashed_tokens: Sequence[str], user_api_key_cache: UserApiKeyCache, @@ -5378,79 +5399,6 @@ async def _granted_model_lists( ) -async def _user_object_or_none( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> LiteLLM_UserTable | None: - try: - return await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - proxy_logging_obj=proxy_logging_obj, - ) - except UserNotFoundError: - return None - - -async def enforced_model_allowlists( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> tuple[Sequence[str], ...]: - """One model allowlist per level that ``common_checks`` enforces on a request from this identity.""" - key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) - if prisma_client is None: - return (key_models, tuple(valid_token.team_models or ())) - team_object: Final = ( - None - if valid_token.team_id is None - else await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - user_object: Final = ( - None - if team_object is not None - else await _user_object_or_none( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - project_object: Final = ( - None - if valid_token.project_id is None - else await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - return ( - key_models, - team_object.models if team_object is not None else (), - await _team_member_granted_models( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - user_object.models if user_object is not None else (), - project_object.models if project_object is not None else (), - ) - - async def collect_matched_model_access_groups( model: str | Sequence[str] | None, valid_token: UserAPIKeyAuth | None, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9350c9c935c..161ff2fdfcd 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1971,6 +1971,7 @@ _MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS: Final = ("/evals",) # can_key_call_model() validates the real target model. _MODEL_ROUTING_SESSION_MODEL_ROUTE_MARKERS: Final = ( "/realtime/client_secrets", + "/realtime/translations/client_secrets", "/realtime/calls", ) _MODEL_ROUTING_ID_FIELDS: Final = ( diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 4bcf020a9c4..73edcd66bd7 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -6,7 +6,7 @@ import copy import functools import itertools import re -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Collection, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from types import MappingProxyType from typing import Final, TypeVar @@ -23,7 +23,12 @@ from litellm.router import Router from litellm.router_api import fallback_chain from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params from litellm.types.utils import LlmProviders -from litellm.utils import ProviderConfigManager, get_valid_models, potential_model_names_in_lookup_order +from litellm.utils import ( + AZURE_MODEL_PLACEHOLDER, + ProviderConfigManager, + get_valid_models, + potential_model_names_in_lookup_order, +) LLM_PROVIDER_PREFIXES: Final = frozenset(llm_provider.value for llm_provider in LlmProviders) _PROVIDERS_SENDING_THE_DEPLOYMENT_PATH_AS_THE_MODEL_ID: Final = frozenset({LlmProviders.TOGETHER_AI.value}) @@ -35,7 +40,9 @@ _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) _MODEL_LISTING_POOL: Final = ThreadPoolExecutor(thread_name_prefix="litellm-model-listing") -_PROVIDERS_LISTED_ONLY_THROUGH_A_DEPLOYMENT: Final = frozenset({LlmProviders.OPENROUTER.value}) +_PROVIDERS_LISTED_ONLY_THROUGH_A_DEPLOYMENT: Final = frozenset( + {LlmProviders.OPENROUTER.value, LlmProviders.HOSTED_VLLM.value} +) _Listed: Final = TypeVar("_Listed") @@ -373,6 +380,10 @@ def _ids_routing_to_catalog_models( catalog_ids: Sequence[str], routes_to_model: Callable[[str, str], bool], ) -> tuple[str, ...]: + if name_pattern.count("*") > deployment_pattern.count("*"): + return _catalog_ids_a_wider_name_matches( + name_pattern, deployment_pattern, provider, catalog_ids, routes_to_model + ) if name_pattern.count("*") != deployment_pattern.count("*") or _lists_name_fragments( name_pattern, deployment_pattern ): @@ -417,6 +428,23 @@ def _ids_routing_to_catalog_models( ) +def _catalog_ids_a_wider_name_matches( + name_pattern: str, + deployment_pattern: str, + provider: str, + catalog_ids: Sequence[str], + routes_to_model: Callable[[str, str], bool], +) -> tuple[str, ...]: + name_regex: Final = _wildcard_name_regex(name_pattern) + return tuple( + catalog_id + for catalog_id in _ids_routing_to_catalog_models( + deployment_pattern, deployment_pattern, provider, catalog_ids, routes_to_model + ) + if name_regex.fullmatch(catalog_id) + ) + + def _model_info_field(deployment: Mapping[str, object], field: str) -> object: model_info: Final = deployment.get("model_info") return model_info.get(field) if is_str_keyed_mapping(model_info) else None @@ -454,11 +482,12 @@ def _routes_to_model( return lambda listed_id, model: routed_models is None or model in routed_models(listed_id, team_id) -def _deployment_pattern_naming_listed_ids(wildcard_suffix: str, litellm_params: LiteLLM_Params | None) -> str | None: +def _deployment_pattern_naming_listed_ids(wildcard_model: str, litellm_params: LiteLLM_Params | None) -> str | None: if litellm_params is None or "/" not in litellm_params.model or not _check_wildcard_routing(litellm_params.model): return None + name_suffix: Final = wildcard_model.split("/", 1)[-1] deployment_suffix: Final = litellm_params.model.partition("/")[2] - return litellm_params.model if wildcard_suffix != "*" or deployment_suffix != "*" else None + return litellm_params.model if name_suffix != "*" or deployment_suffix != "*" else None def get_known_models_from_wildcard( @@ -466,6 +495,43 @@ def get_known_models_from_wildcard( litellm_params: LiteLLM_Params | None = None, routes_to_model: Callable[[str, str], bool] = lambda _listed_id, _model: True, ) -> list[str]: + provider_id_prefix: Final = _custom_llm_provider_prefix(wildcard_model, litellm_params) + if litellm_params is None or provider_id_prefix is None: + return list(_known_models_from_wildcard(wildcard_model, litellm_params, routes_to_model)) + return [ + listed_id.removeprefix(provider_id_prefix) + for listed_id in _known_models_from_wildcard( + wildcard_model, + litellm_params.model_copy(update={"model": f"{provider_id_prefix}{litellm_params.model}"}), + lambda listed_id, model: routes_to_model( + listed_id.removeprefix(provider_id_prefix), model.removeprefix(provider_id_prefix) + ), + ) + ] + + +def _custom_llm_provider_prefix(wildcard_model: str, litellm_params: LiteLLM_Params | None) -> str | None: + if ( + litellm_params is None + or not litellm_params.custom_llm_provider + or "/" in wildcard_model + or wildcard_model == "*" + or "/" in litellm_params.model + or not _check_wildcard_routing(litellm_params.model) + ): + return None + return f"{litellm_params.custom_llm_provider}/" + + +def _without_azure_placeholder(catalog_ids: Sequence[str]) -> tuple[str, ...]: + return tuple(catalog_id for catalog_id in catalog_ids if catalog_id != AZURE_MODEL_PLACEHOLDER) + + +def _known_models_from_wildcard( + wildcard_model: str, + litellm_params: LiteLLM_Params | None, + routes_to_model: Callable[[str, str], bool], +) -> tuple[str, ...]: wildcard_model_to_expand: Final = ( litellm_params.model if wildcard_model == "*" @@ -474,15 +540,14 @@ def get_known_models_from_wildcard( and "/" in litellm_params.model else wildcard_model ) - try: - wildcard_provider_prefix, wildcard_suffix = wildcard_model_to_expand.split("/", 1) - except ValueError: # safely fail - return [] deployment_pattern_naming_listed_ids: Final = ( wildcard_model_to_expand if wildcard_model_to_expand != wildcard_model - else _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params) + else _deployment_pattern_naming_listed_ids(wildcard_model, litellm_params) ) + wildcard_provider_prefix, has_provider_prefix, wildcard_suffix = wildcard_model_to_expand.partition("/") + if not has_provider_prefix and deployment_pattern_naming_listed_ids is None: + return () # Use provider from litellm_params when available, otherwise from wildcard prefix # (e.g., "openai" from "openai/*" - needed for BYOK where wildcard isn't in router) @@ -499,16 +564,14 @@ def get_known_models_from_wildcard( wildcard_models = get_provider_models(provider=provider, litellm_params=litellm_params) if wildcard_models is None: - return [] + return () if deployment_pattern_naming_listed_ids is not None: - return list( - _ids_routing_to_catalog_models( - name_pattern=wildcard_model, - deployment_pattern=deployment_pattern_naming_listed_ids, - provider=provider, - catalog_ids=wildcard_models, - routes_to_model=routes_to_model, - ) + return _ids_routing_to_catalog_models( + name_pattern=wildcard_model, + deployment_pattern=deployment_pattern_naming_listed_ids, + provider=provider, + catalog_ids=wildcard_models if has_provider_prefix else _without_azure_placeholder(wildcard_models), + routes_to_model=routes_to_model, ) if wildcard_suffix != "*": model_prefix: Final = wildcard_suffix.replace("*", "") @@ -524,7 +587,7 @@ def get_known_models_from_wildcard( # add model prefix to wildcard models wildcard_models = [f"{model_prefix}{model.removeprefix(provider_id_prefix)}" for model in wildcard_models] - return list( + return tuple( dict.fromkeys(_under_wildcard_prefix(model, wildcard_provider_prefix, provider) for model in wildcard_models) ) @@ -562,9 +625,17 @@ def _names_no_deployment_serves_by_name( return tuple(name for name in names if not deployed_by_name(name, team_id)) +def _can_expand_to_an_allowed_name(wildcard_pattern: str, allowed_model_names: Collection[str] | None) -> bool: + if allowed_model_names is None: + return True + pattern: Final = _wildcard_name_regex(wildcard_pattern) + return any(pattern.match(name) for name in allowed_model_names) + + def expand_wildcard_deployments_for_model_info( deployments: list[dict[str, object]], llm_router: Router | None = None, + allowed_model_names: Collection[str] | None = None, ) -> list[dict[str, object]]: """Expand wildcard deployments into one row per known provider model. @@ -581,19 +652,11 @@ def expand_wildcard_deployments_for_model_info( raw_params = deployment.get("litellm_params") litellm_params_dict: dict[str, object] = raw_params if isinstance(raw_params, dict) else {} litellm_model = str(litellm_params_dict.get("model") or "") - - # Determine the wildcard pattern to expand. - # Branch order matters: only fall to litellm_model when model_name is - # also a wildcard, so a concrete model_name is never overwritten. - if _check_wildcard_routing(model_name) and ("/" in model_name or model_name == "*"): - wildcard_pattern = model_name - elif _check_wildcard_routing(model_name) and _check_wildcard_routing(litellm_model): - wildcard_pattern = litellm_model - elif _check_wildcard_routing(model_name): - wildcard_pattern = model_name - else: + if not _check_wildcard_routing(model_name): expanded.append(deployment) continue + if not _can_expand_to_an_allowed_name(model_name, allowed_model_names): + continue try: litellm_params = LiteLLM_Params.model_validate(litellm_params_dict) if litellm_params_dict else None @@ -601,7 +664,7 @@ def expand_wildcard_deployments_for_model_info( expanded.append(deployment) continue expanded_names = get_known_models_from_wildcard( - wildcard_model=wildcard_pattern, + wildcard_model=model_name, litellm_params=litellm_params, routes_to_model=_routes_to_model(routed_models, _deployment_team_id(deployment)), ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a863de8b0f8..1dd0647d25f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -148,6 +148,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg populate_request_with_path_params, read_raw_json_body, read_request_body, + resolve_realtime_route_model, rewrite_request_model, safe_get_request_headers, safe_get_request_query_params, @@ -1246,11 +1247,21 @@ def _apply_budget_limits_to_end_user_params( async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth: - return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model")) + return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model", "")) + + +async def user_api_key_auth_realtime_websocket(websocket: WebSocket) -> UserAPIKeyAuth: + route_path: Final[object] = websocket.scope.get("path") + model: Final = resolve_realtime_route_model( + model=websocket.query_params.get("model"), + intent=websocket.query_params.get("intent"), + is_translation=isinstance(route_path, str) and route_path.endswith("/realtime/translations"), + ) + return await user_api_key_auth_websocket_for_model(websocket, model=model or "") async def user_api_key_auth_live_websocket(websocket: WebSocket) -> UserAPIKeyAuth: - return await user_api_key_auth_websocket_for_model(websocket, model=None) + return await user_api_key_auth_websocket_for_model(websocket, model="") def _has_required_baggage_tag_mapping() -> bool: @@ -1266,7 +1277,7 @@ def _has_required_baggage_tag_mapping() -> bool: async def user_api_key_auth_websocket_for_model( - websocket: WebSocket, model: str | None, *, close_on_auth_error: bool = True + websocket: WebSocket, model: str, *, close_on_auth_error: bool = True ) -> UserAPIKeyAuth: ws_scope: Final = websocket.scope or {} scope_headers: Final = list(ws_scope.get("headers") or []) diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 12d97ef6cce..b70432d1e72 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -770,9 +770,10 @@ def codex_model_catalog( A model Codex does not know gets the fallback entry, with the same base instructions Codex itself uses so the agent never runs without a system prompt, and offers the reasoning efforts `groups` gives its model group - that the installed Codex, at `client_version`, can parse. Apart from aliases, the catalog is - the one `/v1/models?client_version=...` serves, built from the installed Codex's own entries - instead of the proxy's vendored ones. + that the installed Codex, at `client_version`, can parse. Apart from aliases, and the window the + deployments' `model_info.max_input_tokens` gives a model Codex knows, the catalog is the one + `/v1/models?client_version=...` serves, built from the installed Codex's own entries instead of + the proxy's vendored ones. """ rows: Final = tuple(_codex_row(m, upstreams.get(m.id), groups) for m in models) body: Final = codex_models_response_json( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9282ee7f703..5c538c0f6bd 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -30,6 +30,7 @@ import orjson from fastapi import HTTPException, Request, status from fastapi.responses import JSONResponse, Response, StreamingResponse from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.types import Receive, Scope, Send from typing_extensions import Never @@ -170,7 +171,9 @@ ProxyRouteType: TypeAlias = Literal[ "_arealtime", "_aresponses_websocket", "acreate_realtime_client_secret", + "acreate_realtime_translation_client_secret", "arealtime_calls", + "arealtime_translation_calls", "aget_responses", "adelete_responses", "acancel_responses", @@ -780,7 +783,7 @@ def resolve_counted_tool_fields( return resolved -def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException: +def proxy_exception_from_http_exception(exc: StarletteHTTPException, headers: dict[str, str]) -> ProxyException: raw_detail: Final = attribute_of(exc, "detail", str(exc)) message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} diff --git a/litellm/proxy/common_utils/anthropic_error_payload.py b/litellm/proxy/common_utils/anthropic_error_payload.py new file mode 100644 index 00000000000..89001194f7f --- /dev/null +++ b/litellm/proxy/common_utils/anthropic_error_payload.py @@ -0,0 +1,172 @@ +"""Auth, budget and model-access failures are raised by ``user_api_key_auth`` before the Messages +handler runs, so the central handlers pick the Anthropic envelope from the request route.""" + +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final + +from fastapi import Request, status +from fastapi.responses import JSONResponse +from starlette.types import Scope + +from litellm.anthropic_interface.exceptions import ( + ANTHROPIC_ERROR_TYPE_MAP, + AnthropicErrorDetail, + AnthropicErrorResponse, + AnthropicErrorType, + AnthropicExceptionMapping, +) +from litellm.proxy._types import ProxyException +from litellm.proxy.auth.auth_utils import get_request_route +from litellm.proxy.common_utils.error_body_call_id import error_body_call_id +from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER, openai_error_type_for_status + +ANTHROPIC_MESSAGES_ROUTES: Final[frozenset[str]] = frozenset( + { + "/v1/messages", + "/v1/messages/count_tokens", + "/claude_code_gateway/v1/messages", + "/claude_code_gateway/v1/messages/count_tokens", + } +) +ANTHROPIC_PASS_THROUGH_PREFIX: Final = "/anthropic/" +ANTHROPIC_SHAPED_MODEL_LIST_ROUTES: Final[frozenset[str]] = frozenset( + {"/v1/models", "/models", "/claude_code_gateway/v1/models"} +) +PROXY_ERROR_TYPE_BY_STATUS: Final[Mapping[int, AnthropicErrorType]] = MappingProxyType( + {**ANTHROPIC_ERROR_TYPE_MAP, 402: "billing_error", 409: "conflict_error", 504: "timeout_error"} +) + + +def proxy_error_type(status_code: int) -> AnthropicErrorType: + named: Final = PROXY_ERROR_TYPE_BY_STATUS.get(status_code) + if named is not None: + return named + if status_code == status.HTTP_408_REQUEST_TIMEOUT or not 400 <= status_code < 500: + return "api_error" + return "invalid_request_error" + + +def uses_anthropic_error_envelope(request: Request) -> bool: + route: Final = get_request_route(request) + if route in ANTHROPIC_MESSAGES_ROUTES or route.startswith(ANTHROPIC_PASS_THROUGH_PREFIX): + return True + return route in ANTHROPIC_SHAPED_MODEL_LIST_ROUTES and "anthropic-version" in request.headers + + +def route_error_content( + scope: Scope, anthropic_type: AnthropicErrorType, message: str, openai_content: Mapping[str, object] +) -> Mapping[str, object]: + request: Final = Request(scope) + if not uses_anthropic_error_envelope(request): + return openai_content + content: Final[AnthropicErrorResponse] = {"type": "error", "error": {"type": anthropic_type, "message": message}} + request_id: Final = request.headers.get("x-request-id") + return content if request_id is None else {**content, "request_id": request_id} + + +def anthropic_error_status_code(exc: ProxyException) -> int: + code: Final = exc.code + if code.isdigit(): + return int(code) + return status.HTTP_500_INTERNAL_SERVER_ERROR + + +def _with_provider_specific_fields(exc: ProxyException, detail: AnthropicErrorDetail) -> AnthropicErrorDetail: + if not exc.provider_specific_fields: + return detail + with_fields: Final[AnthropicErrorDetail] = {**detail, "provider_specific_fields": exc.provider_specific_fields} + return with_fields + + +def _with_litellm_error_type( + exc: ProxyException, status_code: int, detail: AnthropicErrorDetail +) -> AnthropicErrorDetail: + if not exc.type or exc.type in (detail["type"], openai_error_type_for_status(status_code)): + return detail + with_type: Final[AnthropicErrorDetail] = {**detail, "litellm_error_type": exc.type} + return with_type + + +def _with_call_id(detail: AnthropicErrorDetail, call_id: str | None) -> AnthropicErrorDetail: + if call_id is None: + return detail + with_call_id: Final[AnthropicErrorDetail] = {**detail, "litellm_call_id": call_id} + return with_call_id + + +def _anthropic_content( + exc: ProxyException, + status_code: int, + request_id: str | None, + general_settings: Mapping[str, object], + error_type: AnthropicErrorType | None, +) -> AnthropicErrorResponse: + envelope: Final = AnthropicExceptionMapping.transform_to_anthropic_error( + status_code=status_code, + raw_message=exc.message, + request_id=request_id, + error_type=error_type, + ) + body_call_id: Final = error_body_call_id(general_settings, exc.headers.get(LITELLM_CALL_ID_HEADER)) + content: Final[AnthropicErrorResponse] = { + **envelope, + "error": _with_call_id(_with_provider_specific_fields(exc, envelope["error"]), body_call_id), + } + return content + + +def anthropic_handler_error_content( + exc: ProxyException, + status_code: int, + request_id: str | None, + general_settings: Mapping[str, object], +) -> AnthropicErrorResponse: + return _anthropic_content(exc, status_code, request_id, general_settings, None) + + +def anthropic_error_content( + exc: ProxyException, + status_code: int, + request_id: str | None, + general_settings: Mapping[str, object], +) -> AnthropicErrorResponse: + content: Final = _anthropic_content(exc, status_code, request_id, general_settings, proxy_error_type(status_code)) + with_type: Final[AnthropicErrorResponse] = { + **content, + "error": _with_litellm_error_type(exc, status_code, content["error"]), + } + return with_type + + +def _anthropic_json_response( + exc: ProxyException, + request: Request, + build_content: Callable[[ProxyException, int, str | None, Mapping[str, object]], AnthropicErrorResponse], +) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a JSONResponse returned instead of a raised error must stamp the OTel server span like the global ProxyException handler does + general_settings_view, + ) + + status_code: Final = anthropic_error_status_code(exc) + _close_dangling_otel_server_span(request, status_code, exc=exc) + content: Final = build_content(exc, status_code, request.headers.get("x-request-id"), general_settings_view()) + return JSONResponse(status_code=status_code, content=content, headers=exc.headers) + + +def anthropic_error_json_response(exc: ProxyException, request: Request) -> JSONResponse: + return _anthropic_json_response(exc, request, anthropic_error_content) + + +def anthropic_handler_error_json_response(exc: ProxyException, request: Request) -> JSONResponse: + return _anthropic_json_response(exc, request, anthropic_handler_error_content) + + +def anthropic_internal_server_error_response(request: Request) -> JSONResponse: + content: Final = AnthropicExceptionMapping.transform_to_anthropic_error( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + raw_message="Internal server error", + request_id=request.headers.get("x-request-id"), + ) + return JSONResponse(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=content) diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index b2640d2a765..edc8048ee73 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -183,7 +183,7 @@ def evict_local(cache_keys: Sequence[str], user_api_key_cache: "UserApiKeyCache" def _drop_from_memory(cache_key: str, user_api_key_cache: "UserApiKeyCache") -> None: - user_api_key_cache.in_memory_cache_for(cache_key).delete_cache(cache_key) + user_api_key_cache.delete_from_memory(cache_key) @with_service_target(AUTH_OBJECTS_TARGET) diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 63fb27d62cf..7b48aa07365 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -7,8 +7,10 @@ never the OpenAI list shape. This module builds that body from the rows ``/v1/mo lists: a model Codex knows keeps Codex's own entry for it, any other model gets the fallback entry Codex itself uses for an unknown slug, and ``model_info.service_tiers`` becomes the entry's ``service_tiers``, which Codex turns into slash commands (``/ultrafast``) that send ``service_tier`` -upstream. A fallback entry offers the model group's ``supported_reasoning_efforts`` as its reasoning -levels, trimmed to the ones the requesting Codex release can parse. +upstream. A ``model_info.max_input_tokens`` the operator sets replaces a known model's +``context_window``, up to the most Codex lets its own ``model_context_window`` raise that model to. +A fallback entry offers the model group's ``supported_reasoning_efforts`` as its reasoning levels, +trimmed to the ones the requesting Codex release can parse. The stock entries are Codex 0.159.3's bundled catalog, vendored unchanged from https://github.com/openai/codex/blob/rust-v0.159.3/codex-rs/models-manager/models.json (Apache-2.0) @@ -158,6 +160,7 @@ class CodexStockModel(LiteLLMBaseModel): supports_search_tool: bool = False supported_reasoning_levels: tuple[CodexReasoningLevel, ...] = () context_window: int | None = None + max_context_window: int | None = None class CodexStockCatalog(LiteLLMBaseModel): @@ -173,8 +176,9 @@ class CodexCatalogRow: """What the catalog needs to know about one listed model: its public id, the listing's mode and input limit, the upstream model behind it, the reasoning effort its deployments pin, and the `model_info` values the operator set; `service_tiers` holds one raw `model_info.service_tiers` - value per deployment behind the id, None where a deployment sets none. `supported_reasoning_efforts` - is the model group's, None when unknown.""" + value per deployment behind the id, None where a deployment sets none, and + `configured_max_input_tokens` one `model_info.max_input_tokens` per deployment the same way. + `supported_reasoning_efforts` is the model group's, None when unknown.""" id: str mode: str | None = None @@ -185,6 +189,7 @@ class CodexCatalogRow: service_tiers: tuple[object, ...] = () supported_reasoning_efforts: tuple[str, ...] | None = None search_tool_support: tuple[bool | None, ...] = () + configured_max_input_tokens: tuple[int | None, ...] = () @cache @@ -333,6 +338,25 @@ def _alias_context_window(row: CodexCatalogRow, stock_model: CodexStockModel) -> return MappingProxyType({"context_window": min(stock_model.context_window, row.max_input_tokens)}) +def _configured_window(row: CodexCatalogRow, stock_model: CodexStockModel) -> Mapping[str, int]: + """The `context_window` the operator's `model_info.max_input_tokens` gives a known model's entry, + or nothing, which keeps `stock_model`'s own. + + A request to the row can route to any of its deployments, so the window is the smallest one of them + takes: its `max_input_tokens` where the operator set a positive one, Codex's own window where not. A + limit set on only some deployments can therefore lower the window but never raise it. The window never + goes past `stock_model`'s `max_context_window`, the most Codex lets its own `model_context_window` + setting raise that model to.""" + per_deployment: Final = row.configured_max_input_tokens + declared: Final = tuple(limit for limit in per_deployment if limit is not None and limit > 0) + unset_window: Final = () if len(declared) == len(per_deployment) else (stock_model.context_window,) + if not declared or None in unset_window: + return MappingProxyType({}) + smallest: Final = min(limit for limit in (*declared, *unset_window) if limit is not None) + ceiling: Final = stock_model.max_context_window + return MappingProxyType({"context_window": smallest if ceiling is None else min(smallest, ceiling)}) + + def _alias_reasoning_level(row: CodexCatalogRow, stock_model: CodexStockModel) -> str | None: efforts: Final = frozenset(level.effort for level in stock_model.supported_reasoning_levels) return row.reasoning_effort if row.reasoning_effort in efforts else None @@ -398,6 +422,7 @@ def _entry( "upgrade": upgrade, **({} if search_tool is None else {"supports_search_tool": search_tool}), **(_deployment_defaults(row, stock_model) if is_alias and alias_defaults == "deployment" else {}), + **_configured_window(row, stock_model), } if configured is None: return stock_model.model_copy(update=listing) @@ -473,10 +498,12 @@ def codex_models_response_json( byte limit silently and keeps its bundled one, so the body holds only the entries that fit, the ones offering a service tier taken first and the rest in listing order, an entry too large for the bytes left passed over, every kept entry at its listing position; `byte_limit=None` - keeps every entry. A row listed under another name than the Codex model its `upstream_model` - names (an alias) takes that model's entry; `alias_defaults="deployment"` drops the defaults - Codex picks for that model in favor of what the alias's deployments pin, while `"codex"` keeps - them, as Codex would use them for the model's own name. `client_version` is the requesting Codex's version, which decides the reasoning + keeps every entry. A row whose deployments set `model_info.max_input_tokens` gets the smallest + window any of them takes as its known model's `context_window`, capped at that model's + `max_context_window`. A row listed under another name than the Codex model its `upstream_model` names (an alias) takes that model's + entry; `alias_defaults="deployment"` drops the defaults Codex picks for that model in favor of + what the alias's deployments pin, while `"codex"` keeps them, as Codex would use them for the + model's own name. `client_version` is the requesting Codex's version, which decides the reasoning levels a fallback entry may offer; None counts as a version too old for custom levels. """ known: Final = bundled_codex_models() if stock is None else stock @@ -519,6 +546,9 @@ def _catalog_row( search_tool_support=( registry.resolve(lookup_id, kind="search_tool_support", team_id=team_id) if registry is not None else () ), + configured_max_input_tokens=( + registry.group(lookup_id, field="max_input_tokens", team_id=team_id) if registry is not None else () + ), ) @@ -532,9 +562,9 @@ def codex_catalog_rows( """`rows` joined with the router's configured metadata, looked up by each entry's internal id so team-scoped rows resolve the way the Anthropic listing's display names do, a `model_group_alias` reading its target's deployments under its own name; `team_id` is the requesting key's team, so the - upstream model that picks Codex's stock entry and the tiers are read only off the deployments its - requests can route to, and the reasoning efforts are the ones `/model_group/info` reports for the - entry. Display names are read under `label_entries`' ids when given.""" + upstream model that picks Codex's stock entry, the tiers, and the configured input limits are read + only off the deployments its requests can route to, and the reasoning efforts are the ones + `/model_group/info` reports for the entry. Display names are read under `label_entries`' ids when given.""" lookup_ids: Final = MappingProxyType(dict(entries)) label_ids: Final = lookup_ids if label_entries is None else MappingProxyType(dict(label_entries)) return tuple( diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 81ff460613c..baf586acef6 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -37,6 +37,16 @@ _DECODED_JSON: Final = TypeAdapter(object) _JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) +def resolve_realtime_route_model(model: str | None, intent: str | None, is_translation: bool) -> str | None: + if model is not None: + return model + if is_translation: + return "gpt-realtime-translate" + if intent == "transcription": + return "gpt-realtime-whisper" + return None + + def resolve_inference_model( body_model: object, settings: Mapping[str, object], @@ -572,14 +582,16 @@ async def get_form_data(request: Request) -> dict[str, Any]: Handles when OpenAI SDKs pass form keys as `timestamp_granularities[]="word"` instead of `timestamp_granularities=["word", "sentence"]` """ form: Final = await request.form() - parsed_form_data: Final[dict[str, Any]] = {} - for key, value in form.multi_items(): # not dict(form), which keeps only the last repeat - if key.endswith("[]"): - clean_key = key[:-2] - parsed_form_data.setdefault(clean_key, []).append(value) - else: - parsed_form_data[key] = value - return parsed_form_data + form_items: Final = tuple(form.multi_items() if hasattr(form, "multi_items") else form.items()) + array_keys: Final = frozenset(key[:-2] for key, _ in form_items if key.endswith("[]")) + normalized_items: Final = tuple((key.removesuffix("[]"), value) for key, value in form_items) + normalized_keys: Final = tuple(dict.fromkeys(key for key, _ in normalized_items)) + return { + key: [value for item_key, value in normalized_items if item_key == key] + if key in array_keys + else next(value for item_key, value in reversed(normalized_items) if item_key == key) + for key in normalized_keys + } async def convert_upload_files_to_file_data( diff --git a/litellm/proxy/common_utils/model_deprecation.py b/litellm/proxy/common_utils/model_deprecation.py index 21ad8cf1fde..5c764d1a33e 100644 --- a/litellm/proxy/common_utils/model_deprecation.py +++ b/litellm/proxy/common_utils/model_deprecation.py @@ -1,14 +1,15 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import date, datetime, timezone from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, TypeAlias import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetModelCostMap from litellm.types.proxy.model_deprecation import ( DEFAULT_DEPRECATION_WARN_DAYS, DeprecationStatus, @@ -20,6 +21,9 @@ if TYPE_CHECKING: from litellm.router import Router _NO_MODEL_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) +_SHARED_SCHEDULE_PROVIDER: Final[Mapping[str, str]] = MappingProxyType({"azure_ai": "azure"}) + +LoadedCatalog: TypeAlias = Callable[[], Mapping[str, Mapping[str, object]]] @dataclass(frozen=True, slots=True) @@ -27,6 +31,7 @@ class _ResolvedDeprecation: deprecation_date: date litellm_model: str | None litellm_provider: str | None + copied_from_another_row: bool = False def _parse_deprecation_date(raw_value: object) -> date | None: @@ -47,8 +52,8 @@ def _successor_model(raw_value: object) -> str | None: return raw_value.strip() or None if isinstance(raw_value, str) else None -def _cost_map_lookup(model_key: object) -> _ResolvedDeprecation | None: - if not isinstance(model_key, str) or not model_key: +def _cost_map_lookup(model_key: str | None, catalog: Mapping[str, object]) -> _ResolvedDeprecation | None: + if model_key is None: return None entry: Final = litellm.model_cost.get(model_key) if not isinstance(entry, Mapping): @@ -61,6 +66,51 @@ def _cost_map_lookup(model_key: object) -> _ResolvedDeprecation | None: deprecation_date=parsed, litellm_model=model_key, litellm_provider=provider if isinstance(provider, str) else None, + copied_from_another_row=isinstance(entry, AliasedCostEntry) + and "deprecation_date" not in entry.registered_fields + and model_key not in catalog, + ) + + +def _first_own_row(model_keys: Sequence[str | None], catalog: Mapping[str, object]) -> _ResolvedDeprecation | None: + found: Final = tuple( + resolved for resolved in (_cost_map_lookup(key, catalog) for key in model_keys) if resolved is not None + ) + return next((resolved for resolved in found if not resolved.copied_from_another_row), next(iter(found), None)) + + +def _deployment_provider(litellm_params: Mapping[str, object]) -> str | None: + declared: Final = litellm_params.get("custom_llm_provider") + if isinstance(declared, str) and declared: + return declared + model: Final = litellm_params.get("model") + prefix: Final = model.split("/", 1)[0] if isinstance(model, str) and "/" in model else None + return prefix if prefix in litellm.provider_list else None + + +def _provider_qualified(provider: str | None, model_key: str | None) -> str | None: + if provider is None or model_key is None or model_key.startswith(f"{provider}/"): + return None + return f"{provider}/{model_key}" + + +def _as_model_key(raw_value: object) -> str | None: + return raw_value if isinstance(raw_value, str) and raw_value else None + + +def _without_provider(provider: str | None, model_name: str | None) -> str | None: + if provider is None or model_name is None: + return model_name + return model_name.removeprefix(f"{provider}/") + + +def _schedule_keys(provider: str | None, model_name: str | None) -> tuple[str | None, ...]: + shared_provider: Final = _SHARED_SCHEDULE_PROVIDER.get(provider) if provider is not None else None + unprefixed: Final = _without_provider(provider, model_name) + return ( + _provider_qualified(provider, unprefixed), + _provider_qualified(shared_provider, unprefixed), + model_name, ) @@ -70,34 +120,37 @@ def _mapping_field(deployment: Mapping[str, object], key: str) -> Mapping[str, o def _resolve_deployment_deprecation( - deployment: Mapping[str, object], + deployment: Mapping[str, object], catalog: Mapping[str, object] ) -> _ResolvedDeprecation | None: """Resolve a deployment's deprecation date, preferring its explicit override""" model_info: Final = _mapping_field(deployment, "model_info") - raw_model: Final = _mapping_field(deployment, "litellm_params").get("model") + litellm_params: Final = _mapping_field(deployment, "litellm_params") + model_key: Final = _as_model_key(litellm_params.get("model")) override: Final = _parse_deprecation_date(model_info.get("deprecation_date")) if override is not None: - provider: Final = model_info.get("litellm_provider") + override_provider: Final = model_info.get("litellm_provider") return _ResolvedDeprecation( deprecation_date=override, - litellm_model=raw_model if isinstance(raw_model, str) else None, - litellm_provider=provider if isinstance(provider, str) else None, + litellm_model=model_key, + litellm_provider=override_provider if isinstance(override_provider, str) else None, ) - unprefixed: Final = raw_model.split("/", 1)[1] if isinstance(raw_model, str) and "/" in raw_model else None - return next( - ( - resolved - for resolved in ( - _cost_map_lookup(model_info.get("base_model")), - _cost_map_lookup(raw_model), - _cost_map_lookup(unprefixed), - ) - if resolved is not None - ), - None, + provider: Final = _deployment_provider(litellm_params) + base_model: Final = _as_model_key(model_info.get("base_model")) + model_name: Final = model_key.split("/", 1)[1] if model_key is not None and "/" in model_key else model_key + picks: Final = tuple( + pick + for pick in ( + _first_own_row(_schedule_keys(provider, base_model), catalog), + _first_own_row( + (_provider_qualified(provider, model_key), model_key, *_schedule_keys(provider, model_name)), catalog + ), + ) + if pick is not None ) + own_rows: Final = tuple(pick for pick in picks if not pick.copied_from_another_row) + return min(own_rows or picks, key=lambda pick: pick.deprecation_date, default=None) def _classify(days_until: int, warn_within_days: int) -> DeprecationStatus: @@ -124,13 +177,17 @@ def _checked_successor(model_name: str, model_info: Mapping[str, object], llm_ro def _build_info( - deployment: Mapping[str, object], llm_router: Router, today: date, warn_within_days: int + deployment: Mapping[str, object], + llm_router: Router, + today: date, + warn_within_days: int, + catalog: Mapping[str, object], ) -> ModelDeprecationInfo | None: model_name: Final = deployment.get("model_name") if not isinstance(model_name, str) or not model_name: return None - resolved: Final = _resolve_deployment_deprecation(deployment) + resolved: Final = _resolve_deployment_deprecation(deployment, catalog) if resolved is None: return None @@ -182,15 +239,12 @@ def _bucket(models: Sequence[ModelDeprecationInfo], status: DeprecationStatus) - ) -def _collect_from_router(llm_router: Router, today: date, warn_within_days: int) -> tuple[ModelDeprecationInfo, ...]: +def _collect_from_router( + llm_router: Router, today: date, warn_within_days: int, catalog: Mapping[str, object] +) -> tuple[ModelDeprecationInfo, ...]: deployments: Final = llm_router.registry().deployments() or () - deduped: Final = _dedupe( - tuple( - info - for info in (_build_info(deployment, llm_router, today, warn_within_days) for deployment in deployments) - if info is not None - ) - ) + infos: Final = (_build_info(deployment, llm_router, today, warn_within_days, catalog) for deployment in deployments) + deduped: Final = _dedupe(tuple(info for info in infos if info is not None)) verbose_logger.debug( "model_deprecation: %d/%d deployments carry a deprecation date", @@ -204,12 +258,15 @@ def collect_model_deprecations( llm_router: Router | None, warn_within_days: int = DEFAULT_DEPRECATION_WARN_DAYS, today: date | None = None, + loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, ) -> ModelDeprecationResponse: """Bucket every deployment carrying a deprecation date by how urgent it is""" snapshot_time: Final = datetime.now(timezone.utc) effective_today: Final = today or snapshot_time.date() deduped: Final = ( - _collect_from_router(llm_router, effective_today, warn_within_days) if llm_router is not None else () + _collect_from_router(llm_router, effective_today, warn_within_days, loaded_catalog()) + if llm_router is not None + else () ) return ModelDeprecationResponse( diff --git a/litellm/proxy/common_utils/realtime_utils.py b/litellm/proxy/common_utils/realtime_utils.py index c641748843e..d4fd4556ed4 100644 --- a/litellm/proxy/common_utils/realtime_utils.py +++ b/litellm/proxy/common_utils/realtime_utils.py @@ -1,3 +1,4 @@ +import json from functools import lru_cache from typing import Final @@ -10,7 +11,7 @@ def realtime_request_body(model: str | None) -> bytes: Generate the realtime websocket request body. Cached with LRU semantics to avoid repeated string formatting work while keeping memory usage bounded. """ - return f'{{"model": "{model or ""}"}}'.encode() + return json.dumps({"model": model}).encode() _realtime_request_body: Final = realtime_request_body diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index dd04e531f2d..4e3059cfe26 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -56,6 +56,22 @@ class TokenCountSettings: framing: int | None count_model: TokenCountModel | None tokenizer: TokenCountTokenizer | None = None + tool_framing: int | None = None + responses_framing: int | None = None + responses_tool_framing: int | None = None + + @property + def frames_counts(self) -> bool: + return any( + framing is not None + for framing in (self.framing, self.tool_framing, self.responses_framing, self.responses_tool_framing) + ) + + def hidden_prompt_tokens(self, carries_tools: bool, responses_api: bool) -> int: + chat: Final = self.tool_framing if carries_tools else None + responses: Final = self.responses_tool_framing if carries_tools else self.responses_framing + narrower: Final = responses if responses_api else chat + return (narrower if narrower is not None else self.framing) or 0 @dataclass(frozen=True, slots=True) @@ -105,6 +121,15 @@ def _framing(value: object) -> int | None: return None +def _first_framing(source: Mapping[str, object], keys: tuple[str, ...]) -> int | None: + return next((framing for framing in map(_framing, map(source.get, keys)) if framing is not None), None) + + +def _configured_framing(configured: Mapping[str, object], default: Mapping[str, object], *keys: str) -> int | None: + operator_sets_one: Final = _first_framing(configured, (*keys, "token_count_framing")) is not None + return _first_framing(configured if operator_sets_one else default, keys) + + def _non_blank(value: object) -> str | None: return value if isinstance(value, str) and value.strip() else None @@ -175,13 +200,21 @@ def _tokenizer(value: object) -> TokenCountTokenizer | None: def resolve_token_count_settings( configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] ) -> TokenCountSettings: - framing: Final = _framing(configured.get("token_count_framing")) return TokenCountSettings( - framing=framing if framing is not None else _framing(default.get("token_count_framing")), + framing=_configured_framing(configured, default, "token_count_framing"), count_model=_resolved_count_model( configured.get("token_count_model"), default.get("token_count_model"), deployment_region ), tokenizer=_tokenizer(default.get("token_count_tokenizer")), + tool_framing=_configured_framing(configured, default, "token_count_tool_framing"), + responses_framing=_configured_framing(configured, default, "token_count_responses_framing"), + responses_tool_framing=_configured_framing( + configured, + default, + "token_count_responses_tool_framing", + "token_count_responses_framing", + "token_count_tool_framing", + ), ) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index e64fd755ef9..a2272c3f73d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -24,6 +24,7 @@ T = TypeVar("T", bound=BaseModel) from litellm.constants import AUTH_OBJECTS_TARGET as AUTH_OBJECTS_TARGET # noqa: E402, PLC0414 # public re-export _HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") +_BUDGET_WINDOWS_ABSENT_PREFIX: Final = "budget_windows_absent:" def is_user_key_cache_key(key: str) -> bool: @@ -70,6 +71,7 @@ class UserApiKeyCache(DualCache): default_in_memory_ttl: float | None = None, default_redis_ttl: float | None = None, key_object_in_memory_cache: InMemoryCache | None = None, + budget_windows_absent_in_memory_cache: InMemoryCache | None = None, ) -> None: super().__init__( in_memory_cache=in_memory_cache, @@ -83,19 +85,33 @@ class UserApiKeyCache(DualCache): default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl, ) + self.budget_windows_absent_cache: Final = DualCache( + in_memory_cache=budget_windows_absent_in_memory_cache + or InMemoryCache(default_ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL), + default_in_memory_ttl=default_in_memory_ttl, + ) def in_memory_cache_for(self, key: str) -> InMemoryCache: return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache + def delete_from_memory(self, key: str) -> None: + self.in_memory_cache_for(key).delete_cache(key) + if key.startswith(_BUDGET_WINDOWS_ABSENT_PREFIX): + self.budget_windows_absent_cache.delete_cache(key) + def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None) -> None: super().update_cache_ttl(default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl) self.key_object_cache.update_cache_ttl( default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl ) + self.budget_windows_absent_cache.update_cache_ttl( + default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=None + ) def update_in_memory_max_size(self, max_size: int | None) -> None: super().update_in_memory_max_size(max_size) self.key_object_cache.update_in_memory_max_size(max_size) + self.budget_windows_absent_cache.update_in_memory_max_size(max_size) def attach_redis_cache( self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None @@ -272,6 +288,7 @@ class UserApiKeyCache(DualCache): def flush_cache(self) -> None: super().flush_cache() self.key_object_cache.in_memory_cache.flush_cache() + self.budget_windows_absent_cache.flush_cache() async def async_set_cache_pipeline( self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object @@ -400,6 +417,11 @@ def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str: NO_TEAM_MEMBERSHIP_SENTINEL: Final = "__no_team_membership__" +def budget_windows_absent_cache_key(entity_type: str, entity_id: str) -> str: + """Cache key marking a key or team that holds no budget windows, so the spend fallback skips its row read.""" + return f"{_BUDGET_WINDOWS_ABSENT_PREFIX}{entity_type}:{entity_id}" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/config_resolvers/settings_store.py b/litellm/proxy/config_resolvers/settings_store.py index 70486be0068..3be66821915 100644 --- a/litellm/proxy/config_resolvers/settings_store.py +++ b/litellm/proxy/config_resolvers/settings_store.py @@ -42,10 +42,12 @@ _EMPTY_ROWS: Final[Mapping[DbRow, Mapping[str, JsonValue]]] = MappingProxyType({ class SettingsStore(MutableMapping[str, JsonValue]): - def __init__(self, section: Section) -> None: + def __init__(self, section: Section, *, managed_db_keys: frozenset[str] | None = None) -> None: self._section: Final = section + self._managed_db_keys: Final = managed_db_keys self._yaml_values: Mapping[str, JsonValue] = _EMPTY_VALUES self._database_rows: Mapping[DbRow, Mapping[str, JsonValue]] = _EMPTY_ROWS + self._unmanaged_db_keys: frozenset[str] = frozenset() self._runtime_values: Mapping[str, JsonValue] = _EMPTY_VALUES self._deleted_runtime_keys: frozenset[str] = frozenset() @@ -77,15 +79,24 @@ class SettingsStore(MutableMapping[str, JsonValue]): return self.owned_by_config(key) and self._db_value_is_shadowed(key) def apply_db_row(self, row: DbRow, db_row: Mapping[str, JsonValue]) -> None: + applied_row: Final = MappingProxyType({key: db_row[key] for key in db_row if self._db_manages(row, key)}) previous_row: Final = self._database_rows.get(row, _EMPTY_VALUES) changed: Final = frozenset( key - for key in (*previous_row, *db_row) - if previous_row.get(key, ABSENT) != db_row.get(key, ABSENT) # pyright: ignore[reportUnknownArgumentType] # JsonValue vs Absent compare + for key in (*previous_row, *applied_row) + if previous_row.get(key, ABSENT) != applied_row.get(key, ABSENT) # pyright: ignore[reportUnknownArgumentType] # JsonValue vs Absent compare ) - self._database_rows = MappingProxyType({**self._database_rows, row: MappingProxyType(dict(db_row))}) + self._database_rows = MappingProxyType({**self._database_rows, row: applied_row}) + if row == self._section: + self._unmanaged_db_keys = frozenset(key for key in db_row if key not in applied_row) self._clear_runtime_keys(changed) + def unmanaged_db_keys(self) -> tuple[str, ...]: + return tuple(sorted(self._unmanaged_db_keys)) + + def _db_manages(self, row: DbRow, key: str) -> bool: + return row != self._section or self._managed_db_keys is None or key in self._managed_db_keys + def resolved(self) -> Mapping[str, JsonValue]: return MappingProxyType(dict(self)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py index ef007155c46..66a2d427944 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py @@ -261,7 +261,8 @@ class ResponsesStreamRestorer: own window, keyed by the family, the item id and the part index. - Any `*.done` event closes the stream of the same family. Whatever its window still holds goes out first, as a copy of that stream's last delta event -- so it carries - the stream's own ids, and repeats that event's `sequence_number`. Then every text + the stream's own ids. It takes the done event's `sequence_number`, and every later + event's number moves up by one, so the numbers stay unique and in order. Then every text field on the done event is restored in full: its string fields other than identifiers, plus any `part` or `item` it repeats. - `response.completed` / `response.incomplete` repeat the whole reply, and are @@ -273,6 +274,8 @@ class ResponsesStreamRestorer: self._rehydrate: Final = rehydrate self._carries: Final[dict[ResponsesStreamKey, str]] = {} # mutable-ok: per-stream windows advanced in place. self._last_deltas: Final[dict[ResponsesStreamKey, object]] = {} # mutable-ok: newest delta per stream. + self._inserted: int = 0 + self._last_sequence_number: int | None = None async def restore(self, event: object) -> tuple[object, ...]: """The events to emit in place of `event`: any flush, then the event itself.""" @@ -281,7 +284,7 @@ class ResponsesStreamRestorer: return (event,) if kind.endswith(".delta") and kind not in RESPONSES_BINARY_DELTAS: await self._restore_delta(event, kind) - return (event,) + return self._numbered((), event) slots: Final[SlotSink] = [] flushed: Final = await self._flush(responses_stream_key(event, kind)) if kind.endswith(".done") else () if kind.endswith(".done"): @@ -294,12 +297,29 @@ class ResponsesStreamRestorer: for item in read_list(read_field(event, "response"), "output"): collect_response_item(item, slots) await rehydrate_slots(slots, self._rehydrate) - return (*flushed, event) + return self._numbered(flushed, event) async def finish(self) -> tuple[object, ...]: """Flushes every stream the provider never closed, e.g. a truncated reply.""" - flushed: Final = tuple([await self._flush(key) for key in tuple(self._carries)]) - return tuple(itertools.chain.from_iterable(flushed)) + flushed: Final = tuple(itertools.chain.from_iterable([await self._flush(key) for key in tuple(self._carries)])) + last: Final = self._last_sequence_number + if last is not None: + for offset, flush in enumerate(flushed, start=1): + write_field(flush, "sequence_number", last + offset) + return flushed + + def _numbered(self, flushed: tuple[object, ...], event: object) -> tuple[object, ...]: + """Gives each flush the slot `event` arrived in and moves `event` and later events up.""" + number: Final = read_field(event, "sequence_number") + if not isinstance(number, int): + return (*flushed, event) + for offset, flush in enumerate(flushed): + write_field(flush, "sequence_number", number + self._inserted + offset) + self._inserted += len(flushed) + self._last_sequence_number = number + self._inserted + if self._inserted: + write_field(event, "sequence_number", self._last_sequence_number) + return (*flushed, event) async def _restore_delta(self, event: object, kind: str) -> None: text: Final = read_field(event, "delta") diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 9c0a574ae3b..b4652e010ff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -107,9 +107,20 @@ class MCPSecurityGuardrail(CustomGuardrail): from litellm.proxy.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.proxy_server import prisma_client async with global_mcp_server_manager.catalog.operation(): - registry: Final = global_mcp_server_manager.get_registry() - registered_names: Final = set(registry.keys()) - - return frozenset(requested_servers - registered_names) + not_servers: Final = frozenset( + name + for name in requested_servers + if global_mcp_server_manager.get_mcp_server_by_identifier(name) is None + ) + if prisma_client is None: + return not_servers + return frozenset( + [ + name + for name in not_servers + if await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is None + ] + ) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index cd01d15d834..f9e85c31d46 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -13,7 +13,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs +from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_parent_otel_span_from_kwargs from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, @@ -53,9 +53,16 @@ def _response_total_tokens(kwargs: Mapping[str, object], response_obj: object) - _CONNECTION_SCOPED_CALL_TYPES: Final = frozenset({CallTypes.aresponses_websocket.value}) +def _is_websocket_passthrough_row(kwargs: Mapping[str, object]) -> bool: + passthrough_payload: Final = as_str_mapping(kwargs.get("passthrough_logging_payload")) + return passthrough_payload is not None and passthrough_payload.get("request_method") == "WEBSOCKET" + + def _slots_released(kwargs: Mapping[str, object]) -> int: """A websocket turn callback releases nothing; the socket's slot is released once at disconnect.""" - return 0 if kwargs.get("call_type") in _CONNECTION_SCOPED_CALL_TYPES else 1 + if kwargs.get("call_type") in _CONNECTION_SCOPED_CALL_TYPES or _is_websocket_passthrough_row(kwargs): + return 0 + return 1 class CacheObject(TypedDict): diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c83551a815c..1cbe5eb39e8 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -709,7 +709,7 @@ class RequestRateLimiterStash: A single mutable instance is shared by every context forked from the request task (the SDK call, streaming generators, and the logging worker's captured context all see the same object), which is what makes the - ``reservation_released`` flag and ``parallel_slot`` clearing effective + released flags and ``parallel_slot`` clearing effective across sibling callbacks: the first release wins, later callbacks observe the cleared state. @@ -746,9 +746,32 @@ class RequestRateLimiterStash: batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = () reservation_released: bool = False + io_reservation_released: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) socket_agent_windows: tuple[RateLimitDescriptor, ...] = () + def clear_reserved_tokens(self) -> None: + self.reserved_tokens = 0 + self.ptu_reserved_tokens = 0 + self.itpm_reserved_tokens = 0 + self.otpm_reserved_tokens = 0 + + def mark_reservation_released(self) -> None: + self.reservation_released = True + self.io_reservation_released = True + + def discard_released_tpm_reservation(self) -> None: + if self.reservation_released: + self.reserved_tokens = 0 + self.ptu_reserved_tokens = 0 + self.reservation_released = False + + def discard_released_io_reservation(self) -> None: + if self.io_reservation_released: + self.itpm_reserved_tokens = 0 + self.otpm_reserved_tokens = 0 + self.io_reservation_released = False + @dataclass(frozen=True, slots=True) class CounterRefund: @@ -4238,6 +4261,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, ) + stash.discard_released_io_reservation() if itpm_reserved > 0: itpm_scopes: Final = tuple( (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY @@ -4597,6 +4621,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # incremented so post-call reconciliation only applies # the (actual - reserved) delta to those — unreserved # scopes get charged the full actual usage instead. + stash.discard_released_tpm_reservation() stash.reserved_tokens = estimated_tokens stash.ptu_reserved_tokens = ptu_estimated_tokens stash.ptu_reserved_window_identities = frozenset( @@ -5140,7 +5165,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._response_usage(response_obj) ) or self._aggregate_only_total_tokens(self._response_usage(kwargs.get("combined_usage_object"))) - if not response_usage[2] and not combined_usage[2] and aggregate_total <= 0 and not stash.reservation_released: + if ( + not response_usage[2] + and not combined_usage[2] + and aggregate_total <= 0 + and not stash.io_reservation_released + ): return () resolved_usage: Final = ( response_usage @@ -5153,25 +5183,25 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) billable_input, completion_tokens, _ = resolved_usage - if stash.reservation_released or ( + if stash.io_reservation_released or ( not stash.itpm_reserved_window_identities and not stash.otpm_reserved_window_identities ): return self._build_reservation_aware_tpm_ops( targets=tuple(stash.itpm_reserved_scopes), - reserved_scopes=frozenset() if stash.reservation_released else stash.itpm_reserved_scopes, + reserved_scopes=frozenset() if stash.io_reservation_released else stash.itpm_reserved_scopes, actual_tokens=billable_input, - reserved_tokens=0 if stash.reservation_released else itpm_reserved, + reserved_tokens=0 if stash.io_reservation_released else itpm_reserved, ) + self._build_reservation_aware_tpm_ops( targets=tuple(stash.otpm_reserved_scopes), - reserved_scopes=frozenset() if stash.reservation_released else stash.otpm_reserved_scopes, + reserved_scopes=frozenset() if stash.io_reservation_released else stash.otpm_reserved_scopes, actual_tokens=completion_tokens, - reserved_tokens=0 if stash.reservation_released else otpm_reserved, + reserved_tokens=0 if stash.io_reservation_released else otpm_reserved, ) itpm_ops: Final[Sequence[ReservationAwareIncrementOperation]] = ( self._build_project_reservation_ops( targets=tuple(stash.itpm_reserved_scopes), - reserved_scopes=frozenset() if stash.reservation_released else stash.itpm_reserved_scopes, + reserved_scopes=frozenset() if stash.io_reservation_released else stash.itpm_reserved_scopes, actual_tokens=billable_input, reserved_tokens=itpm_reserved, reservation_window_identities=stash.itpm_reserved_window_identities, @@ -5182,7 +5212,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): otpm_ops: Final[Sequence[ReservationAwareIncrementOperation]] = ( self._build_project_reservation_ops( targets=tuple(stash.otpm_reserved_scopes), - reserved_scopes=frozenset() if stash.reservation_released else stash.otpm_reserved_scopes, + reserved_scopes=frozenset() if stash.io_reservation_released else stash.otpm_reserved_scopes, actual_tokens=completion_tokens, reserved_tokens=otpm_reserved, reservation_window_identities=stash.otpm_reserved_window_identities, @@ -5402,7 +5432,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): total_tokens = self._aggregate_only_total_tokens(usage=_usage) stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - reserved_tokens: Final = stash.reserved_tokens if stash is not None else 0 + open_reservation: Final = stash if stash is not None and not stash.reservation_released else None + reserved_tokens: Final = open_reservation.reserved_tokens if open_reservation is not None else 0 reserved_model: Final = stash.reserved_model if stash is not None else None reserved_scopes: Final[frozenset[tuple[str, str]]] = stash.reserved_scopes if stash is not None else frozenset() # Reconciliation must target the same model-scoped counter that the @@ -5464,7 +5495,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reconcile_model=reconcile_model, reserved_scopes=reserved_scopes, reserved_ceiling=stash.ptu_ceiling if stash is not None else None, - reserved_tokens=stash.ptu_reserved_tokens if stash is not None else 0, + reserved_tokens=open_reservation.ptu_reserved_tokens if open_reservation is not None else 0, reserved_windows=stash.ptu_reserved_window_identities if stash is not None else frozenset(), total_tokens=0 if ptu_skips_socket_usage else total_tokens, ) @@ -5561,18 +5592,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else None, ) - await self._apply_usage_operations( - self._build_success_event_pipeline_operations( - kwargs=kwargs, - response_obj=response_obj, - rate_limit_type=rate_limit_type, - ), - litellm_parent_otel_span, + usage_operations: Final = self._build_success_event_pipeline_operations( + kwargs=kwargs, + response_obj=response_obj, + rate_limit_type=rate_limit_type, ) io_token_operations: Final = self._build_io_token_reservation_ops( kwargs=kwargs, response_obj=response_obj, ) + if stash is not None: + stash.clear_reserved_tokens() + await self._apply_usage_operations(usage_operations, litellm_parent_otel_span) if io_token_operations: if isinstance(io_token_operations, list): await self.async_increment_tokens_with_ttl_preservation( @@ -5733,10 +5764,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # already released it (proxy-level rejection that also bubbles up # here as an LLM-error callback). max_parallel_requests is its # own counter and is always decremented per call. - reserved_tokens, itpm_reserved, otpm_reserved = ( - (0, 0, 0) - if stash is None or stash.reservation_released - else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) + reserved_tokens: Final = 0 if stash is None or stash.reservation_released else stash.reserved_tokens + itpm_reserved, otpm_reserved = ( + (0, 0) + if stash is None or stash.io_reservation_released + else (stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) ) tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(failed_call) @@ -5806,6 +5838,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else () ) + if stash is not None and (reserved_tokens > 0 or itpm_reserved > 0 or otpm_reserved > 0): + stash.mark_reservation_released() if pipeline_operations: await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline_post_call( pipeline_operations, parent_otel_span=litellm_parent_otel_span @@ -5820,8 +5854,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): pipeline_operations=project_operations, parent_otel_span=litellm_parent_otel_span, ) - if stash is not None and (reserved_tokens > 0 or itpm_reserved > 0 or otpm_reserved > 0): - stash.reservation_released = True if stash is None and tpm_actual > 0 and _is_socket_turn(failed_call): await self._apply_usage_operations( self._build_success_event_pipeline_operations( @@ -5927,7 +5959,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Idempotent: the slot release clears the stashed acquisition (and slot removal is a no-op ZREM on a second run), and the TPM/ITPM/OTPM - refund is guarded by the stash's ``reservation_released`` flag — if + refund is guarded by the stash's two released flags — if both this hook and async_log_failure_event end up running in the same flow, only the first release/refund applies. A mid-stream failure relayed here with recovered partial usage settles the reservation at @@ -5953,11 +5985,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) stash.batch_tpd_refund_ops = () - if stash.reservation_released: - return - reserved_tokens: Final = stash.reserved_tokens - itpm_reserved: Final = stash.itpm_reserved_tokens - otpm_reserved: Final = stash.otpm_reserved_tokens + reserved_tokens: Final = 0 if stash.reservation_released else stash.reserved_tokens + itpm_reserved: Final = 0 if stash.io_reservation_released else stash.itpm_reserved_tokens + otpm_reserved: Final = 0 if stash.io_reservation_released else stash.otpm_reserved_tokens if reserved_tokens <= 0 and itpm_reserved <= 0 and otpm_reserved <= 0: return tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(request_data) @@ -6024,6 +6054,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): itpm_reserved, otpm_reserved, ) + stash.mark_reservation_released() if combined_ops: await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( increment_list=list(combined_ops), @@ -6040,7 +6071,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): pipeline_operations=project_ops, parent_otel_span=user_api_key_dict.parent_otel_span, ) - stash.reservation_released = True except Exception as e: verbose_proxy_logger.exception("Error releasing TPM reservation on post-call failure: %s", e) return diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 980f0a5d410..8a015d396d0 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -64,7 +64,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # _sanitize_error_information_for_spend_logs, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export deployment_params_provider, get_request_model_access_groups, - sanitize_error_information_for_spend_logs, + sanitize_error_information_for_spend_logs, # pyright: ignore[reportUnusedImport] # backwards-compatible package export should_store_prompts_and_responses_in_spend_logs, ) from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, capture_callback, captured_callback_kwargs @@ -79,7 +79,6 @@ from litellm.types.utils import ( CallTypes, LiteLLMBatch, StandardLoggingPayload, - StandardLoggingPayloadErrorInformation, ) from litellm.utils import get_end_user_id_for_cost_tracking @@ -492,13 +491,6 @@ class _ProxyDBLogger(CustomLogger): # ``.get("traceback")`` / truthy checks, and the TypedDict marks # the field as optional, so omitting is type-safe. _error_information.pop("traceback", None) - # Strip echoed request input + apply DB-size cap before storing in - # the spend-log metadata column (LIT-2992). Result is never None - # here because the input above is constructed non-None. - _error_information = cast( - StandardLoggingPayloadErrorInformation, - sanitize_error_information_for_spend_logs(_error_information, original_exception=original_exception), - ) _metadata["error_information"] = _error_information _metadata = ( # rebind-ok: enriched identity is shared with the earlier immutable capture diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 771d68a0ea0..0269a784445 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -7,23 +7,33 @@ instead of writing immediately on each request. from collections.abc import AsyncGenerator, Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, TypeVar from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import LitellmUserRoles from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, ResponsesAPIResponse, ) -from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypesLiteral, LLMResponseTypes, SpecialEnums +from litellm.types.responses.main import DeleteResponseResult +from litellm.types.utils import ( + ADDRESSED_COMPARISON_RESPONSE_ID_FIELD, + ADDRESSED_RESPONSE_ID_FIELD, + CallTypesLiteral, + LLMResponseTypes, + SpecialEnums, +) if TYPE_CHECKING: from litellm.caching.caching import DualCache @@ -67,6 +77,9 @@ def _proxy_signing_key() -> str | None: _RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_REQUEST_DATA_ADAPTER: Final = TypeAdapter(dict[str, object]) +_ResponseObjectT = TypeVar("_ResponseObjectT", bound=BaseLiteLLMOpenAIResponseObject) +_AddressedResultT = TypeVar("_AddressedResultT", bound=ResponsesAPIResponse | DeleteResponseResult) def _response_payload(response_obj: object) -> Mapping[str, object] | None: @@ -77,9 +90,9 @@ def _response_payload(response_obj: object) -> Mapping[str, object] | None: def _rewrite_advertised_id( - event: BaseLiteLLMOpenAIResponseObject, + event: _ResponseObjectT, rewrite: Callable[[str], str], -) -> BaseLiteLLMOpenAIResponseObject: +) -> _ResponseObjectT: event_id: Final = getattr(event, "id", None) if isinstance(event_id, str) and event_id.startswith("resp_"): setattr(event, "id", rewrite(event_id)) @@ -116,6 +129,35 @@ def _streams_response_ids(request_route: str | None) -> bool: return canonical in _RESPONSES_API_ROUTES or (parent in _RESPONSES_API_ROUTES and response_id != "") +_NO_HELD_IDS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def _upstream_response_id(response_id: str) -> str: + return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id) + + +def _names_the_same_response(response_id: str | None, original_response_id: str) -> bool: + return response_id is not None and _upstream_response_id(response_id) == _upstream_response_id(original_response_id) + + +def _held_id_naming(response_id: str | None, held_ids: Mapping[str, str]) -> str | None: + return next( + ( + held_id + for original_response_id, held_id in held_ids.items() + if _names_the_same_response(response_id, original_response_id) + ), + None, + ) + + +def _id_sent_back(response: ResponsesAPIResponse | DeleteResponseResult, held_ids: Mapping[str, str]) -> str | None: + held_id: Final = _held_id_naming(response.id, held_ids) + if held_id is not None: + return held_id + return response.id if isinstance(response, DeleteResponseResult) else None + + class ResponsesIDSecurity(CustomLogger): def __init__( self, @@ -135,16 +177,64 @@ class ResponsesIDSecurity(CustomLogger): addressed_id_field: Final = _ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE.get(call_type) if addressed_id_field is None: return None - retained_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) + request_data: Final = _REQUEST_DATA_ADAPTER.validate_python(data) + authorized_data: Final = ( + {**request_data, **self._authorized_comparison_response_id_fields(request_data, user_api_key_dict)} + if call_type == "aresponses" + else request_data + ) + retained_id: Final = authorized_data.get(ADDRESSED_RESPONSE_ID_FIELD) addressed_id: Final = ( - retained_id if isinstance(retained_id, str) and retained_id else data.get(addressed_id_field) + retained_id if isinstance(retained_id, str) and retained_id else authorized_data.get(addressed_id_field) ) if not isinstance(addressed_id, str) or not addressed_id: - return data - authorized_id: Final = self.authorize_response_id(addressed_id, user_api_key_dict) - data[addressed_id_field] = authorized_id - data[ADDRESSED_RESPONSE_ID_FIELD] = addressed_id - return data + return authorized_data + return { + **authorized_data, + addressed_id_field: self.authorize_response_id(addressed_id, user_api_key_dict), + ADDRESSED_RESPONSE_ID_FIELD: addressed_id, + } + + def _authorized_comparison_response_id_fields( + self, + data: Mapping[str, object], + user_api_key_dict: "UserAPIKeyAuth", + ) -> Mapping[str, object]: + extra_body: Final = as_str_mapping(data.get("extra_body")) + overriding_extra_body: Final = ( + extra_body if extra_body is not None and "prompt_cache_options" in extra_body else None + ) + options: Final = as_str_mapping( + data.get("prompt_cache_options") + if overriding_extra_body is None + else overriding_extra_body.get("prompt_cache_options") + ) + requested_id: Final = None if options is None else options.get("comparison_response_id") + if options is None or not isinstance(requested_id, str): + return {} + addressed_id: Final = self._comparison_addressed_id( + requested_id, data.get(ADDRESSED_COMPARISON_RESPONSE_ID_FIELD) + ) + authorized_options: Final = { + **options, + "comparison_response_id": self.authorize_response_id(addressed_id, user_api_key_dict), + } + authorized_field: Final[Mapping[str, object]] = ( + {"prompt_cache_options": authorized_options} + if overriding_extra_body is None + else {"extra_body": {**overriding_extra_body, "prompt_cache_options": authorized_options}} + ) + return {**authorized_field, ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: addressed_id} + + def _comparison_addressed_id(self, requested_id: str, retained_id: object) -> str: + if ( + not isinstance(retained_id, str) + or self._is_encrypted_response_id(requested_id) + or not self._is_encrypted_response_id(retained_id) + ): + return requested_id + original_response_id, _, _ = self._decrypt_response_id(retained_id) + return retained_id if original_response_id == requested_id else requested_id def authorize_response_id( self, @@ -275,10 +365,18 @@ class ResponsesIDSecurity(CustomLogger): def _encrypt_response_id( self, - response: BaseLiteLLMOpenAIResponseObject, + response: _ResponseObjectT, user_api_key_dict: "UserAPIKeyAuth", request_cache: dict[str, str] | None = None, - ) -> BaseLiteLLMOpenAIResponseObject: + ) -> _ResponseObjectT: + encrypt: Final = self._response_id_encrypter(user_api_key_dict, request_cache) + return response if encrypt is None else _rewrite_advertised_id(response, encrypt) + + def _response_id_encrypter( + self, + user_api_key_dict: "UserAPIKeyAuth", + request_cache: dict[str, str] | None, + ) -> Callable[[str], str] | None: signing_key: Final = self._get_signing_key() if signing_key is None: verbose_proxy_logger.debug( @@ -287,7 +385,7 @@ class ResponsesIDSecurity(CustomLogger): "Skipping response ID encryption. " "See: https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" ) - return response + return None def encrypt(original_id: str) -> str: cached: Final = request_cache.get(original_id) if request_cache is not None else None @@ -304,47 +402,73 @@ class ResponsesIDSecurity(CustomLogger): request_cache[original_id] = encrypted_id return encrypted_id - return _rewrite_advertised_id(response, encrypt) + return encrypt + + def _streamed_response_id_advertiser( + self, + user_api_key_dict: "UserAPIKeyAuth", + request_data: Mapping[str, object], + ) -> Callable[[str], str] | None: + if not _streams_response_ids(user_api_key_dict.request_route): + return None + if self._general_settings_reader().get("disable_responses_id_security", False): + return None + encrypt: Final = self._response_id_encrypter(user_api_key_dict, request_cache={}) + if encrypt is None: + return None + logging_obj: Final = request_data.get("litellm_logging_obj") + held_ids: Final = self._ids_the_caller_holds(request_data.get(ADDRESSED_RESPONSE_ID_FIELD)) + + def advertise(original_id: str) -> str: + logged_id: Final = encrypt(original_id) + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.client_facing_response_id = logged_id + return _held_id_naming(original_id, held_ids) or logged_id + + return advertise + + def _ids_the_caller_holds(self, addressed_id: object) -> Mapping[str, str]: + if not isinstance(addressed_id, str) or not self._is_encrypted_response_id(addressed_id): + return _NO_HELD_IDS + original_response_id, _, _ = self._decrypt_response_id(addressed_id) + return MappingProxyType({original_response_id: addressed_id}) + + def _advertised_response( + self, + response: _AddressedResultT, + user_api_key_dict: "UserAPIKeyAuth", + request_data: Mapping[str, object], + ) -> _AddressedResultT: + sent_back_id: Final = _id_sent_back( + response, self._ids_the_caller_holds(request_data.get(ADDRESSED_RESPONSE_ID_FIELD)) + ) + logged_response: Final = self._encrypt_response_id(response, user_api_key_dict, request_cache=None) + return logged_response if sent_back_id is None else logged_response.model_copy(update={"id": sent_back_id}) async def async_post_call_success_hook( self, - data: dict, + data: dict[str, object], user_api_key_dict: "UserAPIKeyAuth", response: LLMResponseTypes, ) -> LLMResponseTypes: - """ - Queue response IDs for batch processing instead of writing directly to DB. - - This method adds response IDs to an in-memory queue, which are then - batch-processed by the DBSpendUpdateWriter during regular database update cycles. - """ general_settings: Final = self._general_settings_reader() if general_settings.get("disable_responses_id_security", False): return response - if isinstance(response, ResponsesAPIResponse): - response = cast( - ResponsesAPIResponse, - self._encrypt_response_id(response, user_api_key_dict, request_cache=None), - ) + if isinstance(response, ResponsesAPIResponse | DeleteResponseResult): + return self._advertised_response(response, user_api_key_dict, data) return response async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict + self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict[str, object] ) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]: - general_settings: Final = self._general_settings_reader() - - # Create a request-scoped cache for consistent encryption across streaming chunks. - request_encryption_cache: Final[dict[str, str]] = {} - + advertise: Final = self._streamed_response_id_advertiser(user_api_key_dict, request_data) async for chunk in response: - if ( - isinstance(chunk, BaseLiteLLMOpenAIResponseObject) - and _streams_response_ids(user_api_key_dict.request_route) - and not general_settings.get("disable_responses_id_security", False) - ): - chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) - yield chunk + yield ( + _rewrite_advertised_id(chunk, advertise) + if advertise is not None and isinstance(chunk, BaseLiteLLMOpenAIResponseObject) + else chunk + ) def decrypted_response_id(response_id: str) -> str: diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 6244e5025ce..95816edfecb 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -37,6 +37,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( from litellm.proxy.route_llm_request import route_request from litellm.types.images.main import ImageEditRequestParams from litellm.types.llms.openai import ChatCompletionUserMessage +from litellm.types.utils import CallTypes router: Final = APIRouter() _IMAGE_GENERATION_MODEL_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object) @@ -221,7 +222,10 @@ async def image_generation( return response except Exception as e: await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=data, + dispatched_call_type=CallTypes.aimage_generation.value, ) log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b712ba42a6a..f505406fe0e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -109,6 +109,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo can_org_access_model, delete_cache_key_objects, delete_cache_team_object, + evict_budget_windows_absent_marker, get_jwt_key_mapping_cache_keys_for_tokens, get_org_object, get_team_membership, @@ -612,6 +613,9 @@ async def refresh_cached_team( await publish_auth_cache_invalidation(cache_key=f"team_id:{team_row.team_id}") if team_table.team_alias is not None: await publish_auth_cache_invalidation(cache_key=f"team_alias:{team_table.team_alias}") + await evict_budget_windows_absent_marker( + entity_type=Litellm_EntityType.TEAM, entity_id=team_row.team_id, user_api_key_cache=user_api_key_cache + ) _refresh_cached_team: Final = refresh_cached_team diff --git a/litellm/proxy/mcp_server/server.py b/litellm/proxy/mcp_server/server.py index aeb640a7c8b..bbbbd09e527 100644 --- a/litellm/proxy/mcp_server/server.py +++ b/litellm/proxy/mcp_server/server.py @@ -21,7 +21,7 @@ import httpx from fastapi import FastAPI, HTTPException from pydantic import ConfigDict, TypeAdapter, ValidationError from starlette.requests import Request as StarletteRequest -from starlette.responses import JSONResponse +from starlette.responses import JSONResponse, Response from starlette.routing import Route from starlette.types import Message, Receive, Scope, Send @@ -2489,6 +2489,9 @@ if MCP_AVAILABLE: }, )(scope, receive, send) return + if scope.get("method") == "OPTIONS": + await Response(status_code=204, headers={"Allow": "GET, POST, DELETE, OPTIONS"})(scope, receive, send) + return path: Final = _scope_path(scope) peek: Final = _LazyPeekedBody(receive) if scope.get("method") == "POST" else None if peek is not None: diff --git a/litellm/proxy/middleware/admission_control_middleware.py b/litellm/proxy/middleware/admission_control_middleware.py index 1bb94e23d15..83503d0c3d4 100644 --- a/litellm/proxy/middleware/admission_control_middleware.py +++ b/litellm/proxy/middleware/admission_control_middleware.py @@ -10,6 +10,7 @@ from starlette.responses import JSONResponse from starlette.types import ASGIApp, Receive, Scope, Send from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.anthropic_error_payload import route_error_content ADMISSION_LEASE_SCOPE_KEY: Final = "litellm.admission_lease" @@ -175,7 +176,7 @@ class AdmissionControlMiddleware: state.record_admission() elif state.get_stats().queued >= settings.max_queued_requests: state.record_rejection("queue_full") - await _overloaded_response(state)(scope, receive, send) + await _overloaded_response(state, scope)(scope, receive, send) return else: state.record_queue() @@ -187,7 +188,7 @@ class AdmissionControlMiddleware: except asyncio.TimeoutError: state.record_dequeue() state.record_rejection("queue_timeout") - await _overloaded_response(state)(scope, receive, send) + await _overloaded_response(state, scope)(scope, receive, send) return except asyncio.CancelledError: state.record_dequeue() @@ -313,18 +314,16 @@ def get_admission_control_settings(settings: Mapping[str, object]) -> AdmissionC ) -def _overloaded_response(state: AdmissionControlState) -> JSONResponse: +def _overloaded_response(state: AdmissionControlState, scope: Scope) -> JSONResponse: stats: Final = state.get_stats() + message: Final = f"Worker at capacity: {stats.admitted} in-flight, {stats.queued} queued requests. Retry later." return JSONResponse( status_code=503, headers={"retry-after": "1"}, - content={ - "error": { - "message": ( - f"Worker at capacity: {stats.admitted} in-flight, {stats.queued} queued requests. Retry later." - ), - "type": "overloaded_error", - "code": "503", - } - }, + content=route_error_content( + scope, + "overloaded_error", + message, + {"error": {"message": message, "type": "overloaded_error", "code": "503"}}, + ), ) diff --git a/litellm/proxy/middleware/request_body_decompression_middleware.py b/litellm/proxy/middleware/request_body_decompression_middleware.py index 5eb8a9bfc91..3e144d7bf4c 100644 --- a/litellm/proxy/middleware/request_body_decompression_middleware.py +++ b/litellm/proxy/middleware/request_body_decompression_middleware.py @@ -12,12 +12,14 @@ from starlette.requests import ClientDisconnect, Request from starlette.responses import JSONResponse from starlette.types import ASGIApp, Message, Receive, Scope, Send +from litellm.anthropic_interface.exceptions import AnthropicErrorType from litellm.constants import ( MAX_DECOMPRESSED_REQUEST_BODIES_IN_FLIGHT_MB, MAX_DECOMPRESSED_REQUEST_BODY_SIZE_MB, MAX_REQUEST_BODY_DECOMPRESSION_RATIO, ) from litellm.proxy.auth.auth_utils import get_request_route +from litellm.proxy.common_utils.anthropic_error_payload import proxy_error_type, route_error_content from litellm.proxy.middleware.request_size_limit_middleware import MaxRequestSizeGetter, send_request_too_large RFC_9659_ZSTD_WINDOW_LOG_MAX: Final = 23 @@ -291,13 +293,22 @@ def _replay_body(body: bytes, receive: Receive, more_body: bool = False) -> Rece def _error_response( + scope: Scope, message: str, status_code: int, error_type: str = "invalid_request_error", headers: Mapping[str, str] | None = None, ) -> JSONResponse: + anthropic_type: Final[AnthropicErrorType] = ( + "overloaded_error" if error_type == "overloaded_error" else proxy_error_type(status_code) + ) return JSONResponse( - {"error": {"message": message, "type": error_type, "param": None, "code": str(status_code)}}, + route_error_content( + scope, + anthropic_type, + message, + {"error": {"message": message, "type": error_type, "param": None, "code": str(status_code)}}, + ), status_code=status_code, headers=headers, ) @@ -356,6 +367,7 @@ class RequestBodyDecompressionMiddleware: if codec is None: supported: Final = ", ".join(sorted(self.codecs)) await _error_response( + scope, f"This proxy cannot decode {coding} request bodies. Supported encodings: {supported}", 415, headers={"accept-encoding": supported}, @@ -372,7 +384,7 @@ class RequestBodyDecompressionMiddleware: case _Body(data=data): await self._decode_within_budget(scope, receive, send, codec, data, limits) case _TooLarge(): - await send_request_too_large(send=send, max_request_size_mb=max_mb) + await send_request_too_large(scope=scope, send=send, max_request_size_mb=max_mb) case _ClientDisconnected(): return @@ -382,6 +394,7 @@ class RequestBodyDecompressionMiddleware: reservation: Final = _max_decoded_length(codec, data, limits) if not self.budget.try_reserve(reservation): await _error_response( + scope, "The proxy is decoding too many compressed request bodies right now. " "Retry later, or send this body uncompressed", 503, @@ -407,12 +420,12 @@ class RequestBodyDecompressionMiddleware: case _Body(data=body): await self.app(_identity_scope(scope, len(body)), _replay_body(body, receive), send) case _DecodedTooLarge(): - await _error_response(f"Request body is larger than {limits.max_mb:g} MB when decoded", 413)( + await _error_response(scope, f"Request body is larger than {limits.max_mb:g} MB when decoded", 413)( scope, receive, send ) case _OverRatio(): - await _error_response(f"Request body expands more than {limits.max_ratio} times when decoded", 413)( - scope, receive, send - ) + await _error_response( + scope, f"Request body expands more than {limits.max_ratio} times when decoded", 413 + )(scope, receive, send) case _Undecodable(message=message): - await _error_response(message, 400)(scope, receive, send) + await _error_response(scope, message, 400)(scope, receive, send) diff --git a/litellm/proxy/middleware/request_size_limit_middleware.py b/litellm/proxy/middleware/request_size_limit_middleware.py index cbbc79cc2ca..15bbf05bd01 100644 --- a/litellm/proxy/middleware/request_size_limit_middleware.py +++ b/litellm/proxy/middleware/request_size_limit_middleware.py @@ -4,6 +4,8 @@ from typing import Final from starlette.types import ASGIApp, Message, Receive, Scope, Send +from litellm.proxy.common_utils.anthropic_error_payload import route_error_content + MaxRequestSizeGetter = Callable[[], int | float | None] RequestSizeLimitEnabledGetter = Callable[[], bool] @@ -44,7 +46,7 @@ class RequestSizeLimitMiddleware: content_length: Final = _get_content_length(scope=scope) if content_length is not None and content_length > max_request_size_bytes: - await send_request_too_large(send=send, max_request_size_mb=max_request_size_mb) + await send_request_too_large(scope=scope, send=send, max_request_size_mb=max_request_size_mb) return received_body_bytes = 0 @@ -74,7 +76,7 @@ class RequestSizeLimitMiddleware: except RequestEntityTooLarge: if response_started: raise - await send_request_too_large(send=send, max_request_size_mb=max_request_size_mb) + await send_request_too_large(scope=scope, send=send, max_request_size_mb=max_request_size_mb) def _mb_to_bytes(max_request_size_mb: float | None) -> int | None: @@ -98,11 +100,13 @@ def _get_content_length(scope: Scope) -> int | None: async def send_request_too_large( + scope: Scope, send: Send, max_request_size_mb: float | None, ) -> None: + message: Final = f"Request size is too large. Max size is {max_request_size_mb} MB" body: Final = json.dumps( - {"error": f"Request size is too large. Max size is {max_request_size_mb} MB"}, + route_error_content(scope, "request_too_large", message, {"error": message}), separators=(",", ":"), ).encode("utf-8") await send( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index c7983f40aec..be9f80694b8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,11 +15,11 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast import httpx import openai @@ -51,6 +51,7 @@ from litellm.constants import ( from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, attribute_or_none from litellm.litellm_core_utils.fallback_generalizations import match_routing_generalization from litellm.llms.anthropic.common_utils import ( AnthropicModelInfo, @@ -85,7 +86,12 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.llms.vllm.common_utils import endpoint_under_api_base from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import enforced_model_allowlists +from litellm.proxy.auth.auth_checks import ( + can_agent_call_model, + can_key_call_resolved_model, + key_model_aliases_for_auth_check, + team_model_aliases_for_auth_check, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -93,6 +99,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul get_bearer_token, is_no_auth_dev_mode, user_api_key_auth, + user_api_key_auth_accepted_websocket_for_model, user_api_key_auth_websocket, user_api_key_auth_websocket_for_model, ) @@ -114,6 +121,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_headers, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.model_listing_utils import alias_map from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( @@ -129,16 +137,26 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, without_injected_usage_frames, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + STORED_PROMPT_PROBE_MODEL, HttpPassThroughEndpointHelpers, + WebsocketFrameModelAllowed, + WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, + WebsocketUpstreamFrameClassifier, YieldsToConfiguredPassThroughRoute, create_pass_through_route, create_websocket_passthrough_route, + refuse_websocket, truncated_close_reason, websocket_passthrough_request, ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType -from litellm.proxy.utils import is_known_model +from litellm.proxy.utils import is_known_model, normalize_route_for_root_path from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, @@ -198,7 +216,7 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li """ try: model: Final = request_body.get("model") - return is_known_model(model, llm_router) + return isinstance(model, str) and is_known_model(model, llm_router) except Exception: return False @@ -228,9 +246,8 @@ def is_passthrough_request_streaming(request_body: object) -> bool: A JSON body need not be an object, so a list or scalar can reach here; it carries no streaming flag. """ - if not isinstance(request_body, dict): - return False - return bool(request_body.get("stream", False)) + body: Final = as_str_mapping(request_body) + return body is not None and bool(body.get("stream", False)) def _optional_str(value: object) -> str | None: @@ -3387,19 +3404,6 @@ def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm return joined_path_str -_OPENAI_WS_ALL_MODEL_ACCESS: Final = frozenset( - { - SpecialModelNames.all_proxy_models.value, - SpecialModelNames.all_team_models.value, - "*", - } -) - - -def _has_model_restrictions(model_allowlists: tuple[Sequence[str], ...]) -> bool: - return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for allowlist in model_allowlists for model in allowlist) - - @dataclass(frozen=True, slots=True) class _OpenAIWebsocketRefusal: close_reason: str @@ -3416,6 +3420,16 @@ class _OpenAIWebsocketErrorFrame(TypedDict): error: ReadOnly[_OpenAIWebsocketErrorDetail] +_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = ( + "This gateway could not authorize the model named in this frame, so the session was closed. Try again, " + "and ask a proxy admin to check the gateway logs if it keeps happening." +) + +_OPENAI_WS_ALIASED_MODEL_REFUSAL: Final = ( + "This key cannot use a model alias over OpenAI websocket passthrough. The passthrough sends each frame to " + "OpenAI unchanged, so name the model this key may use directly." +) + _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( close_reason="OpenAI websocket passthrough is disabled", message=( @@ -3424,14 +3438,6 @@ _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( ), ) -_OPENAI_WS_MODEL_RESTRICTED_REFUSAL: Final = _OpenAIWebsocketRefusal( - close_reason="Keys with model restrictions cannot use OpenAI websocket passthrough", - message=( - "Keys with model restrictions cannot use OpenAI websocket passthrough, because this route " - "relays frames to the provider without reading which model they ask for." - ), -) - def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, object]) -> bool: setting: Final = general_settings.get("enable_openai_websocket_passthrough") @@ -3440,19 +3446,9 @@ def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, obje return setting is True -class _OpenAIWebsocketModelAllowlists(Protocol): - async def __call__(self, valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: ... - - -async def _openai_websocket_refusal( - user_api_key_dict: UserAPIKeyAuth, - general_settings: Mapping[str, object], - model_allowlists: _OpenAIWebsocketModelAllowlists, -) -> _OpenAIWebsocketRefusal | None: +def _openai_websocket_refusal(general_settings: Mapping[str, object]) -> _OpenAIWebsocketRefusal | None: if not _is_openai_websocket_passthrough_enabled(general_settings): return _OPENAI_WS_DISABLED_REFUSAL - if _has_model_restrictions(await model_allowlists(user_api_key_dict)): - return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL return None @@ -3467,6 +3463,8 @@ class _WebsocketRelay(Protocol): forward_headers: bool, endpoint: str, accept_websocket: bool, + client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, meter: WebSocketSessionMeter | None = None, ) -> None: ... @@ -3475,18 +3473,95 @@ def _websocket_relay() -> _WebsocketRelay: return websocket_passthrough_request -def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists: - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache +def _proxy_error_type(denial: ProxyException) -> str: + return denial.type.value if isinstance(denial.type, ProxyErrorTypes) else denial.type - async def resolve(valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: - return await enforced_model_allowlists( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + +FrameModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]] + + +def _frame_model_authorizer() -> FrameModelAuthorizer: + return user_api_key_auth_accepted_websocket_for_model + + +def _proxy_model_list() -> Sequence[object] | None: + from litellm.proxy import proxy_server + + return as_object_list(attribute_or_none(proxy_server, "llm_model_list")) + + +def _gateway_remaps_model(model: str, frame_auth: UserAPIKeyAuth, llm_router: Router | None) -> bool: + alias_maps: Final = ( + key_model_aliases_for_auth_check(frame_auth) or {}, + team_model_aliases_for_auth_check(frame_auth) or {}, + alias_map(litellm.model_alias_map), + ) + targets: Final = ( + *(aliases.get(model) for aliases in alias_maps), + llm_router.registry().resolve(model, kind="alias") if llm_router is not None else None, + ) + return any(target not in (None, model) for target in targets) + + +def _aliased_model_verdict( + allowed: WebsocketFrameModelAllowed, probe: WebsocketFrameModelVerdict +) -> WebsocketFrameModelVerdict: + match probe: + case WebsocketFrameModelRefused(error_type=error_type) if error_type.endswith("_model_access_denied"): + return WebsocketFrameModelRefused(message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=error_type) + case WebsocketFrameModelRefused(): + return probe + case WebsocketFrameModelAllowed(): + return allowed + + +def _proxy_frame_model_gate( + authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)], +) -> WebsocketFrameModelGate: + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation + + async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelVerdict: + connect_reservation: Final = attribute_or_none(websocket.state, "budget_reservation") + await release_or_invalidate_budget_reservation( + budget_reservation=connect_reservation if isinstance(connect_reservation, dict) else None ) + frame_auth: Final = await authorize_frame_model(websocket, model) + access_groups: Final = tuple(frame_auth.matched_model_access_groups or ()) + llm_router: Final = proxy_server.llm_router + try: + await can_key_call_resolved_model( + model=model, llm_model_list=_proxy_model_list(), valid_token=frame_auth, llm_router=llm_router + ) + await can_agent_call_model( + model=model, + valid_token=frame_auth, + llm_router=llm_router, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + finally: + await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation) + allowed: Final = WebsocketFrameModelAllowed(access_groups=access_groups) + if not _gateway_remaps_model(model, frame_auth, llm_router): + return allowed + return _aliased_model_verdict(allowed, await gate(websocket, STORED_PROMPT_PROBE_MODEL)) - return resolve + async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: + try: + return await authorize(websocket, model) + except ProxyException as denial: + return WebsocketFrameModelRefused(message=denial.message, error_type=_proxy_error_type(denial)) + except HTTPException as denial: + return WebsocketFrameModelRefused(message=str(denial.detail), error_type="invalid_request_error") + except litellm.BudgetExceededError as exhausted: + return WebsocketFrameModelRefused(message=exhausted.message, error_type="budget_exceeded") + except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band + verbose_proxy_logger.exception("OpenAI websocket passthrough: the frame model gate failed on %s", model) + return WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error") + + return gate def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: @@ -3498,6 +3573,14 @@ def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: return requested_subprotocols[0] if requested_subprotocols else None +async def _opening_model_refusal( + websocket: WebSocket, query_string: str, frame_model_gate: WebsocketFrameModelGate +) -> WebsocketFrameModelRefused | None: + opening_models: Final = dict.fromkeys(httpx.QueryParams(query_string).get_list("model")) + verdicts: Final = tuple([await frame_model_gate(websocket, model) for model in opening_models]) + return next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) + + @router.websocket("/openai_passthrough/{endpoint:path}") @router.websocket("/openai/{endpoint:path}") async def openai_websocket_proxy_route( @@ -3506,12 +3589,12 @@ async def openai_websocket_proxy_route( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], - model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)], + frame_model_gate: Annotated[WebsocketFrameModelGate, Depends(_proxy_frame_model_gate)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket) - refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists) + refusal: Final = _openai_websocket_refusal(general_settings) if refusal is not None: await websocket.accept(subprotocol=negotiated_subprotocol) error_frame: Final[_OpenAIWebsocketErrorFrame] = { @@ -3557,14 +3640,21 @@ async def openai_websocket_proxy_route( await websocket.accept(subprotocol=negotiated_subprotocol) + opening_refusal: Final = await _opening_model_refusal(websocket, query_string, frame_model_gate) + if opening_refusal is not None: + await refuse_websocket(websocket, opening_refusal) + return + await relay( websocket=websocket, target=wss_target, custom_headers=custom_headers, user_api_key_dict=user_api_key_dict, forward_headers=False, - endpoint=websocket.url.path, + endpoint=normalize_route_for_root_path(websocket.url.path) or websocket.url.path, accept_websocket=False, + client_frame_model_gate=frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) @@ -3882,7 +3972,8 @@ async def _tinyfish_json_object_field_names(request: Request) -> frozenset[str] parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_body)) except (json.JSONDecodeError, UnicodeDecodeError): return None - return frozenset(parsed) if isinstance(parsed, dict) else None + fields: Final = as_str_mapping(parsed) + return None if fields is None else frozenset(fields) def _tinyfish_route_timeout() -> float | None: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py new file mode 100644 index 00000000000..f3bad54ffb2 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -0,0 +1,262 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import Final, Literal, TypeAlias, cast # noqa: TID251 # cost helpers type raw frames as a TypedDict list + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.cost_calculator import RealtimeAPITokenUsageProcessor, handle_realtime_stream_cost_calculation +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponsesAPIResponse +from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind +from litellm.types.utils import LlmProviders, ModelResponse, Usage + +WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]] + +OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value + +_SESSION_CREATED: Final = "session.created" +_SESSION_UPDATED: Final = "session.updated" +_RESPONSE_CREATED: Final = "response.created" +_ERROR: Final = "error" +_REALTIME_RESPONSE_DONE: Final = "response.done" +_TRANSCRIPTION_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" +_RESPONSES_TURN_ENDS: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) +_TURN_ENDS: Final = _RESPONSES_TURN_ENDS | {_REALTIME_RESPONSE_DONE} +_MODEL_NAMING_EVENTS: Final = _TURN_ENDS | {_SESSION_CREATED, _SESSION_UPDATED, _RESPONSE_CREATED} +_REALTIME_BILLED_EVENTS: Final = frozenset({_REALTIME_RESPONSE_DONE, _TRANSCRIPTION_COMPLETED}) +_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" +_FRAME_KINDS: Final[Mapping[str, WebsocketUpstreamFrameKind]] = MappingProxyType( + { + _SESSION_CREATED: "opens_session", + _SESSION_UPDATED: "updates_session", + _RESPONSE_CREATED: "starts_turn", + _TRANSCRIPTION_COMPLETED: "bills_usage", + _ERROR: "refuses_turn", + } +) + + +class _NamedModel(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + + +class _AudioInput(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + transcription: _NamedModel | None = None + + +class _SessionAudio(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + input: _AudioInput | None = None + + +class _NamedSession(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + audio: _SessionAudio | None = None + + +class _ModelNamingEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + session: _NamedSession | None = None + response: _NamedModel | None = None + + +class _ResponsesTurnEndEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: Literal["response.completed", "response.incomplete", "response.failed"] + response: ResponsesAPIResponse + + +_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) +_RESPONSES_TURN_END_EVENT: Final = TypeAdapter(_ResponsesTurnEndEvent) + + +@dataclass(frozen=True, slots=True) +class _Billing: + usage: Usage + cost: float + + +@dataclass(frozen=True, slots=True) +class _CompletedTurn: + model: str + response: ResponsesAPIResponse + usage: Usage + + +def _event_model(event: Mapping[str, object]) -> str | None: + if event.get("type") not in _MODEL_NAMING_EVENTS: + return None + try: + naming: Final = _MODEL_NAMING_EVENT.validate_python(event) + except ValidationError: + return None + session: Final = naming.session or _NamedSession() + response: Final = naming.response or _NamedModel() + audio_input: Final = (session.audio or _SessionAudio()).input or _AudioInput() + transcription: Final = audio_input.transcription or _NamedModel() + return session.model or response.model or transcription.model or None + + +def observed_websocket_model(messages: WebsocketMessages) -> str | None: + return next((model for model in map(_event_model, messages) if model is not None), None) + + +def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstreamFrameKind: + frame_type: Final = frame.get("type") + if frame_type in _TURN_ENDS: + return "ends_turn" + if not isinstance(frame_type, str): + return "relayed" + return _FRAME_KINDS.get(frame_type, "relayed") + + +def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None: + if not any(event.get("type") in _REALTIME_BILLED_EVENTS for event in messages): + return None + results: Final = cast( # cast-ok: the realtime cost helpers read raw provider frames through their TypedDict view + OpenAIRealtimeStreamList, tuple(event for event in messages if "type" in event) + ) + usage: Final = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results) + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=usage, + custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + litellm_model_name=model, + litellm_logging_obj=logging_obj, + ) + return _Billing(usage=usage, cost=cost) + + +def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None: + if event.get("type") not in _RESPONSES_TURN_ENDS: + return None + try: + ended: Final = _RESPONSES_TURN_END_EVENT.validate_python(event) + except ValidationError: + verbose_proxy_logger.warning( + "OpenAI websocket passthrough: a %s frame carried no readable usage, so it is not billed", event.get("type") + ) + return None + if ended.response.usage is None: + return None + return _CompletedTurn( + model=ended.response.model or session_model, + response=ended.response, + usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter + ended.response.usage + ), + ) + + +def _responses_billing(messages: WebsocketMessages, model: str, requested_model: str | None) -> _Billing | None: + turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None) + if not turns: + return None + usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns]) + return _Billing( + usage=usage, + cost=sum(_responses_cost(turn.response, (turn.model, requested_model)) for turn in turns), + ) + + +def _model_cost(response: ResponsesAPIResponse, model: str) -> float | None: + try: + return litellm.completion_cost( + completion_response=response, + model=model, + custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + call_type="responses", + ) + except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log + return None + + +def _responses_cost(response: ResponsesAPIResponse, price_as: Sequence[str | None]) -> float: + costs: Final = (_model_cost(response, model) for model in dict.fromkeys(price_as) if model is not None) + cost: Final = next((cost for cost in costs if cost), None) + if cost is None: + verbose_proxy_logger.warning( + "OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", price_as + ) + return 0.0 + return cost + + +def _requested_model(kwargs: Mapping[str, object]) -> str | None: + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) + metadata: Final = None if litellm_params is None else as_str_mapping(litellm_params.get("metadata")) + model_group: Final = None if metadata is None else metadata.get("model_group") + return model_group if isinstance(model_group, str) else None + + +def _priced_result(model: str, billing: _Billing, start_time: datetime, call_id: str) -> ModelResponse: + result: Final = ModelResponse( + id=call_id, + object="chat.completion", + created=int(start_time.timestamp()), + model=model, + usage=billing.usage, + ) + hidden: Final = {"additional_headers": MappingProxyType({_RESPONSE_COST_HEADER: billing.cost})} + result._hidden_params = hidden # pyright: ignore[reportPrivateUsage] # the one settled-cost channel + return result + + +class OpenAIWebsocketPassthroughLoggingHandler: + def openai_websocket_passthrough_handler( + self, + websocket_messages: WebsocketMessages, + logging_obj: LiteLLMLoggingObj, + url_route: str, + start_time: datetime, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict: + model: Final = observed_websocket_model(websocket_messages) + if model is None: + verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no model named, skipping cost", url_route) + unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}} + return unpriced + + logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object + logging_obj.model_call_details.update( + MappingProxyType({"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}) + ) + + billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing( + websocket_messages, model, _requested_model(kwargs) + ) + if billing is None: + verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no usage reported, skipping cost", url_route) + unbilled: Final[PassThroughEndpointLoggingTypedDict] = { + "result": None, + "kwargs": {**kwargs, "model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}, + } + return unbilled + + logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost})) + priced: Final[PassThroughEndpointLoggingTypedDict] = { + "result": _priced_result(model, billing, start_time, logging_obj.litellm_call_id), + "kwargs": { + **kwargs, + "model": model, + "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER, + "response_cost": billing.cost, + }, + } + return priced diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 970cd466337..7f5983d5799 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -6,15 +6,17 @@ import posixpath import re import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from datetime import datetime -from itertools import chain, count, groupby +from functools import reduce +from itertools import chain, count, groupby, takewhile from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast -from urllib.parse import urlencode, urlparse +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast +from urllib.parse import parse_qs, urlencode, urlparse import httpx +import orjson from fastapi import ( APIRouter, Depends, @@ -24,6 +26,7 @@ from fastapi import ( Response, UploadFile, WebSocket, + WebSocketDisconnect, status, ) from fastapi.responses import StreamingResponse @@ -34,7 +37,7 @@ from starlette.routing import BaseRoute, Match, Route from starlette.types import Scope from starlette.websockets import WebSocketState from typing_extensions import ReadOnly -from websockets.asyncio.client import connect +from websockets.asyncio.client import ClientConnection, connect from websockets.exceptions import ( ConnectionClosedError, ConnectionClosedOK, @@ -60,6 +63,8 @@ from litellm.integrations.custom_guardrail import ( from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + as_str_mapping, bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -69,6 +74,7 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_errors import realtime_error_event from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.anthropic.common_utils import ( @@ -154,6 +160,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, WebSocketSessionMeter, WebSocketSessionUsage, + WebsocketUpstreamFrameKind, ) from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage @@ -2574,11 +2581,16 @@ def _resolved_vertex_live_setup( return {**setup_data, "model": setup_model_rewriter(setup_model)} +def _decoded_frame(frame: str | bytes) -> object: + try: + return _DECODED_JSON.validate_python(orjson.loads(frame)) + except orjson.JSONDecodeError: + return _DECODED_JSON.validate_python(json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))) + + def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: try: - decoded: Final[object] = _DECODED_JSON.validate_python( - json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")) - ) + decoded: Final = _decoded_frame(frame) except (json.JSONDecodeError, UnicodeDecodeError): return None return decoded if isinstance(decoded, dict) else None @@ -2637,6 +2649,359 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +@dataclass(frozen=True, slots=True) +class WebsocketFrameModelAllowed: + access_groups: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class WebsocketFrameModelRefused: + message: str + error_type: str + + +WebsocketFrameModelVerdict: TypeAlias = WebsocketFrameModelAllowed | WebsocketFrameModelRefused + + +class WebsocketFrameModelGate(Protocol): + async def __call__(self, websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: ... + + +WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind] +_SESSION_FRAMES_EVERY_ROW_CARRIES: Final[frozenset[WebsocketUpstreamFrameKind]] = frozenset( + {"opens_session", "updates_session"} +) +_FrameT: Final = TypeVar("_FrameT", bound=Mapping[str, object]) + + +@dataclass(frozen=True, slots=True) +class _TurnAuthorization: + requested_model: str | None + access_groups: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class _ForwardClientFrame: + starts_turn: bool + turn: _TurnAuthorization | None + lane: str | None + + +@dataclass(frozen=True, slots=True) +class _RefuseClientFrame: + refusal: WebsocketFrameModelRefused + + +@dataclass(frozen=True, slots=True) +class _ClientFrameRequests: + starts_turn: bool + lane: str | None + turn_model: str | None + models: tuple[str, ...] + unverifiable_model_refusal: str | None + + +_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests( + starts_turn=False, lane=None, turn_model=None, models=(), unverifiable_model_refusal=None +) +_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None, lane=None) +_UNREADABLE_CLIENT_FRAME: Final = _RefuseClientFrame( + refusal=WebsocketFrameModelRefused( + message="The gateway cannot parse this frame, so it cannot check which model it asks for.", + error_type="invalid_request_error", + ) +) +_TURN_STARTING_FRAME: Final = "response.create" +_FRAME_TEXT_WORTH_PARSING: Final = (f'"{_TURN_STARTING_FRAME}"', '"model"', '"prompt"', '"tools"', "\\") +_TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model")) +_TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( + ("session", "input_audio_transcription", "model"), + ("session", "audio", "input", "transcription", "model"), +) +_STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt")) +_TOOL_LIST_PLACEMENTS: Final = (("tools",), ("response", "tools"), ("session", "tools")) +_TOOL_THAT_RUNS_A_MODEL: Final = "image_generation" +STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model" +_STORED_PROMPT_REFUSAL: Final = ( + "This key is limited to specific models, so it cannot run a stored prompt. " + "The gateway cannot check which model a stored prompt runs." +) +_UNNAMED_TOOL_MODEL_REFUSAL: Final = ( + "This key is limited to specific models, so an image_generation tool has to name its model. " + "The gateway cannot check which model an unnamed one runs." +) +_FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") +_LANE_NAME: Final = re.compile(r"[A-Za-z0-9_.-]{1,256}") +_UNROUTABLE_LANE_REFUSAL: Final = _RefuseClientFrame( + refusal=WebsocketFrameModelRefused( + message=( + "This stream_id is not a lane name the provider accepts. " + "The provider's error for it names no lane, so the gateway cannot tell which turn it refuses." + ), + error_type="invalid_request_error", + ) +) + + +def _may_name_a_model(text_data: str) -> bool: + return any(marker in text_data for marker in _FRAME_TEXT_WORTH_PARSING) + + +def _frame_lane(frame: Mapping[str, object]) -> str | None: + stream_id: Final = frame.get("stream_id") + return stream_id if isinstance(stream_id, str) else None + + +def _names_an_unroutable_lane(frame: Mapping[str, object]) -> bool: + stream_id: Final = frame.get("stream_id") + return stream_id is not None and not (isinstance(stream_id, str) and _LANE_NAME.fullmatch(stream_id) is not None) + + +def _frames_of_lane(lane: str | None) -> Callable[[Mapping[str, object]], bool]: + def in_lane(frame: Mapping[str, object]) -> bool: + return _frame_lane(frame) == lane + + return in_lane + + +def _every_frame(_frame: Mapping[str, object]) -> bool: + return True + + +def _without_earlier_session_updates( + upstream_frame_kind: WebsocketUpstreamFrameClassifier, + frames: Sequence[_FrameT], + latest_update: _FrameT, +) -> tuple[_FrameT, ...]: + return tuple( + frame + for index, frame in enumerate(frames) + if frame is latest_update + or upstream_frame_kind(frame) != "updates_session" + or _unbilled_usage_ran_under(map(upstream_frame_kind, frames[index + 1 :])) + ) + + +def _unbilled_usage_ran_under(later_kinds: Iterable[WebsocketUpstreamFrameKind]) -> bool: + return "bills_usage" in takewhile(lambda kind: kind != "updates_session", later_kinds) + + +def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None: + return None if frame is None else as_str_mapping(frame.get(key)) + + +def _placed_value(frame: Mapping[str, object], placement: Sequence[str]) -> object: + *parents, leaf = placement + container: Final = reduce(_nested_frame, parents, frame) + return None if container is None else container.get(leaf) + + +def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[str]]) -> tuple[str, ...]: + placed: Final = (_placed_value(frame, placement) for placement in placements) + return tuple(value for value in placed if isinstance(value, str)) + + +def _placed_tools(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + tool_lists: Final = (_placed_value(frame, placement) for placement in _TOOL_LIST_PLACEMENTS) + tools: Final = chain.from_iterable(as_object_list(tool_list) or () for tool_list in tool_lists) + return tuple(tool for tool in map(as_str_mapping, tools) if tool is not None) + + +def _unverifiable_model_refusal(frame: Mapping[str, object], tools: Sequence[Mapping[str, object]]) -> str | None: + prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS) + if any(isinstance(prompt, Mapping) for prompt in prompts): + return _STORED_PROMPT_REFUSAL + if any(tool.get("type") == _TOOL_THAT_RUNS_A_MODEL and not isinstance(tool.get("model"), str) for tool in tools): + return _UNNAMED_TOOL_MODEL_REFUSAL + return None + + +def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests | _RefuseClientFrame: + if isinstance(frame_data, str) and not _may_name_a_model(frame_data): + return _NO_CLIENT_FRAME_REQUESTS + try: + frame: Final = _json_object_frame(frame_data) + except (RecursionError, ValueError): + return _UNREADABLE_CLIENT_FRAME + if frame is None: + return _NO_CLIENT_FRAME_REQUESTS + if _names_an_unroutable_lane(frame): + return _UNROUTABLE_LANE_REFUSAL + turn_models: Final = _placed_models(frame, _TURN_MODEL_PLACEMENTS) + tools: Final = _placed_tools(frame) + tool_models: Final = (model for model in (tool.get("model") for tool in tools) if isinstance(model, str)) + return _ClientFrameRequests( + starts_turn=frame.get("type") == _TURN_STARTING_FRAME, + lane=_frame_lane(frame), + turn_model=next(iter(turn_models), None), + models=tuple( + dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS), *tool_models)) + ), + unverifiable_model_refusal=_unverifiable_model_refusal(frame, tools), + ) + + +def _unverifiable_model_probe_refusal(probe: WebsocketFrameModelRefused, message: str) -> WebsocketFrameModelRefused: + if probe.error_type.endswith("_model_access_denied"): + return WebsocketFrameModelRefused(message=message, error_type=probe.error_type) + return probe + + +async def _client_frame_verdict( + websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None +) -> _ForwardClientFrame | _RefuseClientFrame: + if gate is None: + return _FORWARD_WITHOUT_TURN + requests: Final = _client_frame_requests(frame_data) + if isinstance(requests, _RefuseClientFrame): + return requests + verdicts: Final = tuple([await gate(websocket, model) for model in requests.models]) + refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) + if refusal is not None: + return _RefuseClientFrame(refusal=refusal) + if requests.unverifiable_model_refusal is not None and isinstance( + probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused + ): + return _RefuseClientFrame(refusal=_unverifiable_model_probe_refusal(probe, requests.unverifiable_model_refusal)) + allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed)) + return _ForwardClientFrame( + starts_turn=requests.starts_turn, + lane=requests.lane, + turn=( + _TurnAuthorization(requested_model=requests.turn_model, access_groups=allowed[0].access_groups) + if requests.turn_model is not None and allowed + else None + ), + ) + + +def _refusal_close_reason(message: str) -> str: + first_sentence: Final = _FIRST_SENTENCE.match(message) + return truncated_close_reason(message if first_sentence is None else first_sentence.group(0)) + + +async def refuse_websocket(websocket: WebSocket, refusal: WebsocketFrameModelRefused) -> None: + await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type)) + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) + + +def _stamp_billed_turn(kwargs: Mapping[str, object], turn: _TurnAuthorization) -> None: + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) + metadata: Final = None if litellm_params is None else litellm_params.get("metadata") + if not isinstance(metadata, MutableMapping): + return + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(turn.access_groups) + if turn.requested_model is not None: + metadata["model_group"] = turn.requested_model + + +def _opening_model(target: str) -> str | None: + return next(iter(parse_qs(urlparse(target).query).get("model", ())), None) + + +class _WebsocketTurnLedger: + def __init__(self, opening: _TurnAuthorization, started: datetime) -> None: + self.opening = opening + self.queued_turns: Mapping[str | None, tuple[_TurnAuthorization | None, ...]] = MappingProxyType({}) + self.session_opened = False + self.started = started + self.logged_turns = 0 + self.unstarted_turns = 0 + + def current(self, lane: str | None) -> _TurnAuthorization: + oldest: Final = next(iter(self.queued_turns.get(lane, ())), None) + return self.opening if oldest is None else oldest + + def queue_turn(self, lane: str | None, turn: _TurnAuthorization | None) -> None: + self.unstarted_turns += 1 + if not self.session_opened: + self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*self.queued_turns.get(lane, ()), turn)}) + + def open_session(self) -> None: + self.session_opened = True + self.queued_turns = MappingProxyType({}) + + def close_turn(self, lane: str | None, end_time: datetime) -> None: + self.logged_turns += 1 + self.started = end_time + self._forget_turn(lane, 0) + + def start_turn(self) -> None: + self.unstarted_turns = max(0, self.unstarted_turns - 1) + + def drop_refused_turn(self, lane: str | None, lane_is_running_a_turn: bool) -> None: + self.start_turn() + self._forget_turn(lane, 1 if lane_is_running_a_turn else 0) + + def _forget_turn(self, lane: str | None, position: int) -> None: + turns: Final = self.queued_turns.get(lane, ()) + self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*turns[:position], *turns[position + 1 :])}) + + +def _session_end_needs_a_row( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, + unlogged_frames: Sequence[Mapping[str, object]], + ledger: _WebsocketTurnLedger, +) -> bool: + if upstream_frame_kind is None: + return True + kinds: Final = frozenset(map(upstream_frame_kind, unlogged_frames)) + return "starts_turn" in kinds or (ledger.logged_turns == 0 and bool(kinds - {"opens_session"})) + + +def _turn_in_flight( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, unlogged_frames: Sequence[Mapping[str, object]] +) -> bool: + return upstream_frame_kind is not None and any( + upstream_frame_kind(frame) == "starts_turn" for frame in unlogged_frames + ) + + +def _turn_owed( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, + unlogged_frames: Sequence[Mapping[str, object]], + ledger: _WebsocketTurnLedger, +) -> bool: + return _turn_in_flight(upstream_frame_kind, unlogged_frames) or ( + upstream_frame_kind is not None and ledger.unstarted_turns > 0 + ) + + +def _lanes_in_flight( + upstream_frame_kind: WebsocketUpstreamFrameClassifier, unlogged_frames: Sequence[Mapping[str, object]] +) -> tuple[str | None, ...]: + started: Final = (frame for frame in unlogged_frames if upstream_frame_kind(frame) == "starts_turn") + return tuple(dict.fromkeys(map(_frame_lane, started))) + + +async def _relay_client_frame( + websocket: WebSocket, + upstream_ws: ClientConnection, + checked_frame: str | bytes, + forwarded_frame: str | bytes, + gate: WebsocketFrameModelGate | None, + ledger: _WebsocketTurnLedger, +) -> _ForwardClientFrame | _RefuseClientFrame: + verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate) + match verdict: + case _RefuseClientFrame(refusal=refusal): + await refuse_websocket(websocket, refusal) + case _ForwardClientFrame(starts_turn=True, turn=turn, lane=lane): + ledger.queue_turn(lane, turn) + await upstream_ws.send(forwarded_frame) + case _ForwardClientFrame(): + await upstream_ws.send(forwarded_frame) + return verdict + + +_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"}) +WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0 + + +def _metered_client_text(meter: WebSocketSessionMeter | None, client_text: str) -> int: + return 0 if meter is None else meter.measure(client_text) + + def _log_session_usage(logging_obj: LiteLLMLoggingObj, usage: WebSocketSessionUsage) -> None: logging_obj.model = usage.model # rebind-ok: the spend logger reads model and cost off the shared logging object logging_obj.model_call_details["model"] = usage.model # rebind-ok: same shared logging object @@ -2654,6 +3019,9 @@ async def websocket_passthrough_request( cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, + walkout_turn_drain_seconds: float = WALKOUT_TURN_DRAIN_SECONDS, builtin_provider_route: bool = False, model: str | None = None, meter: WebSocketSessionMeter | None = None, @@ -2670,6 +3038,11 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + client_frame_model_gate: Optional per-frame check of every model a client frame names + upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones + that open the session, so each turn logs its own spend and keeps the session's provider-named model + walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client + left or was refused, so the provider's usage for it still reaches spend builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider model: The model the caller named before connecting, logged on a failure before the setup frame arrives meter: Prices the session from the client text frames delivered upstream, for an upstream that reports no usage @@ -2709,26 +3082,31 @@ async def websocket_passthrough_request( proxy_config=proxy_config, route_description="websocket_passthrough", ) - logging_obj: Final = Logging( - model="unknown", - messages=[{"role": "user", "content": "WebSocket connection"}], - stream=True, # WebSockets are inherently streaming - call_type="pass_through_endpoint", - start_time=start_time, - litellm_call_id=litellm_call_id, - function_id="websocket_passthrough", - dynamic_success_callbacks=team_callbacks.success_callbacks, - dynamic_failure_callbacks=team_callbacks.failure_callbacks, - kwargs=team_callbacks.logging_kwargs, - ) - # Create passthrough logging payload - passthrough_logging_payload: Final = PassthroughStandardLoggingPayload( - url=target, - request_body={}, # WebSocket doesn't have a traditional request body - request_method="WEBSOCKET", - cost_per_request=cost_per_request, - ) + def new_call_logging(call_start: datetime, call_id: str) -> LiteLLMLoggingObj: + return Logging( + model="unknown", + messages=[{"role": "user", "content": "WebSocket connection"}], + stream=True, + call_type="pass_through_endpoint", + start_time=call_start, + litellm_call_id=call_id, + function_id="websocket_passthrough", + dynamic_success_callbacks=team_callbacks.success_callbacks, + dynamic_failure_callbacks=team_callbacks.failure_callbacks, + kwargs=team_callbacks.logging_kwargs, + ) + + def new_call_payload() -> PassthroughStandardLoggingPayload: + return PassthroughStandardLoggingPayload( + url=target, + request_body={}, + request_method="WEBSOCKET", + cost_per_request=cost_per_request, + ) + + logging_obj: Final = new_call_logging(start_time, litellm_call_id) + passthrough_logging_payload: Final = new_call_payload() # Create a dummy request object for WebSocket connections to maintain compatibility # with the existing _init_kwargs_for_pass_through_endpoint function @@ -2750,38 +3128,136 @@ async def websocket_passthrough_request( ) # Initialize kwargs for logging using the same pattern as HTTP passthrough - kwargs: Final = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( - user_api_key_dict=user_api_key_dict, - _parsed_body={}, # WebSocket doesn't have a traditional request body - passthrough_logging_payload=passthrough_logging_payload, - litellm_call_id=litellm_call_id, - request=dummy_request, - logging_obj=logging_obj, + def init_call_kwargs( + token: UserAPIKeyAuth, + call_payload: PassthroughStandardLoggingPayload, + call_id: str, + call_logging: LiteLLMLoggingObj, + ) -> dict[str, object]: # mutable-ok: the Vertex relay stamps the model from setup frames into these kwargs + return HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( + user_api_key_dict=token, + _parsed_body={}, # WebSocket doesn't have a traditional request body + passthrough_logging_payload=call_payload, + litellm_call_id=call_id, + request=dummy_request, + logging_obj=call_logging, + ) + + kwargs: Final = init_call_kwargs(user_api_key_dict, passthrough_logging_payload, litellm_call_id, logging_obj) + + def start_call_logging(call_logging: LiteLLMLoggingObj, call_kwargs: Mapping[str, object]) -> None: + call_logging.update_environment_variables( + model="unknown", + user="unknown", + optional_params={}, + litellm_params=dict(as_str_mapping(call_kwargs.get("litellm_params")) or {}), + call_type="pass_through_endpoint", + ) + call_logging.pre_call( + input=[{"role": "user", "content": "WebSocket connection"}], + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": target, + "headers": upstream_headers, + }, + ) + + start_call_logging(logging_obj, kwargs) + + class MockWebSocketRequest: + def __init__(self, target_url: str): + self.method = "WEBSOCKET" + self.url = target_url + + class MockWebSocketResponse: + def __init__(self, target_url: str): + self.status_code = 200 + self.text = "WebSocket connection successful" + self.headers: dict[str, str] = {} + self.request = MockWebSocketRequest(target_url) + + ledger: Final = _WebsocketTurnLedger( + opening=_TurnAuthorization( + requested_model=_opening_model(target), + access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), + ), + started=start_time, ) - # Update logging environment variables - logging_obj.update_environment_variables( - model="unknown", - user="unknown", - optional_params={}, - litellm_params=dict(kwargs.get("litellm_params", {})), - call_type="pass_through_endpoint", - ) - logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + def log_frames( + call_logging: LiteLLMLoggingObj, + call_payload: PassthroughStandardLoggingPayload, + call_kwargs: Mapping[str, object], + end_time: datetime, + lane: str | None, + in_row: Callable[[Mapping[str, object]], bool], + ) -> None: + all_frames: Final = list(websocket_messages) + frames: Final = [frame for frame in all_frames if in_row(frame)] + websocket_messages[:] = ( + () + if upstream_frame_kind is None + else ( + frame + for frame in all_frames + if not in_row(frame) or upstream_frame_kind(frame) in _SESSION_FRAMES_EVERY_ROW_CARRIES + ) + ) + logged_payload: Final[PassthroughStandardLoggingPayload] = { + **call_payload, + "response_body": frames, + "end_time": end_time, + } + if client_frame_model_gate is not None: + _stamp_billed_turn(call_kwargs, ledger.current(lane)) + success_kwargs: Final = { + key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS + } + custom_llm_provider: Final = call_kwargs.get("custom_llm_provider") + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler( + httpx_response=MockWebSocketResponse(target), + response_body=frames, + url_route=endpoint or "", + result="websocket_connection_successful", + start_time=ledger.started, + end_time=end_time, + logging_obj=call_logging, + cache_hit=False, + request_body={}, + passthrough_logging_payload=logged_payload, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + finalize_before_dispatch=False, + **success_kwargs, + ) + ) + bind_budget_reservation_to_callbacks(call_logging.litellm_params) - # Pre-call logging - logging_obj.pre_call( - input=[{"role": "user", "content": "WebSocket connection"}], - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": target, - "headers": upstream_headers, - }, - ) + def log_turn(lane: str | None, in_row: Callable[[Mapping[str, object]], bool]) -> None: + end_time: Final = datetime.now() + if upstream_frame_kind is None: + log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time, lane, in_row) + else: + turn_call_id: Final = str(uuid.uuid4()) + turn_logging: Final = new_call_logging(ledger.started, turn_call_id) + turn_payload: Final = new_call_payload() + turn_kwargs: Final = init_call_kwargs( + user_api_key_dict + if ledger.logged_turns == 0 + else user_api_key_dict.model_copy(update={"budget_reservation": None}), + turn_payload, + turn_call_id, + turn_logging, + ) + start_call_logging(turn_logging, turn_kwargs) + log_frames(turn_logging, turn_payload, turn_kwargs, end_time, lane, in_row) + ledger.close_turn(lane, end_time) ### CALL HOOKS ### - modify incoming data / reject request before calling the model - websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True) + websocket_hook_data: Final = guardrail_request_data_with_streaming( + MappingProxyType({"litellm_call_id": litellm_call_id}), is_streaming=True + ) await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=websocket_hook_data, @@ -2808,7 +3284,6 @@ async def websocket_passthrough_request( message = await websocket.receive() message_type = message.get("type") if message_type == "websocket.disconnect": - await upstream_ws.close() break text_data: str | None = message.get("text") @@ -2871,11 +3346,29 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) - if meter is not None: - metered_total += meter.measure(text_data) + if isinstance( + await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_model_gate, + ledger, + ), + _RefuseClientFrame, + ): + break + metered_total += _metered_client_text(meter, text_data) elif bytes_data is not None: - await upstream_ws.send(bytes_data) + if isinstance( + await _relay_client_frame( + websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate, ledger + ), + _RefuseClientFrame, + ): + break + if not _turn_owed(upstream_frame_kind, websocket_messages, ledger): + await upstream_ws.close() except asyncio.CancelledError: raise except Exception: @@ -2902,11 +3395,7 @@ async def websocket_passthrough_request( is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint) json_frame_ordinal: Final = count() - async def relay_upstream_frame(upstream_message: str | bytes) -> None: - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - else: - await websocket.send_text(upstream_message) + def record_upstream_frame(upstream_message: str | bytes) -> None: message_data: Final = _json_object_frame(upstream_message) if message_data is None: return @@ -2914,6 +3403,43 @@ async def websocket_passthrough_request( _extract_vertex_live_model_from_setup_response(message_data) return websocket_messages.append(message_data) + match None if upstream_frame_kind is None else upstream_frame_kind(message_data): + case "opens_session": + ledger.open_session() + case "starts_turn": + ledger.start_turn() + case "updates_session" if upstream_frame_kind is not None: + websocket_messages[:] = _without_earlier_session_updates( + upstream_frame_kind, websocket_messages, message_data + ) + case "ends_turn" if upstream_frame_kind is not None: + if _frame_lane(message_data) not in _lanes_in_flight(upstream_frame_kind, websocket_messages): + ledger.start_turn() + log_turn(_frame_lane(message_data), _frames_of_lane(_frame_lane(message_data))) + case "bills_usage" if not _turn_in_flight(upstream_frame_kind, websocket_messages): + log_turn(_frame_lane(message_data), _frames_of_lane(_frame_lane(message_data))) + case "refuses_turn" if upstream_frame_kind is not None: + lane: Final = _frame_lane(message_data) + ledger.drop_refused_turn( + lane, lane in _lanes_in_flight(upstream_frame_kind, websocket_messages) + ) + case _: + pass + + async def relay_upstream_frame(upstream_message: str | bytes) -> None: + record_upstream_frame(upstream_message) + if not _client_socket_is_open(websocket): + if not _turn_owed(upstream_frame_kind, websocket_messages, ledger): + await upstream_ws.close() + return + try: + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + else: + await websocket.send_text(upstream_message) + except WebSocketDisconnect: + if not _turn_owed(upstream_frame_kind, websocket_messages, ledger): + raise async def forward_upstream_to_client() -> Close | None: try: @@ -2922,6 +3448,11 @@ async def websocket_passthrough_request( except (ConnectionClosedOK, ConnectionClosedError) as e: verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e) return e.rcvd + except WebSocketDisconnect: + verbose_proxy_logger.debug( + "WebSocket passthrough (%s): client left while upstream was still sending", endpoint + ) + return None except asyncio.CancelledError: verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client") raise @@ -2933,12 +3464,15 @@ async def websocket_passthrough_request( raise # Create tasks for bidirectional message forwarding - tasks: Final = [ - asyncio.create_task(forward_client_to_upstream()), - asyncio.create_task(forward_upstream_to_client()), - ] + client_task: Final = asyncio.create_task(forward_client_to_upstream()) + upstream_task: Final = asyncio.create_task(forward_upstream_to_client()) + tasks: Final = (client_task, upstream_task) - done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + if client_task.done() and _turn_owed(upstream_frame_kind, websocket_messages, ledger): + await asyncio.wait((upstream_task,), timeout=walkout_turn_drain_seconds) + done: Final = tuple(task for task in tasks if task.done()) + pending: Final = tuple(task for task in tasks if not task.done()) # Cancel remaining tasks for task in pending: @@ -2961,64 +3495,17 @@ async def websocket_passthrough_request( reason=truncated_close_reason(upstream_close.reason), ) - end_time: Final = datetime.now() - - # Update passthrough logging payload with response data - passthrough_logging_payload["response_body"] = websocket_messages - passthrough_logging_payload["end_time"] = end_time - - # Remove logging_obj from kwargs to avoid duplicate keyword argument - success_kwargs: Final = kwargs.copy() - success_kwargs.pop("logging_obj", None) - - # # Add user authentication context for database logging - # if user_api_key_dict: - # success_kwargs.setdefault('litellm_params', {}) - # success_kwargs['litellm_params'].update({ - # 'proxy_server_request': { - # 'body': { - # 'user': user_api_key_dict.user_id, - # 'team_id': user_api_key_dict.team_id, - # 'end_user_id': user_api_key_dict.end_user_id, - # } - # } - # }) - # # Also add the user_api_key for direct access - # success_kwargs['user_api_key'] = user_api_key_dict.api_key - - # Create a dummy httpx.Response for WebSocket connections - class MockWebSocketResponse: - def __init__(self, target_url: str): - self.status_code = 200 - self.text = "WebSocket connection successful" - self.headers: dict[str, str] = {} - self.request = MockWebSocketRequest(target_url) - - class MockWebSocketRequest: - def __init__(self, target_url: str): - self.method = "WEBSOCKET" - self.url = target_url - - mock_response: Final = MockWebSocketResponse(target) if meter is not None: _log_session_usage(logging_obj, meter.price(metered_total)) - # Use the same success handler as HTTP passthrough endpoints - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler( - httpx_response=mock_response, - response_body=websocket_messages, - url_route=endpoint or "", - result="websocket_connection_successful", - start_time=start_time, - end_time=end_time, - logging_obj=logging_obj, - cache_hit=False, - request_body={}, - **success_kwargs, + if _session_end_needs_a_row(upstream_frame_kind, websocket_messages, ledger): + lanes_in_flight: Final = ( + () if upstream_frame_kind is None else _lanes_in_flight(upstream_frame_kind, websocket_messages) ) - ) - bind_budget_reservation_to_callbacks(logging_obj.litellm_params) + *earlier_lanes, last_lane = lanes_in_flight or (None,) + for lane in earlier_lanes: + log_turn(lane, _frames_of_lane(lane)) + log_turn(last_lane, _every_frame) # Call the proxy logging success hook if proxy_logging_obj: @@ -3076,6 +3563,7 @@ async def websocket_passthrough_request( if _client_socket_is_open(websocket): await websocket.close(code=1011, reason="WebSocket passthrough error") finally: + await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict) if _client_socket_is_open(websocket): await websocket.close() diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 8b09b1cf7cd..8f55906498b 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -1,5 +1,6 @@ import asyncio import json +from collections.abc import Mapping, Sequence from contextlib import suppress from datetime import datetime from types import MappingProxyType @@ -13,7 +14,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy._types import PassThroughEndpointLoggingResultValues +from litellm.proxy._types import PassThroughEndpointLoggingResultValues, PassThroughEndpointLoggingTypedDict from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) @@ -70,6 +71,14 @@ def _safe_response_text(httpx_response: httpx.Response) -> str: return "" +def _websocket_frames( + response_body: Mapping[str, object] | Sequence[Mapping[str, object]] | None, +) -> tuple[Mapping[str, object], ...]: + return tuple( + message for message in (response_body if isinstance(response_body, list) else ()) if isinstance(message, dict) + ) + + class PassThroughEndpointLogging: def __init__( self, @@ -120,6 +129,7 @@ class PassThroughEndpointLogging: # Vertex AI Live API WebSocket self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"] + self.TRACKED_OPENAI_WEBSOCKET_ROUTES: Final = ("/openai/", "/openai_passthrough/") @property def _log_dispatch(self) -> PassThroughLogDispatch: @@ -422,21 +432,18 @@ class PassThroughEndpointLogging: standard_logging_response_object = vertex_ai_live_handler_result["result"] kwargs = vertex_ai_live_handler_result["kwargs"] - elif DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route): - deepgram_handler_result: Final = ( - DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( - websocket_messages=tuple( - message - for message in (response_body if isinstance(response_body, list) else ()) - if isinstance(message, dict) - ), - logging_obj=logging_obj, - upstream_url=str(httpx_response.request.url), - kwargs=kwargs, - ) + elif ( + websocket_relay_result := self._websocket_relay_logging_result( + url_route=url_route, + websocket_frames=_websocket_frames(response_body), + logging_obj=logging_obj, + httpx_response=httpx_response, + start_time=start_time, + kwargs=kwargs, ) - standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain - kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract + ) is not None: + standard_logging_response_object = websocket_relay_result["result"] # rebind-ok: elif-chain + kwargs = websocket_relay_result["kwargs"] # rebind-ok: elif-chain contract elif FalAIPassthroughLoggingHandler.is_fal_ai_route(url_route, custom_llm_provider): fal_ai_handler_result: Final = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( response_body=response_body if isinstance(response_body, dict) else MappingProxyType({}), @@ -452,6 +459,36 @@ class PassThroughEndpointLogging: return_dict["kwargs"] = kwargs return return_dict + def _websocket_relay_logging_result( + self, + url_route: str, + websocket_frames: tuple[Mapping[str, object], ...], + logging_obj: LiteLLMLoggingObj, + httpx_response: httpx.Response, + start_time: datetime, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict | None: + if self.is_openai_websocket_route(url_route): + from .llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + OpenAIWebsocketPassthroughLoggingHandler, + ) + + return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( + websocket_messages=websocket_frames, + logging_obj=logging_obj, + url_route=url_route, + start_time=start_time, + kwargs=kwargs, + ) + if DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route): + return DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=websocket_frames, + logging_obj=logging_obj, + upstream_url=str(httpx_response.request.url), + kwargs=kwargs, + ) + return None + async def pass_through_async_success_handler( self, httpx_response: httpx.Response, @@ -625,6 +662,9 @@ class PassThroughEndpointLogging: return True return False + def is_openai_websocket_route(self, url_route: str) -> bool: + return url_route.startswith(self.TRACKED_OPENAI_WEBSOCKET_ROUTES) + def is_cursor_route(self, url_route: str, custom_llm_provider: str | None = None): """Check if the URL route is a Cursor Cloud Agents API route.""" if custom_llm_provider == "cursor": diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5fcbfb2564..6fcf789dc1b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -26,6 +26,7 @@ from collections.abc import ( Collection, Container, Iterable, + Iterator, Mapping, MutableMapping, Sequence, @@ -33,7 +34,7 @@ from collections.abc import ( from collections.abc import ( Set as AbstractSet, ) -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType, TracebackType, UnionType @@ -59,8 +60,11 @@ import anyio import anyio.lowlevel import websockets import websockets.exceptions +from openai.types.audio import TranscriptionStreamEvent from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config from pydantic.fields import FieldInfo, PydanticUndefined +from pydantic.main import IncEx +from pydantic_core import PydanticSerializationError, to_jsonable_python from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._uuid import uuid @@ -88,6 +92,7 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, + VLLM_CHOICE_FIELDS, ) from litellm.litellm_core_utils.access_windows import access_windows_config_error from litellm.litellm_core_utils.add_retry_fallback_headers import ( @@ -105,6 +110,7 @@ from litellm.litellm_core_utils.litellm_logging import ( initialize_sentry, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker +from litellm.litellm_core_utils.prompt_templates.chat_template_tools import chat_template_tool_section from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.routing_groups import parse_routing_groups from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -113,6 +119,7 @@ from litellm.litellm_core_utils.token_counter import ( UnsupportedContentTypeError, deployment_openai_model, offload_token_count, + sends_function_tools, ) from litellm.llms.bedrock.count_tokens.kimi_k3 import ( HUGGING_FACE_SOURCE, @@ -197,9 +204,11 @@ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.rust_bridge.catalog import Target from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( PRICING_OVERRIDES_KEY, + CallTypes, CountContextManagement, LlmProviders, ModelResponse, @@ -207,6 +216,7 @@ from litellm.types.utils import ( TextCompletionResponse, TokenCountMessagesFormat, TokenCountResponse, + TranscriptionResponse, echoed_cost_map_pricing_fields, is_server_derived_pricing_key, pricing_override_fields, @@ -316,6 +326,7 @@ from litellm._internal_context import service_target, with_service_target from litellm._logging import ( STARTUP_OUTPUT_SETTINGS, print_startup_message, + redact_internal_details_from_client_message, redact_string, verbose_proxy_logger, verbose_router_logger, @@ -365,6 +376,11 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) +from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( + TRANSCRIPTION_STREAM_DONE_FRAME, + GuardedTranscriptionReplay, + guarded_transcription_events, +) from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( as_str_mapping, @@ -387,6 +403,12 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_sensitive_keys, ) from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact +from litellm.llms.anthropic.messages.context_management.inference_view import ( + InferenceRequestView, + InferenceRequestViewError, + count_context_edits, + inference_request_view, +) from litellm.llms.anthropic.messages.convertible_blocks import ( deployment_drops_params, drop_non_string_tool_result_texts, @@ -397,6 +419,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai.live import parse_live_session_start from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.messages.dispatch import messages_target from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling @@ -456,6 +479,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401, RUF100 # lega user_api_key_auth, user_api_key_auth_accepted_websocket_for_model, user_api_key_auth_live_websocket, + user_api_key_auth_realtime_websocket, user_api_key_auth_websocket, ) from litellm.proxy.batches_endpoints.endpoints import router as batches_router @@ -473,12 +497,18 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l is_azure_model_router_request, log_llm_api_exception, open_sse_before_first_byte, + proxy_exception_from_http_exception, proxy_exception_from_route_error, resolve_litellm_call_id, should_return_raw_model_name, sse_error_payload, ttft_keepalive_interval, ) +from litellm.proxy.common_utils.anthropic_error_payload import ( + anthropic_error_json_response, + anthropic_internal_server_error_response, + uses_anthropic_error_envelope, +) from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, ) @@ -515,6 +545,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401, RUF10 resolve_inference_model, safe_get_request_headers, ) +from litellm.proxy.common_utils.http_parsing_utils import ( + resolve_realtime_route_model as _resolve_realtime_route_model, +) from litellm.proxy.common_utils.load_config_utils import get_config_from_bucket from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( @@ -534,9 +567,11 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.common_utils.openai_error_payload import ( LITELLM_CALL_ID_HEADER, attribute_of, + error_status_code, headers_with_litellm_call_id, litellm_call_id_headers, openai_error_param, + openai_error_type, proxy_exception_for, route_failure_proxy_exception, with_litellm_call_id, @@ -575,6 +610,7 @@ from litellm.proxy.common_utils.token_count_log import ( ) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + budget_windows_absent_cache_key, end_user_cache_key, get_management_object_ttl, model_access_group_cache_key, @@ -997,6 +1033,7 @@ from litellm.proxy.video_endpoints.endpoints import router as video_router from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.responses.main import provider_stores_the_conversation from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.router import ( AssistantsTypedDict, @@ -1148,6 +1185,27 @@ server_root_path: Final = get_server_root_path() _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() premium_user_data: Optional["EnterpriseLicenseData"] = _license_check.airgapped_license_data + + +def _set_premium_user(is_premium: bool) -> None: + global premium_user, premium_user_data + premium_user = is_premium # rebind-ok: config load can switch the license in use + premium_user_data = _license_check.airgapped_license_data # rebind-ok: tracks the license in use + + +def _use_config_license(config_license: str, env_license_is_premium: bool) -> bool: + env_license: Final = _license_check.license_str + env_license_data: Final = _license_check.airgapped_license_data + _license_check.license_str = config_license + if _license_check.is_premium(): + return True + if not env_license_is_premium: + return False + _license_check.license_str = env_license + _license_check.airgapped_license_data = env_license_data + return True + + global_max_parallel_request_retries_env: Final[str | None] = os.getenv("LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES") proxy_state: Final = ProxyState() SENSITIVE_DATA_MASKER: Final = SensitiveDataMasker() @@ -2258,6 +2316,8 @@ async def openai_exception_handler(request: Request, exc: ProxyException): # NOTE: DO NOT MODIFY THIS, its crucial to map to Openai exceptions log_model_access_denial(exc) headers: Final = exc.headers + if uses_anthropic_error_envelope(request): + return anthropic_error_json_response(exc, request) error_dict: Final = with_call_id( JSON_OBJECT.validate_python(exc.to_dict()), error_body_call_id(general_settings_view(), headers.get(LITELLM_CALL_ID_HEADER)), @@ -2293,6 +2353,10 @@ async def otlp_http_exception_handler(connection: Request | WebSocket, exc: Star if response is not None: _close_dangling_otel_server_span(connection, exc.status_code, exc=exc) return response + if uses_anthropic_error_envelope(connection): + return anthropic_error_json_response( + proxy_exception_from_http_exception(exc, dict(exc.headers or {})), connection + ) return await http_exception_handler(connection, exc) @@ -2500,6 +2564,8 @@ async def _unhandled_exception_response(request: Request, exc: Exception) -> Res otlp_response: Final = tracing_endpoints.otlp_error_response(request, 500) if otlp_response is not None: return otlp_response + if uses_anthropic_error_envelope(request): + return anthropic_internal_server_error_response(request) return JSONResponse( status_code=500, content={ @@ -3622,6 +3688,52 @@ def _cached_budget_limits(cached: object) -> object: return _UNTYPED_VALUE.validate_python(json.loads(stored)) if isinstance(stored, str) else stored +def _has_budget_windows(budget_limits: object) -> bool: + return isinstance(budget_limits, list) and budget_limits != [] + + +def _budget_windows_absent_marker(absent_key: str) -> object: + return _UNTYPED_VALUE.validate_python(user_api_key_cache.budget_windows_absent_cache.get_cache(key=absent_key)) + + +def _finish_pending_read(absent_key: str, pending_read: str, remember_absent: bool) -> None: + if _budget_windows_absent_marker(absent_key) != pending_read: + return + if remember_absent: + user_api_key_cache.budget_windows_absent_cache.set_cache(key=absent_key, value=True) + return + user_api_key_cache.budget_windows_absent_cache.delete_cache(absent_key) + + +async def _stored_budget_limits(entity_type: Litellm_EntityType, entity_id: str) -> object: + """Read only: caching the row here could race a delete and put a just-deleted key back into the auth cache.""" + if prisma_client is None: + return None + absent_key: Final = budget_windows_absent_cache_key(entity_type.value, entity_id) + if _budget_windows_absent_marker(absent_key) is True: + return None + pending_read: Final = f"reading:{secrets.token_hex(8)}" + user_api_key_cache.budget_windows_absent_cache.set_cache(key=absent_key, value=pending_read) + try: + row: Final[SupportsModelDump | None] = ( + await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": entity_id}) + if entity_type == Litellm_EntityType.KEY + else await TeamRepository(prisma_client).table.find_unique(where={"team_id": entity_id}) + ) + except Exception as e: # noqa: BLE001 # spend tracking must never fail the cost callback + verbose_proxy_logger.warning( + "Unable to read budget windows for %s=%s; this request's window spend is not counted: %s", + entity_type.value, + entity_id, + e, + ) + _finish_pending_read(absent_key, pending_read, remember_absent=False) + return None + budget_limits: Final[object] = None if row is None else _cached_budget_limits(row) + _finish_pending_read(absent_key, pending_read, remember_absent=not _has_budget_windows(budget_limits)) + return budget_limits + + async def _increment_spend_counters_batched( token: str | None, team_id: str | None, @@ -3652,7 +3764,68 @@ async def _increment_spend_counters_batched( cost: Final[float] = response_cost - async def _key_scope(key_token: str) -> tuple[PendingSpendIncrement | BaseException, ...]: + async def _window_increment( + entity_type: Litellm_EntityType, entity_id: str, window: object + ) -> PendingSpendIncrement | None: + duration: Final = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + window_reset_at: Final = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + window_counter: Final = f"spend:{entity_type.value}:{entity_id}:window:{duration}" + window_start: Final = get_budget_window_start(window) + pending_window: Final = ( + await _prepare_window_spend_counter_increment( + counter_key=window_counter, + entity_type=entity_type.value.capitalize(), + entity_id=entity_id, + window_duration=duration, + window_start=window_start, + increment=cost, + ) + if window_counter not in reserved_counter_keys + else None + ) + await _enqueue_window_spend_row_update( + entity_type=entity_type, + entity_id=entity_id, + reset_at=window_reset_at, + window_duration=duration, + window_start=window_start, + increment=cost, + request_started_at=request_started_at, + ) + return pending_window + + async def _window_increments( + entity_type: Litellm_EntityType, entity_id: str, budget_limits: object + ) -> tuple[PendingSpendIncrement | BaseException, ...]: + if not isinstance(budget_limits, list): + return () + window_pending: Final = await asyncio.gather( + *(_window_increment(entity_type, entity_id, window) for window in budget_limits), return_exceptions=True + ) + return tuple(item for item in window_pending if item is not None) + + async def _cached_window_increments( + entity_type: Litellm_EntityType, entity_id: str, cache_key: str + ) -> tuple[_SpendScopeItem, ...]: + cached: Final[object] = _UNTYPED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key)) + if cached is None: + return (_UncachedBudgetWindows(entity_type=entity_type, entity_id=entity_id),) + return await _window_increments(entity_type, entity_id, _cached_budget_limits(cached)) + + async def _stored_window_increments( + uncached: _UncachedBudgetWindows, + ) -> tuple[PendingSpendIncrement | BaseException, ...]: + return await _window_increments( + uncached.entity_type, + uncached.entity_id, + await _stored_budget_limits(uncached.entity_type, uncached.entity_id), + ) + + async def _key_scope(key_token: str) -> tuple[_SpendScopeItem, ...]: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow # hashes raw "sk-..." keys before they reach the callback). The # startswith("sk-") check is a safety net matching update_cache — @@ -3675,52 +3848,9 @@ async def _increment_spend_counters_batched( ) ) - async def _key_window_increment(window: object) -> PendingSpendIncrement | None: - duration = ( - window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) - ) - key_window_reset_at = ( - window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) - ) - key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}" - key_window_start = get_budget_window_start(window) - pending_window: Final = ( - await _prepare_window_spend_counter_increment( - counter_key=key_window_counter, - entity_type="Key", - entity_id=hashed_token, - window_duration=duration, - window_start=key_window_start, - increment=cost, - ) - if key_window_counter not in reserved_counter_keys - else None - ) - await _enqueue_window_spend_row_update( - entity_type=Litellm_EntityType.KEY, - entity_id=hashed_token, - reset_at=key_window_reset_at, - window_duration=duration, - window_start=key_window_start, - increment=cost, - request_started_at=request_started_at, - ) - return pending_window + return key_pending + await _cached_window_increments(Litellm_EntityType.KEY, hashed_token, hashed_token) - key_obj: Final[object] = _UNTYPED_VALUE.validate_python( - await user_api_key_cache.async_get_cache(key=hashed_token) - ) - if key_obj is None: - return key_pending - key_budget_limits: Final = _cached_budget_limits(key_obj) - if not isinstance(key_budget_limits, list): - return key_pending - window_pending: Final = await asyncio.gather( - *(_key_window_increment(window) for window in key_budget_limits), return_exceptions=True - ) - return key_pending + tuple(item for item in window_pending if item is not None) - - async def _team_scope(scope_team_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]: + async def _team_scope(scope_team_id: str) -> tuple[_SpendScopeItem, ...]: team_counter_key: Final = f"spend:team:{scope_team_id}" team_pending: Final[tuple[PendingSpendIncrement, ...]] = ( () @@ -3734,50 +3864,9 @@ async def _increment_spend_counters_batched( ) ) - async def _team_window_increment(window: object) -> PendingSpendIncrement | None: - duration = ( - window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) - ) - team_window_reset_at = ( - window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) - ) - team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}" - team_window_start = get_budget_window_start(window) - pending_window: Final = ( - await _prepare_window_spend_counter_increment( - counter_key=team_window_counter, - entity_type="Team", - entity_id=scope_team_id, - window_duration=duration, - window_start=team_window_start, - increment=cost, - ) - if team_window_counter not in reserved_counter_keys - else None - ) - await _enqueue_window_spend_row_update( - entity_type=Litellm_EntityType.TEAM, - entity_id=scope_team_id, - reset_at=team_window_reset_at, - window_duration=duration, - window_start=team_window_start, - increment=cost, - request_started_at=request_started_at, - ) - return pending_window - - team_obj: Final[object] = _UNTYPED_VALUE.validate_python( - await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") + return team_pending + await _cached_window_increments( + Litellm_EntityType.TEAM, scope_team_id, f"team_id:{scope_team_id}" ) - if team_obj is None: - return team_pending - team_budget_limits: Final = _cached_budget_limits(team_obj) - if not isinstance(team_budget_limits, list): - return team_pending - window_pending: Final = await asyncio.gather( - *(_team_window_increment(window) for window in team_budget_limits), return_exceptions=True - ) - return team_pending + tuple(item for item in window_pending if item is not None) async def _team_member_scope( scope_user_id: str, scope_team_id: str @@ -3848,22 +3937,30 @@ async def _increment_spend_counters_batched( # return_exceptions so a failing scope does not leave its siblings running # as orphaned tasks that race the caller's reservation-counter invalidation; # all scopes settle, then the first error propagates as before. - scope_results: Final = await asyncio.gather(*scope_coros, return_exceptions=True) - scope_errors: Final = tuple( - item - for scope in scope_results - for item in (scope if isinstance(scope, tuple) else (scope,)) - if isinstance(item, BaseException) + scope_items: Final = tuple(_flatten_spend_scopes(await asyncio.gather(*scope_coros, return_exceptions=True))) + uncached_windows: Final = tuple(item for item in scope_items if isinstance(item, _UncachedBudgetWindows)) + await _apply_spend_counter_increments( + pending=reservation_update.pending + + tuple(item for item in scope_items if isinstance(item, PendingSpendIncrement)) ) - pending: Final = tuple( - item - for scope in scope_results - if not isinstance(scope, BaseException) - for item in scope - if not isinstance(item, BaseException) + stored_window_items: Final = ( + tuple( + _flatten_spend_scopes( + await asyncio.gather( + *(_stored_window_increments(uncached) for uncached in uncached_windows), return_exceptions=True + ) + ) + ) + if uncached_windows + else () ) - await _apply_spend_counter_increments(pending=reservation_update.pending + pending) + stored_window_pending: Final = tuple( + item for item in stored_window_items if isinstance(item, PendingSpendIncrement) + ) + if stored_window_pending: + await _apply_spend_counter_increments(pending=stored_window_pending) stamp_budget_reservation_actual_cost(budget_reservation=budget_reservation, actual_cost=response_cost) + scope_errors: Final = tuple(item for item in scope_items + stored_window_items if isinstance(item, BaseException)) if scope_errors: raise scope_errors[0] @@ -3871,6 +3968,25 @@ async def _increment_spend_counters_batched( budget_reservation["finalized"] = True +@dataclass(frozen=True, slots=True) +class _UncachedBudgetWindows: + entity_type: Litellm_EntityType + entity_id: str + + +_SpendScopeItem: TypeAlias = PendingSpendIncrement | BaseException | _UncachedBudgetWindows + + +def _flatten_spend_scopes( + scope_results: Sequence[tuple[_SpendScopeItem, ...] | BaseException], +) -> Iterator[_SpendScopeItem]: + for scope in scope_results: + if isinstance(scope, BaseException): + yield scope + continue + yield from scope + + @dataclass(frozen=True, slots=True) class _ReservationCounterUpdate: """The reserved counters the direct increment must skip, and the adjustments that settle them on the actual @@ -5651,11 +5767,18 @@ def _with_config_file_pass_through_endpoints( def _reload_settings_store(section: Section, store: SettingsStore, section_config: object) -> None: serving_pass_throughs: Final = store.get("pass_through_endpoints") store.load_yaml(_as_settings_mapping(section_config)) - store.apply_db_row(section, _EMPTY_SETTINGS_MAPPING) if is_resource_list(section, "pass_through_endpoints") and serving_pass_throughs is not None: store["pass_through_endpoints"] = serving_pass_throughs +def _general_settings_row_is_readable(row: _ConfigParamValueRow | None) -> bool: + return row is None or row.param_value is None or isinstance(row.param_value, Mapping) + + +def _db_managed_general_settings_keys() -> frozenset[str]: + return frozenset((*ConfigGeneralSettings.model_fields, *_CLEANUP_SCHEDULE_KEYS, "allowed_ips", "email_settings")) + + def _bind_general_settings_store(settings: SettingsStore) -> None: global general_settings general_settings = settings # pyright: ignore[reportAssignmentType] # legacy global accepts mappings @@ -5806,7 +5929,9 @@ class ProxyConfig: # whether an existing request predates the prices it just fetched, and re-serving one # costs a single fetch where skipping one leaves it priced wrong indefinitely self.model_cost_map_applied_revision: int = 0 - self.settings: Final[SettingsStore] = SettingsStore("general_settings") + self.settings: Final[SettingsStore] = SettingsStore( + "general_settings", managed_db_keys=_db_managed_general_settings_keys() + ) self.router_settings: Final[SettingsStore] = SettingsStore("router_settings") self.litellm_settings: Final[SettingsStore] = SettingsStore("litellm_settings") self.environment_variables: Final[SettingsStore] = SettingsStore("environment_variables") @@ -5817,6 +5942,8 @@ class ProxyConfig: self._db_environment_keys: frozenset[str] = frozenset() self._sso_environment_removals: frozenset[str] = frozenset() self._warned_shadowed_keys: frozenset[tuple[Section, str]] = frozenset() + self._warned_unmanaged_keys: frozenset[str] = frozenset() + self._warned_unreadable_db_general_settings: bool = False self._db_config_callback_entries: Mapping[tuple[str, str], _DbCallbackRegistration] = MappingProxyType({}) self._settings_stores: Final[Mapping[Section, SettingsStore]] = MappingProxyType( { @@ -6462,6 +6589,9 @@ class ProxyConfig: prisma_client=prisma_client, store_model_in_db=store_model_in_db, ) + else: + for section, store in self._settings_stores.items(): + store.apply_db_row(section, _EMPTY_SETTINGS_MAPPING) config = self._config_with_resolved_settings(config) @@ -6577,7 +6707,6 @@ class ProxyConfig: def _load_environment_variables(self, config: dict): ## ENVIRONMENT VARIABLES - global premium_user environment_variables: Final = config.get("environment_variables", None) if environment_variables: for key, value in environment_variables.items(): @@ -6617,7 +6746,7 @@ class ProxyConfig: # check if litellm_license in general_settings if "LITELLM_LICENSE" in environment_variables: _license_check.license_str = os.getenv("LITELLM_LICENSE", None) - premium_user = _license_check.is_premium() + _set_premium_user(_license_check.is_premium()) def _set_config_environment_variable(self, key: str, value: str) -> None: if key in os.environ and os.environ[key] != value: @@ -6669,7 +6798,6 @@ class ProxyConfig: prompt_injection_detection_obj, \ redis_usage_cache, \ store_model_in_db, \ - premium_user, \ open_telemetry_logger, \ health_check_details, \ proxy_batch_polling_interval, \ @@ -6682,6 +6810,12 @@ class ProxyConfig: self._load_environment_variables(config=config) + config_license: Final = _GENERAL_SETTINGS_VIEW.validate_python( + _GENERAL_SETTINGS_VIEW.validate_python(config).get("general_settings") or {} + ).get("litellm_license") + if isinstance(config_license, str): + _set_premium_user(_use_config_license(config_license, env_license_is_premium=premium_user)) + ## Coordination Redis (before cache init, so the explicit block wins) coordination_redis_cache: Final = self._init_coordination_redis(config=config) if coordination_redis_cache is not None: @@ -7258,7 +7392,7 @@ class ProxyConfig: allowed_ips: Final = general_settings.get("allowed_ips", None) if allowed_ips is not None and premium_user is False: raise ValueError( - "allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your envionment." + "allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment." ) ## BUDGET RESCHEDULER ## proxy_budget_rescheduler_min_time = general_settings.get( @@ -7314,11 +7448,6 @@ class ProxyConfig: if general_settings.get("enforced_params") is not None and premium_user is not True: raise ValueError("Trying to use `enforced_params`" + CommonProxyErrors.not_premium_user.value) - # check if litellm_license in general_settings - if "litellm_license" in general_settings: - _license_check.license_str = general_settings["litellm_license"] - premium_user = _license_check.is_premium() - router_params: Final[dict] = { "cache_responses": litellm.cache is not None, # cache if user passed in cache values } @@ -8340,15 +8469,20 @@ class ProxyConfig: general_settings.get("maximum_spend_logs_retention_interval", "1d"), ) - async def _update_general_settings(self, db_general_settings: Mapping[str, SettingsJsonValue] | None) -> None: + async def _update_general_settings(self, db_general_settings: Mapping[str, SettingsJsonValue] | str | None) -> None: global general_settings if db_general_settings is None: return + if not isinstance(db_general_settings, Mapping): + self._warn_about_unreadable_db_general_settings() + return + self._warned_unreadable_db_general_settings = False if not isinstance(general_settings, SettingsStore): self.settings.load_yaml(_as_settings_mapping(general_settings)) cache_size_was_db: Final = self.settings.source("user_api_key_cache_max_size") == "db" previous_cleanup_schedule: Final = self._resolved_cleanup_schedule() self.settings.apply_db_row("general_settings", db_general_settings) + self._warn_about_unmanaged_db_settings() _bind_general_settings_store(self.settings) await self._apply_general_settings_side_effects( db_general_settings, @@ -8540,12 +8674,24 @@ class ProxyConfig: environment_read: Final = next(self._environment_read_numbers) sections: Final = tuple(self._settings_stores) responses: Final = await asyncio.gather(*(get_config_param(prisma_client, section) for section in sections)) + general_settings_response: Final = responses[sections.index("general_settings")] + general_settings_value: Final[Mapping[str, SettingsJsonValue]] = ( + self._prepared_db_settings_values("general_settings", general_settings_response.param_value) + if general_settings_response is not None and general_settings_response.param_value is not None + else _EMPTY_SETTINGS_MAPPING + ) + self._settings_stores["general_settings"].apply_db_row("general_settings", general_settings_value) + response: _ConfigParamValueRow | None for section, response in zip(sections, responses): + if section == "general_settings": + continue if response is None: + self._settings_stores[section].apply_db_row(section, _EMPTY_SETTINGS_MAPPING) continue param_value: JsonValue = getattr(response, "param_value", None) if param_value is None: + self._settings_stores[section].apply_db_row(section, _EMPTY_SETTINGS_MAPPING) continue verbose_proxy_logger.debug( @@ -8556,6 +8702,7 @@ class ProxyConfig: if section == "litellm_settings": self._apply_litellm_settings_db_values(self._prepared_db_settings_values(section, param_value)) elif section == "environment_variables": + self._settings_stores[section].apply_db_row(section, _EMPTY_SETTINGS_MAPPING) await self._apply_db_environment_variables_row( prisma_client, response, environment_before_db, environment_read ) @@ -8566,6 +8713,7 @@ class ProxyConfig: ) self._warn_about_shadowed_db_settings() + self._warn_about_unmanaged_db_settings() return self._config_with_resolved_settings(config) async def _apply_db_environment_variables_row( @@ -8624,6 +8772,30 @@ class ProxyConfig: ) self._warned_shadowed_keys = shadowed + def _warn_about_unmanaged_db_settings(self) -> None: + unmanaged: Final = frozenset( + key for key in self.settings.unmanaged_db_keys() if key not in _GENERAL_SETTINGS_UI_LITELLM_FIELDS + ) + for key in sorted(unmanaged - self._warned_unmanaged_keys): + verbose_proxy_logger.warning( + "general_settings.%s is stored in the database, but the proxy does not apply that stored value, so it " + "is ignored. %s", + key, + "The config file's value applies. Remove the stored copy with POST /config/field/delete." + if self.settings.owned_by_config(key) + else "Set it in the config file to use it, or remove it with POST /config/field/delete.", + ) + self._warned_unmanaged_keys = unmanaged + + def _warn_about_unreadable_db_general_settings(self) -> None: + if self._warned_unreadable_db_general_settings: + return + verbose_proxy_logger.warning( + "general_settings is stored in the database as something other than a JSON object, so the proxy " + "ignores every stored general setting." + ) + self._warned_unreadable_db_general_settings = True + def _prepared_db_settings_values(self, section: Section, value: object) -> Mapping[str, SettingsJsonValue]: scrubbed: Final = _scrub_db_overlay_remote_module_loads(section=section, db_value=value) return _as_settings_mapping(scrubbed) @@ -8791,7 +8963,10 @@ class ProxyConfig: ) # initialize vector stores, guardrails, etc. table in db - await self._init_non_llm_objects_in_db(prisma_client=prisma_client) + await self._init_non_llm_objects_in_db( + prisma_client=prisma_client, + refresh_pass_through_endpoints=_general_settings_row_is_readable(db_general_settings), + ) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - %s", e) @@ -8860,7 +9035,9 @@ class ProxyConfig: except Exception as e: # noqa: BLE001 # best-effort: a failing stop must not break proxy shutdown verbose_proxy_logger.error("Error stopping auth cache invalidation subscriber: %s", e) - async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): + async def _init_non_llm_objects_in_db( + self, prisma_client: PrismaClient, *, refresh_pass_through_endpoints: bool = True + ): """ Use this to read non-llm objects from the db and initialize them @@ -8884,7 +9061,7 @@ class ProxyConfig: if self._should_load_db_object(object_type="agents"): await self._init_agents_in_db(prisma_client=prisma_client) - if self._should_load_db_object(object_type="pass_through_endpoints"): + if refresh_pass_through_endpoints and self._should_load_db_object(object_type="pass_through_endpoints"): await self._init_pass_through_endpoints_in_db() if self._should_load_db_object(object_type="prompts"): @@ -10186,6 +10363,8 @@ def _event_with_restamped_nested_model(event: Mapping[str, object], target_model _SSE_DATA_FIELD: Final = b"data:" _SSE_NESTED_RESPONSE_MARKER: Final = b'"response"' _SSE_EVENT_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_JSON_OBJECT: Final = TypeAdapter(dict[object, object]) +_JSON_ARRAY: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) def _restamp_nested_response_model_in_sse_frame(frame: bytes, target_model: str) -> bytes: @@ -10211,91 +10390,110 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) - return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending -def _fast_serialize_simple_model_response_stream( - chunk: ModelResponseStream, -) -> bytes | None: - """ - Serialize the common OpenAI text streaming chunk without the full Pydantic - serializer. Fall back for richer chunks so tool calls, logprobs, usage, and - provider-specific fields keep the canonical model_dump_json behavior. - """ - if ( - getattr(chunk, "provider_specific_fields", None) is not None - or getattr(chunk, "system_fingerprint", None) is not None - or getattr(chunk, "usage", None) is not None - ): - return None - - choices: Final = getattr(chunk, "choices", None) - if not isinstance(choices, list) or len(choices) != 1: - return None - - choice: Final = choices[0] - if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None: - return None - - delta: Final[object] = getattr(choice, "delta", None) - if delta is None: - return None - - unsupported_delta_fields: Final = ( - "function_call", - "tool_calls", - "audio", - "images", - "annotations", - "reasoning_content", - "thinking_blocks", - "provider_specific_fields", - "refusal", - ) - if any(getattr(delta, field, None) is not None for field in unsupported_delta_fields): - return None - - delta_dict: Final[dict] = {} - role: Final = getattr(delta, "role", None) - content: Final = getattr(delta, "content", None) - if role is not None: - delta_dict["role"] = role - if content is not None: - delta_dict["content"] = content - - choice_dict: Final = {"index": getattr(choice, "index", 0), "delta": delta_dict} - finish_reason: Final = getattr(choice, "finish_reason", None) - if finish_reason is not None: - choice_dict["finish_reason"] = finish_reason - - # Match the canonical ``model_dump_json(exclude_none=True)`` shape — if a - # field is None, omit it entirely rather than emitting ``"key": null``. - # Strict OpenAI-compatible clients reject ``null`` for optional fields like - # ``model``, so diverging here would surface as a client-side regression - # only on the fast path. Fall back to the slow path if a required-looking - # top-level identifier is missing. - model: Final = getattr(chunk, "model", None) - if model is None: - return None - - payload: Final[dict] = { - "id": getattr(chunk, "id", None), - "object": getattr(chunk, "object", None), - "created": getattr(chunk, "created", None), - "model": model, - "service_tier": getattr(chunk, "service_tier", None), - "choices": [choice_dict], - } - for top_level_key in ("id", "object", "created", "service_tier"): - if payload[top_level_key] is None: - payload.pop(top_level_key) - return orjson.dumps(payload) +def _vllm_choice_fields(choice: BaseModel) -> Mapping[str, object]: + choice_extras: Final = choice.model_extra or {} + return MappingProxyType({field: choice_extras[field] for field in VLLM_CHOICE_FIELDS if field in choice_extras}) -def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes: - if isinstance(chunk, ModelResponseStream): - serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk) +def _json_encodes(value: object) -> bool: + try: + to_jsonable_python(value) + except ValueError: + return False + return True + + +def _unencodable_part(value: object, ancestors: frozenset[int]) -> IncEx | bool: + if id(value) in ancestors: + return True + lineage: Final = ancestors | {id(value)} + if isinstance(value, BaseModel): + return _unencodable_keys(iter(value), lineage) or True + if isinstance(value, dict): + return _unencodable_entries(_JSON_OBJECT.validate_python(value), lineage) + if isinstance(value, (list, tuple)): + return _unencodable_indexes(enumerate(_JSON_ARRAY.validate_python(value)), lineage) or True + return True + + +def _unencodable_entries(entries: Mapping[object, object], lineage: frozenset[int]) -> IncEx | bool: + by_name: Final = {key: item for key, item in entries.items() if isinstance(key, str)} + if len(by_name) == len(entries): + return _unencodable_keys(by_name.items(), lineage) or True + by_index: Final = {key: item for key, item in entries.items() if isinstance(key, int)} + if len(by_index) == len(entries): + return _unencodable_indexes(by_index.items(), lineage) or True + return True + + +def _unencodable_keys(items: Iterable[tuple[str, object]], lineage: frozenset[int]) -> Mapping[str, IncEx | bool]: + return {key: _unencodable_part(item, lineage) for key, item in items if not _json_encodes(item)} + + +def _unencodable_indexes(items: Iterable[tuple[int, object]], lineage: frozenset[int]) -> Mapping[int, IncEx | bool]: + return {index: _unencodable_part(item, lineage) for index, item in items if not _json_encodes(item)} + + +def _serialize_chat_chunk_leniently(chunk: ModelResponseStream) -> bytes: + exclude: Final = _unencodable_keys(iter(chunk), frozenset({id(chunk)})) + if exclude: + verbose_proxy_logger.warning("Dropping streamed chunk fields that JSON cannot encode: %s", exclude) + jsonable: Final = chunk.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=exclude) + return json.dumps(jsonable, ensure_ascii=True, separators=(",", ":")).encode() + + +def _serialize_chat_chunk(chunk: ModelResponseStream) -> bytes: + try: + return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() + except PydanticSerializationError: + return _serialize_chat_chunk_leniently(chunk) + + +def _serialize_streaming_chunk(chunk: BaseModel) -> bytes: + choices: Final = _streaming_chunk_choices(chunk) + null_choice_fields: Final = tuple(_null_vllm_choice_fields(choice) for choice in choices) + if any(null_choice_fields): + serialized_chunk: Final = _serialize_keeping_null_choice_fields(chunk, choices, null_choice_fields) if serialized_chunk is not None: return serialized_chunk + if isinstance(chunk, ModelResponseStream): + return _serialize_chat_chunk(chunk) + return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() - return chunk.model_dump_json(exclude_none=True, exclude_unset=True) + +def _serialize_keeping_null_choice_fields( + chunk: BaseModel, choices: Sequence[BaseModel], null_choice_fields: Sequence[Mapping[str, None]] +) -> bytes | None: + try: + return json.dumps( + { + **_dump_set_fields(chunk, exclude=frozenset({"choices"})), + "choices": [ + {**_dump_set_fields(choice), **nulls} + for choice, nulls in zip(choices, null_choice_fields, strict=True) + ], + }, + separators=(",", ":"), + ensure_ascii=False, + ).encode() + except ValueError: + return None + + +def _streaming_chunk_choices(chunk: BaseModel) -> Sequence[BaseModel]: + match chunk: + case ModelResponseStream() | TextCompletionResponse(): + return chunk.choices + case _: + return () + + +def _null_vllm_choice_fields(choice: BaseModel) -> Mapping[str, None]: + return MappingProxyType({field: None for field, value in _vllm_choice_fields(choice).items() if value is None}) + + +def _dump_set_fields(model: BaseModel, exclude: AbstractSet[str] = frozenset()) -> Mapping[str, object]: + return model.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=set(exclude)) async def _apply_streaming_chunk_hooks( @@ -10551,6 +10749,26 @@ def _make_keepalive_resolver(request_data: Mapping[str, object]) -> Callable[[ob return _resolve +def _transcription_sse_error_frame(client_error: Exception) -> str: + status_code: Final = error_status_code(client_error, 500) + error_returned: Final = json.dumps( + { + "error": sse_error_payload(client_error)[1] + if isinstance(client_error, HTTPException) + else ProxyException( + message=redact_internal_details_from_client_message( + str(getattr(client_error, "message", client_error)) + ), + type=openai_error_type(client_error, status_code), + param=openai_error_param(client_error), + code=status_code, + ).to_dict() + }, + default=str, + ) + return f"data: {error_returned}\n\n" + + async def async_data_generator( response, user_api_key_dict: UserAPIKeyAuth, @@ -13411,7 +13629,10 @@ async def moderations( return response except Exception as e: await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=data, + dispatched_call_type=CallTypes.amoderation.value, ) log_llm_api_exception(e, litellm_call_id) if isinstance(e, ProxyException): @@ -13556,6 +13777,7 @@ async def audio_speech( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data, + dispatched_call_type=CallTypes.aspeech.value, ) log_llm_api_exception(e, litellm_call_id) if isinstance(e, ProxyException): @@ -13583,11 +13805,13 @@ def _form_fields_without_file(form_data: Mapping[str, object]) -> Mapping[str, o "/v1/audio/transcriptions", dependencies=[Depends(user_api_key_auth)], tags=["audio"], + responses={200: {"model": TranscriptionResponse, "content": {"text/event-stream": {"schema": {"type": "string"}}}}}, ) @router.post( "/audio/transcriptions", dependencies=[Depends(user_api_key_auth)], tags=["audio"], + responses={200: {"model": TranscriptionResponse, "content": {"text/event-stream": {"schema": {"type": "string"}}}}}, ) async def audio_transcriptions( request: Request, @@ -13675,39 +13899,80 @@ async def audio_transcriptions( cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" response_cost: Final = hidden_params.get("response_cost", None) or "" - response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or litellm_call_id additional_headers: Final[dict] = hidden_params.get("additional_headers", {}) or {} + stream_logging_obj: Final = getattr(response, "logging_obj", None) - fastapi_response.headers.update( - ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - call_id=response_call_id, - request_data=data, - hidden_params=hidden_params, - **additional_headers, - ) + custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + response_cost=response_cost, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + call_id=response_call_id, + request_data=data, + hidden_params=hidden_params, + litellm_logging_obj=stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None, + **additional_headers, ) # Call response headers hook (matches base_process_llm_request behavior) callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook( data=data, user_api_key_dict=user_api_key_dict, - response=response, + response=response.response if isinstance(response, GuardedTranscriptionReplay) else response, request_headers=dict(request.headers), ) - if callback_headers: - fastapi_response.headers.update(callback_headers) + response_headers: Final = {**custom_headers, **(callback_headers or {})} + + if hasattr(response, "__aiter__"): + + async def transcription_event_stream( + stream: AsyncIterator[TranscriptionStreamEvent], + ) -> AsyncGenerator[str, None]: + try: + async for event in stream: + yield f"data: {event.model_dump_json(exclude_unset=True)}\n\n" + yield TRANSCRIPTION_STREAM_DONE_FRAME + except Exception as stream_error: # noqa: BLE001 # stream boundary: a mid-stream failure ends the SSE body as an error event + replacement: Final = await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=stream_error, + request_data=( + {**data, "litellm_logging_obj": stream_logging_obj} + if isinstance(stream_logging_obj, LiteLLMLoggingObj) + else data + ), + ) + yield _transcription_sse_error_frame(stream_error if replacement is None else replacement) + yield TRANSCRIPTION_STREAM_DONE_FRAME + finally: + with anyio.CancelScope(shield=True): + close: Final = getattr(stream, "aclose", None) or getattr(stream, "close", None) + if callable(close): + close_result: Final = close() + if inspect.isawaitable(close_result): + await close_result + + return StreamingResponse( + guarded_transcription_events(response.response) + if isinstance(response, GuardedTranscriptionReplay) + else transcription_event_stream(cast(AsyncIterator[TranscriptionStreamEvent], response)), + media_type="text/event-stream", + headers=response_headers, + ) + + fastapi_response.headers.update(response_headers) return response except Exception as e: await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=data, + dispatched_call_type=CallTypes.atranscription.value, ) log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): @@ -13848,6 +14113,8 @@ async def _reject_realtime_session( @app.websocket("/openai/v1/realtime") @app.websocket("/v1/realtime") @app.websocket("/realtime") +@app.websocket("/v1/realtime/translations") +@app.websocket("/realtime/translations") async def realtime_websocket_endpoint( websocket: WebSocket, model: str | None = fastapi.Query(None, description="The model to use for the websocket connection."), @@ -13856,7 +14123,7 @@ async def realtime_websocket_endpoint( None, description="Comma-separated list of guardrail names to apply to this request.", ), - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_realtime_websocket), ): requested_protocols: Final = tuple( protocol.strip() @@ -13864,15 +14131,13 @@ async def realtime_websocket_endpoint( if protocol.strip() ) - route_model = model + is_translation: Final = websocket.url.path.endswith("/realtime/translations") + route_model: Final = _resolve_realtime_route_model(model, intent, is_translation) if route_model is None: - if intent == "transcription": - route_model = "gpt-realtime-whisper" - else: - await _reject_realtime_session( - websocket, user_api_key_dict, code=1008, reason="model query parameter is required" - ) - return + await _reject_realtime_session( + websocket, user_api_key_dict, code=1008, reason="model query parameter is required" + ) + return assert route_model is not None try: await can_key_call_resolved_model( @@ -13893,6 +14158,7 @@ async def realtime_websocket_endpoint( model=model, intent=intent, guardrails=guardrails, + is_translation=is_translation, ) @@ -13906,13 +14172,20 @@ async def _route_realtime_websocket_session( guardrails: str | None, request_path: str | None = None, live_session_start: Mapping[str, object] | None = None, + is_translation: bool = False, ) -> None: # Only use explicit parameters, not all query params - query_params: Final = cast(RealtimeQueryParams, dict(_realtime_query_params_template(model, intent))) + query_model: Final = route_model if is_translation else model + query_params: Final = cast( # cast-ok: cached tuples contain only the declared realtime query keys + RealtimeQueryParams, + dict(_realtime_query_params_template(query_model, intent)), + ) + data: dict[str, object] = { "model": route_model, "websocket": websocket, "query_params": query_params, # Only explicit params + "realtime_mode": "translation" if is_translation else "realtime", } # Pass guardrails into data so pre-call guardrail processing picks them up if guardrails: @@ -14805,6 +15078,7 @@ from litellm.repositories.table_repositories import ( ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository from litellm.responses.litellm_completion_transformation.count_request import ( ChatCountBridgeError, chat_count_request, @@ -14950,7 +15224,7 @@ async def _try_provider_token_count( deployment: Mapping[str, object] | None, request_model: str, tools: list | None = None, - system: str | None = None, + system: object | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: CountContextManagement | None = None, @@ -15020,6 +15294,7 @@ async def _provider_token_count( request_model: str, settings: TokenCountSettings, local_estimate_allowed: bool, + forwarded_edits: bool = False, ) -> Union["TokenCountResponse", ProviderCountFailure, None]: """Count with the provider, using the token_count_model in its region when one is set. @@ -15043,7 +15318,7 @@ async def _provider_token_count( pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model region_hint: Final = ( None - if pinned_count_model is None + if pinned_count_model is None or (forwarded_edits and counted.status_code == 400) else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. " "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" ) @@ -15051,7 +15326,7 @@ async def _provider_token_count( framing_can_stand_in: Final = (count_model is None or not count_model.configured) and ( failure.cannot_count_model or (pinned_count_model is not None and failure.denied_in_region) ) - if not (framing_can_stand_in and settings.framing is not None and local_estimate_allowed): + if not (framing_can_stand_in and settings.frames_counts and local_estimate_allowed): raise failure.to_proxy_exception() if pinned_count_model is not None: verbose_proxy_logger.warning( @@ -15251,6 +15526,7 @@ async def _tokenizer_or_provider_error( async def _local_token_estimate( model: str, + template_model: str, prompt: str | None, messages: Sequence[AllMessageValues] | None, tools: list[ChatCompletionToolParam] | None, @@ -15263,7 +15539,7 @@ async def _local_token_estimate( else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure) ) count: Final = offload_token_count(litellm.token_counter)( - model=model, + model=template_model, text=prompt, messages=messages, tools=tools, @@ -15305,6 +15581,9 @@ class _TokenCountDeployment(TypedDict, total=False): _TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment) +_COUNT_TOOLS: Final[TypeAdapter[tuple[Mapping[str, object], ...] | None]] = TypeAdapter( + tuple[Mapping[str, object], ...] | None +) _ROUTED_DEPLOYMENT: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -15326,10 +15605,15 @@ def _deployment_aws_region(deployment: _TokenCountDeployment) -> str: ) +def _deployment_base_model(deployment: _TokenCountDeployment) -> str | None: + return _non_empty_str(deployment.get("model_info", {}).get("base_model")) or _non_empty_str( + deployment.get("litellm_params", {}).get("base_model") + ) + + def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]: - model_info: Final = deployment.get("model_info", {}) params: Final = deployment.get("litellm_params", {}) - base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model")) + base_model: Final = _deployment_base_model(deployment) if base_model is not None: return base_model, None model: Final = _non_empty_str(params.get("model")) @@ -15374,12 +15658,15 @@ def _count_context_management( model, general_settings_view(), user_model if isinstance(user_model, str) else None ) router_defaults: Final = _router_default_litellm_params(llm_router) - sent_params: Final = tuple( - {**deployment["litellm_params"], **router_defaults, **request_litellm_params} + deployment_params: Final = tuple( + deployment["litellm_params"] for deployment in llm_router.registry().deployments( inference_model if isinstance(inference_model, str) else model, source="routable" ) ) + sent_params: Final = tuple( + {**params, **router_defaults, **request_litellm_params} for params in deployment_params or ({},) + ) if any(context_management_explicitly_dropped(params.get("additional_drop_params")) for params in sent_params): return None return CountContextManagement( @@ -15389,17 +15676,104 @@ def _count_context_management( ) +_GATEWAY_EDIT_ROUTES: Final = frozenset({Target.CHAT_COMPLETIONS, Target.RESPONSES}) + + +def _messages_route(deployment: _TokenCountDeployment, request_model: str) -> Target: + params: Final = deployment.get("litellm_params", {}) + return messages_target( + _non_empty_str(params.get("model")) or request_model, + _non_empty_str(params.get("custom_llm_provider")), + deployment.get("model_info"), + _deployment_api_base(deployment), + ) + + +def _rejected_view(status_code: int, message: str) -> ProxyException: + return ProviderCountFailure( + status_code=status_code, message=message, rejected_before_provider_call=True + ).to_proxy_exception() + + +async def _bridged_request_view( + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + messages_format: TokenCountMessagesFormat, + edits: CountContextManagement | None, + deployment: _TokenCountDeployment, + fallback_model: str, +) -> InferenceRequestView | InferenceRequestViewError | None: + view: Final = await inference_request_view( + messages, + system, + messages_format, + count_context_edits(edits, tools, deployment.get("litellm_params", {}), fallback_model), + ) + return None if isinstance(view, InferenceRequestView) and not view.messages else view + + +async def _responses_request_view( + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + edits: CountContextManagement, + deployment: _TokenCountDeployment, + fallback_model: str, +) -> InferenceRequestView | InferenceRequestViewError: + from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError + from litellm.llms.anthropic.messages.responses_adapters.handler import count_view_messages + + try: + edited: Final = await count_view_messages( + model=_non_empty_str(deployment.get("litellm_params", {}).get("model")) or fallback_model, + messages=messages, + tools=tools, + system=system, + context_management=edits, + ) + except AnthropicContextManagementError as error: + if error.status_code >= 500: + return InferenceRequestView(messages=messages, system=system) + return InferenceRequestViewError(error_message=error.message, status_code=error.status_code) + return InferenceRequestView(messages=edited, system=system) + + +async def _gateway_request_view( + route: Target, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + messages_format: TokenCountMessagesFormat, + edits: CountContextManagement | None, + deployment: _TokenCountDeployment, + fallback_model: str, +) -> InferenceRequestView | InferenceRequestViewError | None: + match route: + case Target.CHAT_COMPLETIONS: + return await _bridged_request_view( + messages, system, tools, messages_format, edits, deployment, fallback_model + ) + case Target.RESPONSES if edits is not None: + return await _responses_request_view(messages, system, tools, edits, deployment, fallback_model) + case _: + return None + + def _responses_routing_kwargs(responses_request: Mapping[str, object] | None) -> Mapping[str, object]: previous_response_id: Final = None if responses_request is None else responses_request.get("previous_response_id") return {"previous_response_id": previous_response_id} if isinstance(previous_response_id, str) else {} +def _inference_params( + litellm_params: Mapping[str, object], responses_request: Mapping[str, object] +) -> Mapping[str, object]: + return {**litellm_params, **_router_default_litellm_params(llm_router), **responses_request} + + def _count_drop_policy(deployment: Mapping[str, object] | None, responses_request: Mapping[str, object]) -> DropPolicy: litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {} - router_defaults: Final[Mapping[str, object]] = ( - {} if llm_router is None else _router_default_litellm_params(llm_router) - ) - return drop_policy({**litellm_params, **router_defaults, **responses_request}) + return drop_policy(_inference_params(litellm_params, responses_request)) def _native_responses_count( @@ -15422,8 +15796,54 @@ def _native_responses_count( ) -def _deployment_api_base(deployment: _TokenCountDeployment) -> str | None: - return _non_empty_str(deployment.get("litellm_params", {}).get("api_base")) +def _deployment_api_base( + deployment: _TokenCountDeployment, responses_request: Mapping[str, object] | None = None +) -> str | None: + return _non_empty_str( + _inference_params(deployment.get("litellm_params", {}), responses_request or {}).get("api_base") + ) + + +def _deployment_provider( + deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] | None +) -> str | None: + params: Final = deployment.get("litellm_params", {}) + try: + return litellm.get_llm_provider( + model=_non_empty_str(params.get("model")) or request_model, + custom_llm_provider=_non_empty_str(params.get("custom_llm_provider")), + api_base=_deployment_api_base(deployment, responses_request), + )[1] + except litellm.BadRequestError: + return None + + +def _serves_native_responses( + deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] +) -> bool: + params: Final = deployment.get("litellm_params", {}) + inference: Final = _inference_params(params, responses_request) + try: + return provider_stores_the_conversation( + _non_empty_str(params.get("model")) or request_model, + _non_empty_str(params.get("custom_llm_provider")), + { + "use_chat_completions_api": bool(inference.get("use_chat_completions_api")), + "model_info": dict(deployment.get("model_info", {})), + "api_base": _non_empty_str(inference.get("api_base")), + }, + ) + except litellm.BadRequestError: + return False + + +_AZURE_PROVIDERS: Final = frozenset((litellm.LlmProviders.AZURE.value, litellm.LlmProviders.AZURE_AI.value)) + + +def _with_strict_by_default(tools: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + return tuple( + {**tool, "strict": True} if tool.get("type") == "function" and "strict" not in tool else tool for tool in tools + ) def _system_message(system: object) -> ChatCompletionSystemMessage | None: @@ -15582,6 +16002,40 @@ async def _count_request_tokens( ) prompt_only: Final = bool(prompt) and conversation_empty request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment) + local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment) + malformed_messages: Final = ( + None + if local_count_error is None + else ProxyException(message=local_count_error, type="token_counting_error", param="messages", code=400) + ) + context_edits: Final = _count_context_management(context_management, request.model, request_litellm_params) + route: Final = ( + _messages_route(deployment_view, request.model) + if request_messages is not None and messages_format == "anthropic" + else None + ) + if context_edits is not None and route in _GATEWAY_EDIT_ROUTES and malformed_messages is not None: + raise malformed_messages + edited_request: Final = ( + None + if route is None or request_messages is None + else await _gateway_request_view( + route, + request_messages, + system, + _COUNT_TOOLS.validate_python(tools), + messages_format, + context_edits, + deployment_view, + model_to_use, + ) + ) + if isinstance(edited_request, InferenceRequestViewError): + raise _rejected_view(edited_request.status_code, edited_request.error_message) + inference_messages: Final = ( + request_messages if edited_request is None else [dict(message) for message in edited_request.messages] + ) + inference_system: Final = system if edited_request is None else edited_request.system contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents)) prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment) @@ -15594,7 +16048,7 @@ async def _count_request_tokens( else await chat_count_request( model_to_use, custom_llm_provider, - _deployment_api_base(deployment_view), + _deployment_api_base(deployment_view, routed_request), routed_request, _count_drop_policy(routed_deployment, routed_request), ) @@ -15621,30 +16075,30 @@ async def _count_request_tokens( if chat_count is not None else [{"role": "user", "content": prompt}] if prompt_as_message - else request_messages + else inference_messages ), contents=None if prompt_as_message else contents, deployment=count_deployment, request_model=request.model, tools=tools if chat_count is None else [dict(tool) for tool in chat_count.tools or ()] or None, - system=system, + system=inference_system, messages_format=messages_format, output_format=output_format, - context_management=_count_context_management(context_management, request.model, request_litellm_params), + context_management=None if edited_request is not None else context_edits, chat_params=MappingProxyType({}) if chat_count is None else chat_count.params, ), model_to_use=model_to_use, request_model=request.model, settings=settings, local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, + forwarded_edits=context_edits is not None and edited_request is None, ) ) if isinstance(ladder, TokenCountResponse): return ladder - local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment) - if local_count_error is not None: - raise ProxyException(message=local_count_error, type="token_counting_error", param="messages", code=400) + if malformed_messages is not None: + raise malformed_messages # Check if token counter is disabled before fallback if litellm.disable_token_counter is True: @@ -15655,6 +16109,27 @@ async def _count_request_tokens( code=503, ) + local_view: Final = ( + None + if route is not Target.MESSAGES or context_edits is None or request_messages is None + else await _bridged_request_view( + request_messages, + system, + _COUNT_TOOLS.validate_python(tools), + messages_format, + replace(context_edits, drop_invalid_edits=False), + deployment_view, + model_to_use, + ) + ) + if isinstance(local_view, InferenceRequestViewError): + raise _rejected_view(local_view.status_code, local_view.error_message) + counted_view: Final = edited_request if local_view is None else local_view + local_messages: Final = ( + request_messages if counted_view is None else [dict(message) for message in counted_view.messages] + ) + local_system: Final = system if counted_view is None else counted_view.system + # Default LiteLLM token counting custom_tokenizer: CustomHuggingfaceTokenizer | None = None if model_info is not None: @@ -15662,22 +16137,51 @@ async def _count_request_tokens( CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - system_message: Final = _system_message(system) + system_message: Final = _system_message(local_system) prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes Sequence[AllMessageValues] | None, - [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else request_messages), + [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else local_messages), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) ) - counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats + native_responses: Final = responses_request is not None and _serves_native_responses( + deployment_view, request.model, responses_request + ) + sent_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 + openai_model: Final = deployment_openai_model(deployment_view) + template_model: Final = openai_model or _deployment_base_model(deployment_view) or model_to_use + tool_section: Final = chat_template_tool_section(template_model) + counted_tools: Final = ( + cast(list[ChatCompletionToolParam], list(_with_strict_by_default(sent_tools))) # cast-ok: strict added + if sent_tools is not None + and native_responses + and _deployment_provider(deployment_view, request.model, responses_request) in _AZURE_PROVIDERS + and tool_section is not None + and tool_section.renders_strict_only_when_true + else sent_tools + ) + framing: Final = ( + 0 + if counted_messages is None + else settings.hidden_prompt_tokens( + carries_tools=sends_function_tools(counted_tools), responses_api=native_responses + ) + ) kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None exact: Final = ( - await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools) + await _kimi_k3_count( + request + if counted_view is None + else request.model_copy(update={"messages": local_messages, "system": local_system}), + messages_format, + model_to_use, + kimi_k3_messages, + counted_tools, + ) if kimi_k3_messages is not None and custom_tokenizer is None else None ) @@ -15690,7 +16194,8 @@ async def _count_request_tokens( count_source="local_estimate", ) estimate: Final = await _local_token_estimate( - model=deployment_openai_model(deployment_view) or model_to_use, + model=openai_model or model_to_use, + template_model=template_model, prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, @@ -18006,8 +18511,6 @@ async def model_info_v1( alias_models: Final = copy.deepcopy(registry.deployments(source="alias")) all_models.extend(alias_models) - all_models = await run_model_listing(partial(expand_wildcard_deployments_for_model_info, all_models, llm_router)) - allowed_model_names: Final = await _get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=user_api_key_dict, llm_router=llm_router, @@ -18016,6 +18519,15 @@ async def model_info_v1( proxy_logging_obj=proxy_logging_obj, ) + all_models = await run_model_listing( + partial( + expand_wildcard_deployments_for_model_info, + all_models, + llm_router, + allowed_model_names=allowed_model_names, + ) + ) + all_models = _filter_v1_model_info_deployments( all_models=all_models, allowed_model_names=allowed_model_names, @@ -20792,53 +21304,56 @@ async def delete_config_general_settings( if data.field_name in _GENERAL_SETTINGS_UI_LITELLM_FIELDS: return await _reset_general_settings_ui_litellm_field(data.field_name, user_api_key_dict) - if data.field_name not in ConfigGeneralSettings.model_fields: + config_table: Final[TableActions[_ConfigParamValueRow]] = _config_param_table(prisma_client) + stored_row: Final = await config_table.find_first(where={"param_name": "general_settings"}) + stored: Final = ( + _SETTINGS_MAPPING.validate_python(stored_row.param_value) + if stored_row is not None and isinstance(stored_row.param_value, Mapping) + else None + ) + + if not _general_settings_row_is_readable(stored_row): + raise HTTPException( + status_code=400, + detail={ + "error": "general_settings is stored in the database as something other than a JSON object, so no " + "field can be removed from it" + }, + ) + + if data.field_name not in ConfigGeneralSettings.model_fields and (stored is None or data.field_name not in stored): raise HTTPException( status_code=400, detail={"error": f"Invalid field={data.field_name} passed in."}, ) - proxy_config.reject_config_owned_deletes(section_name="general_settings", keys=(data.field_name,)) + if data.field_name in _db_managed_general_settings_keys(): + proxy_config.reject_config_owned_deletes(section_name="general_settings", keys=(data.field_name,)) - ## get general settings from db - db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( - where={"param_name": "general_settings"} - ) - ### pop the value - - if db_general_settings is None or db_general_settings.param_value is None: + if stored is None: raise HTTPException( status_code=400, detail={"error": f"Field name={data.field_name} not in config"}, ) - else: - general_settings = dict(db_general_settings.param_value) - - before_general_settings: Final = copy.deepcopy(general_settings) - - ## update db - - general_settings.pop(data.field_name, None) + remaining: Final = {key: value for key, value in stored.items() if key != data.field_name} response: Final = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, data={ "create": { "param_name": "general_settings", - "param_value": json.dumps(general_settings), + "param_value": json.dumps(remaining), }, - "update": {"param_value": json.dumps(general_settings)}, + "update": {"param_value": json.dumps(remaining)}, }, ) await invalidate_config_param("general_settings") - proxy_config.settings.apply_db_row("general_settings", general_settings) + proxy_config.settings.apply_db_row("general_settings", remaining) if is_resource_list("general_settings", data.field_name): - stored_endpoints: Final = general_settings.get("pass_through_endpoints") + stored_endpoints: Final = remaining.get("pass_through_endpoints") await proxy_config.serve_pass_through_endpoints(stored_endpoints if isinstance(stored_endpoints, list) else ()) asyncio.create_task( - create_config_audit_log( - "general_settings", "deleted", before_general_settings, general_settings, user_api_key_dict - ) + create_config_audit_log("general_settings", "deleted", dict(stored), remaining, user_api_key_dict) ) return response diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 044b8a1de6e..aa7554c4a5a 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2641,6 +2641,34 @@ ], "default_model_placeholder": "nace/drex-latest" }, + { + "provider": "PARALLEL_AI", + "provider_display_name": "Parallel AI", + "litellm_provider": "parallel_ai", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.parallel.ai", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "parallel_ai/decision" + }, { "provider": "NEBIUS", "provider_display_name": "Nebius", diff --git a/litellm/proxy/realtime_endpoints/endpoints.py b/litellm/proxy/realtime_endpoints/endpoints.py index 9dbf6bcc43d..17c52fc54bd 100644 --- a/litellm/proxy/realtime_endpoints/endpoints.py +++ b/litellm/proxy/realtime_endpoints/endpoints.py @@ -2,19 +2,24 @@ import json import time -from collections.abc import Mapping +from collections.abc import Generator, Mapping +from contextlib import contextmanager from typing import TYPE_CHECKING, Final import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi import status as http_status -from pydantic import ConfigDict, TypeAdapter, with_config +from fastapi.responses import PlainTextResponse +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict +import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.custom_openapi_spec import inline_request_body from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -43,7 +48,26 @@ router: Final = APIRouter() _REALTIME_TOKEN_VERSION: Final = "realtime_v1" _DEFAULT_REALTIME_MODEL: Final = "gpt-4o-realtime-preview" _DEFAULT_TRANSCRIPTION_MODEL: Final = "gpt-realtime-whisper" -_ALLOWED_SESSION_TYPES: Final = ("realtime", "transcription") +_ALLOWED_SESSION_TYPES: Final = ("realtime", "transcription", "translation") + + +@contextmanager +def _invalid_request_body_answers_400() -> Generator[None, None, None]: + try: + yield + except ValidationError as error: + raise ProxyException(message=str(error), type="invalid_request_error", param=None, code=400) from error + + +_TRANSLATION_SECRET_BODY: Final = { + "requestBody": inline_request_body( + RealtimeClientSecretRequest, + {"model": "gpt-realtime-translate", "session": {"type": "translation", "model": "gpt-realtime-translate"}}, + ) +} +_TRANSLATION_SDP_BODY: Final = { + "requestBody": {"required": True, "content": {"application/sdp": {"schema": {"type": "string"}}}} +} @with_config(ConfigDict(strict=True, extra="allow")) @@ -140,19 +164,60 @@ def _set_transcription_model_on_session( } +def nested_transcription_upstream_model(model: str, team_id: str | None, llm_router: "Router | None") -> str: + deployment_model: Final = ( + llm_router.registry().resolve(model, kind="upstream", team_id=team_id) if llm_router is not None else None + ) + if deployment_model is None: + return model + try: + return litellm.get_llm_provider(model=deployment_model)[0] + except litellm.BadRequestError: + return deployment_model + + +async def _authorize_and_bind_nested_transcription_models( + session_data: dict, # mutable-ok: session payload is rewritten in place for provider serialization + user_api_key_dict: UserAPIKeyAuth, + llm_model_list: list | None, # mutable-ok: inherited auth helper accepts the proxy model list + llm_router: "Router | None", +) -> None: + nested_models: Final = tuple(_transcription_model_candidates_from_session(session_data)) + for nested_model in nested_models: + await can_key_call_resolved_model( + model=nested_model, + valid_token=user_api_key_dict, + llm_model_list=llm_model_list, + llm_router=llm_router, + ) + if nested_models: + _set_transcription_model_on_session( + session=session_data, + model=nested_transcription_upstream_model(nested_models[0], user_api_key_dict.team_id, llm_router), + ) + + async def _prepare_client_secret_session( req: RealtimeClientSecretRequest, user_api_key_dict: UserAPIKeyAuth, llm_model_list: list | None, llm_router: "Router | None", + forced_session_type: str | None = None, ) -> tuple[str, dict | None, str]: - session_type: Final = _coerce_realtime_session_type(req.session.type if req.session else None) - session_data: Final[dict | None] = req.session.model_dump(exclude_none=True) if req.session else None + requested_session_type: Final = req.session.type if req.session else None + if forced_session_type is None and requested_session_type == "translation": + raise HTTPException(status_code=400, detail="Translation sessions require the translations endpoint") + session_type: Final = forced_session_type or _coerce_realtime_session_type(requested_session_type) + session_data: Final[dict | None] = ( + req.session.model_dump(exclude_unset=True) if req.session else ({} if session_type == "translation" else None) + ) if session_data is not None: session_data["type"] = session_type session_model: Final = req.session.model if req.session else None - model: str = session_model or req.model or _DEFAULT_REALTIME_MODEL + requested_model: Final = session_model or req.model + default_model: Final = "gpt-realtime-translate" if session_type == "translation" else _DEFAULT_REALTIME_MODEL + model: str = requested_model or default_model if session_type != "transcription": await can_key_call_resolved_model( model=model, @@ -160,6 +225,15 @@ async def _prepare_client_secret_session( llm_model_list=llm_model_list, llm_router=llm_router, ) + if session_data is not None and (requested_model or session_type == "translation"): + session_data["model"] = model + if session_type == "translation": + await _authorize_and_bind_nested_transcription_models( + session_data=session_data, + user_api_key_dict=user_api_key_dict, + llm_model_list=llm_model_list, + llm_router=llm_router, + ) return model, session_data, session_type transcription_model_candidates: Final = _transcription_model_candidates_from_session(session_data or {}) @@ -194,12 +268,13 @@ def _encode_realtime_token_payload( team_id: str | None, expires_at: int | None, session_type: str = "realtime", + caller_auth: Mapping[str, object] | None = None, ) -> str: """ Encode metadata with the upstream ephemeral key so /realtime/calls can route without requiring model as a query param. """ - payload: Final[dict[str, str | int | None]] = { + payload: Final[Mapping[str, object]] = { "v": _REALTIME_TOKEN_VERSION, "ephemeral_key": ephemeral_key, "model_id": model_id, @@ -207,10 +282,24 @@ def _encode_realtime_token_payload( "team_id": team_id or "", "expires_at": expires_at, "session_type": session_type, + **({} if caller_auth is None else {"caller_auth": caller_auth}), } return json.dumps(payload, separators=(",", ":")) +def _caller_auth_for_token(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: + return user_api_key_dict.model_dump(mode="json", exclude_none=True, exclude={"api_key", "parent_otel_span"}) + + +def _caller_auth_from_token( + decoded_payload: Mapping[str, object] | None, user_id: str | None, team_id: str | None +) -> UserAPIKeyAuth: + caller_auth: Final = None if decoded_payload is None else as_str_mapping(decoded_payload.get("caller_auth")) + if caller_auth is None: + return UserAPIKeyAuth(user_id=user_id, team_id=team_id) + return UserAPIKeyAuth.model_validate({**caller_auth, "api_key": caller_auth.get("token")}) + + def _header_value(headers: Mapping[str, str], name: str, default: str) -> str: return headers.get(name, default) @@ -249,11 +338,24 @@ def _decode_realtime_token_payload( dependencies=[Depends(user_api_key_auth)], tags=["realtime"], ) +@router.post( + "/v1/realtime/translations/client_secrets", + openapi_extra=_TRANSLATION_SECRET_BODY, + dependencies=[Depends(user_api_key_auth)], + tags=["realtime"], +) +@router.post( + "/realtime/translations/client_secrets", + openapi_extra=_TRANSLATION_SECRET_BODY, + dependencies=[Depends(user_api_key_auth)], + tags=["realtime"], +) async def create_realtime_client_secret( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> RealtimeClientSecretResponse: + """Create a short-lived client token for a Realtime conversation, transcription, or translation session""" from litellm.proxy.proxy_server import ( add_litellm_data_to_request, general_settings, @@ -269,13 +371,16 @@ async def create_realtime_client_secret( data: dict = {} try: body: Final = await read_request_body(request=request) - req: Final = RealtimeClientSecretRequest(**body) + with _invalid_request_body_answers_400(): + req: Final = RealtimeClientSecretRequest.model_validate(body) + is_translation_request: Final = "/realtime/translations/client_secrets" in request.url.path model, session_data, session_type = await _prepare_client_secret_session( req=req, user_api_key_dict=user_api_key_dict, llm_model_list=llm_model_list, llm_router=llm_router, + forced_session_type="translation" if is_translation_request else None, ) data = {"model": model} @@ -299,17 +404,20 @@ async def create_realtime_client_secret( proxy_config=proxy_config, ) + call_type: Final = ( + "acreate_realtime_translation_client_secret" if is_translation_request else "acreate_realtime_client_secret" + ) data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=data, - call_type="acreate_realtime_client_secret", + call_type=call_type, ) verbose_proxy_logger.debug("WebRTC: /v1/realtime/client_secrets (model=%s)", model) llm_call: Final = await route_request( data=data, - route_type="acreate_realtime_client_secret", + route_type=call_type, llm_router=llm_router, user_model=user_model, ) @@ -366,6 +474,7 @@ async def create_realtime_client_secret( team_id=getattr(user_api_key_dict, "team_id", None), expires_at=expires_at if isinstance(expires_at, int) else None, session_type=session_type, + caller_auth=_caller_auth_for_token(user_api_key_dict) if session_type == "translation" else None, ) encrypted_token: Final[str] = encrypt_value_helper(token_payload) upstream_json["value"] = encrypted_token @@ -392,10 +501,27 @@ async def create_realtime_client_secret( "/openai/v1/realtime/calls", tags=["realtime"], ) +@router.post( + "/v1/realtime/translations/calls", + status_code=http_status.HTTP_201_CREATED, + response_class=PlainTextResponse, + response_description="WebRTC SDP answer", + openapi_extra=_TRANSLATION_SDP_BODY, + tags=["realtime"], +) +@router.post( + "/realtime/translations/calls", + status_code=http_status.HTTP_201_CREATED, + response_class=PlainTextResponse, + response_description="WebRTC SDP answer", + openapi_extra=_TRANSLATION_SDP_BODY, + tags=["realtime"], +) async def proxy_realtime_calls( request: Request, fastapi_response: Response, ) -> Response: + """Exchange a WebRTC SDP offer for an answer using the ephemeral token from client_secrets""" from litellm.proxy.proxy_server import ( add_litellm_data_to_request, general_settings, @@ -417,6 +543,7 @@ async def proxy_realtime_calls( media_type="application/json", ) + is_translation_request: Final = "/realtime/translations/calls" in request.url.path encrypted_token: Final = auth_header.removeprefix("Bearer ").strip() decrypted_token_value: Final = decrypt_value_helper( value=encrypted_token, @@ -434,33 +561,39 @@ async def proxy_realtime_calls( if decoded_payload is not None: # Check token expiry expires_at: Final = decoded_payload.get("expires_at") - if expires_at is not None and isinstance(expires_at, int): - if time.time() > expires_at: - return Response( - content=json.dumps({"error": "Token has expired"}), - status_code=http_status.HTTP_401_UNAUTHORIZED, - media_type="application/json", - ) + if isinstance(expires_at, int) and time.time() > expires_at: + return Response( + content=json.dumps({"error": "Token has expired"}), + status_code=http_status.HTTP_401_UNAUTHORIZED, + media_type="application/json", + ) openai_ephemeral_key = decoded_payload.get("ephemeral_key", "") model = decoded_payload.get("model_id") or request.query_params.get("model") or _DEFAULT_REALTIME_MODEL user_id = decoded_payload.get("user_id") or None team_id = decoded_payload.get("team_id") or None - session_type = _coerce_realtime_session_type(decoded_payload.get("session_type")) + raw_session_type: Final = decoded_payload.get("session_type") + session_type = _coerce_realtime_session_type(raw_session_type) + if is_translation_request != (raw_session_type == "translation"): + return Response( + content=json.dumps({"error": "Token is not valid for this Realtime endpoint"}), + status_code=http_status.HTTP_401_UNAUTHORIZED, + media_type="application/json", + ) else: - # Backward compatibility: older tokens contained only encrypted upstream key. + if is_translation_request: + return Response( + content=json.dumps({"error": "Token is not valid for this Realtime endpoint"}), + status_code=http_status.HTTP_401_UNAUTHORIZED, + media_type="application/json", + ) openai_ephemeral_key = decrypted_token_value model = request.query_params.get("model", _DEFAULT_REALTIME_MODEL) user_id = None team_id = None session_type = "realtime" - # Build a minimal UserAPIKeyAuth with user/team IDs from the token - # so spend tracking and budget enforcement work correctly. - minimal_auth: Final = UserAPIKeyAuth( - user_id=user_id, - team_id=team_id, - ) + token_caller_auth: Final = _caller_auth_from_token(decoded_payload, user_id, team_id) data: dict = {} try: @@ -487,22 +620,23 @@ async def proxy_realtime_calls( data=data, request=request, general_settings=general_settings, - user_api_key_dict=minimal_auth, + user_api_key_dict=token_caller_auth, version=version, proxy_config=proxy_config, ) + call_type: Final = "arealtime_translation_calls" if is_translation_request else "arealtime_calls" data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=minimal_auth, + user_api_key_dict=token_caller_auth, data=data, - call_type="arealtime_calls", + call_type=call_type, ) verbose_proxy_logger.debug("WebRTC: /v1/realtime/calls (model=%s)", model) llm_call: Final = await route_request( data=data, - route_type="arealtime_calls", + route_type=call_type, llm_router=llm_router, user_model=user_model, ) @@ -510,7 +644,7 @@ async def proxy_realtime_calls( except Exception as e: await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=minimal_auth, + user_api_key_dict=token_caller_auth, original_exception=e, request_data=data, ) @@ -581,7 +715,8 @@ async def create_realtime_transcription_session( data: dict = {} try: body: Final = await read_request_body(request=request) - req: Final = RealtimeTranscriptionSessionRequest(**body) + with _invalid_request_body_answers_400(): + req: Final = RealtimeTranscriptionSessionRequest.model_validate(body) model: Final[str] = req.resolved_model() or "gpt-realtime-whisper" await can_key_call_resolved_model( @@ -592,6 +727,7 @@ async def create_realtime_transcription_session( ) transcription_session: Final = {k: v for k, v in body.items() if k != "model"} + _set_transcription_model_on_session(session=transcription_session, model=model, create_if_missing=True) data = {"model": model, "transcription_session": transcription_session} data = await add_litellm_data_to_request( diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 4f2daed15ed..4905751a85d 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -20,6 +20,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_param, openai_error_type, ) +from litellm.types.utils import CallTypes router: Final = APIRouter() @@ -114,7 +115,10 @@ async def rerank( return response except Exception as e: await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=data, + dispatched_call_type=CallTypes.arerank.value, ) log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index ecc31f3b0a7..148d465099f 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -621,7 +621,9 @@ RouteType = Literal[ "alist_input_items", "_arealtime", # private function for realtime API "acreate_realtime_client_secret", + "acreate_realtime_translation_client_secret", "arealtime_calls", + "arealtime_translation_calls", "acreate_realtime_transcription_session", "_aresponses_websocket", # private function for responses WebSocket mode "aimage_edit", diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index fc6fb266d04..dcfd6148f54 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -432,7 +432,7 @@ async def invalidate_budget_reservation_counters( async def release_or_invalidate_budget_reservation( - budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict + budget_reservation: dict[str, object] | None, # mutable-ok: stamps finalized on the caller's reservation ) -> None: """Reconcile a still-open reservation on a terminal path that settles no cost. diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index ba5f873c4e9..a9478a1f889 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -20,7 +20,6 @@ from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, EMPTY_MAPPING, - INTERNAL_CALL_ORIGIN_METADATA_KEY, KUBERNETES_POD_ROUTING_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, LITELLM_TRUNCATED_PAYLOAD_FIELD, @@ -30,7 +29,6 @@ from litellm.constants import ( MAX_SPEND_LOG_MODEL_NAME_LENGTH, MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB, MCP_SPEND_LOG_MODEL_PREFIX, - NON_INFERENCE_CALL_TYPES, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL, @@ -53,7 +51,11 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider from litellm.litellm_core_utils.hidden_params import served_from_cache -from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call +from litellm.litellm_core_utils.internal_call_metadata import ( + BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX, + is_unbilled_non_inference_call, + settles_background_response, +) from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, UsageObservation, @@ -79,13 +81,8 @@ from litellm.responses.litellm_completion_transformation.compaction_items import summary_only_compaction_content, with_compaction_summary, ) -from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - SETTLES_BACKGROUND_RESPONSE_KWARG, - BackgroundResponseSettlement, -) from litellm.types.router import DeploymentTypedDict, LiteLLM_Params, SilentExperimentFlag from litellm.types.utils import ( - BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, PROMPT_CARRYING_GUARDRAIL_FIELDS, AzureSpillover, CallTypes, @@ -347,24 +344,21 @@ def _get_spend_logs_metadata( BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" -BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" -def _cost_row_request_id_suffix( - call_type: str, metadata: Mapping[str, object], settles_background_response: bool -) -> str: +def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str: if call_type == CallTypes.aretrieve_batch.value: return BATCH_COST_REQUEST_ID_SUFFIX - if settles_background_response or ( - call_type in NON_INFERENCE_CALL_TYPES - and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - ): + if settles_background_response(call_type, metadata, kwargs): return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX return "" def get_spend_logs_id( - call_type: str, response_obj: dict, kwargs: Mapping[str, object], metadata: Mapping[str, object] = EMPTY_MAPPING + call_type: str, + response_obj: Mapping[str, object], + kwargs: Mapping[str, object], + metadata: Mapping[str, object] = EMPTY_MAPPING, ) -> str | None: standard_logging_payload = kwargs.get("standard_logging_object") candidate_ids: Final = ( @@ -377,10 +371,7 @@ def get_spend_logs_id( ) if resolved_id is None: return None - settles_background_response: Final = ( - kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES - ) - return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, settles_background_response)}" + return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" _MISSING_ATTRIBUTE: Final = object() @@ -607,6 +598,7 @@ def get_logging_payload( kwargs = {} rejected_as_unknown_model: Final = isinstance(response_obj, ProxyModelNotFoundError) + failure_exception: Final = response_obj if isinstance(response_obj, BaseException) else None response_served_from_cache: Final = served_from_cache(response_obj) if response_obj is None: response_obj = {} @@ -897,14 +889,11 @@ def get_logging_payload( additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens clean_metadata["additional_usage_values"] = additional_usage_values drop_requester_ip_address: Final = requester_ip_address_logging_disabled() - error_information: Final = clean_metadata["error_information"] if drop_requester_ip_address: clean_metadata["requester_ip_address"] = None - if drop_requester_ip_address and isinstance(error_information, Mapping): - clean_metadata["error_information"] = cast( # cast-ok: same keys, only the str values are masked - StandardLoggingPayloadErrorInformation, - _error_information_without_ip_addresses(_OBJECT_MAPPING.validate_python(error_information)), - ) + clean_metadata["error_information"] = sanitize_error_information_for_spend_logs( + clean_metadata["error_information"], original_exception=failure_exception + ) if litellm.cache is None: cache_key = "Cache OFF" diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index e5348684d78..6d31fa3f226 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -16,6 +16,7 @@ from typing import ( from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile +from prisma.errors import PrismaError from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model, field_validator from pydantic.fields import FieldInfo, PydanticUndefined from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -26,6 +27,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.validation_error_body import public_validation_errors from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError @@ -1247,6 +1249,37 @@ async def get_sso_settings(): } +_ENVIRONMENT_VARIABLES_PARAM: Final = TypeAdapter(dict[str, object]) + + +def _parse_environment_variables_param(param_value: str | Mapping[str, object] | None) -> Mapping[str, object]: + if isinstance(param_value, str): + return _ENVIRONMENT_VARIABLES_PARAM.validate_python(json.loads(param_value)) + return param_value or {} + + +async def _load_environment_variables_param( + repo: ConfigRepository, +) -> Mapping[str, object] | None | PrismaError | ValueError: + try: + entry: Final = await _config_param_db(repo).find_unique(where={"param_name": "environment_variables"}) + param_value: Final = None if entry is None else _parse_environment_variables_param(entry.param_value) + except (PrismaError, ValueError) as e: + return e + else: + return param_value + + +def _decrypted_sso_fields(environment_variables: Mapping[str, object]) -> Mapping[str, str | None]: + return MappingProxyType( + { + field_name: decrypt_value_helper(value=value, key=env_var, return_original_value=True) + for field_name, env_var in SSO_FIELD_ENV_VARS.items() + if isinstance(value := environment_variables.get(env_var), str) + } + ) + + @router.patch( "/update/sso_settings", tags=["SSO Settings"], @@ -1310,9 +1343,24 @@ async def update_sso_settings( if "general_settings" not in config: config["general_settings"] = {} - # Update environment variables in config and in memory - sso_data: Final = sso_config.model_dump() - for field_name, value in sso_data.items(): + # PATCH semantics: omitted fields keep their stored value, an explicit null clears it. SSO values that + # still live only in the environment_variables row are carried over, since that row is stripped below + environment_variables_param: Final = await _load_environment_variables_param( + ConfigRepository(prisma_client, use_writer=True) + ) + legacy_env_values: Final[Mapping[str, str | None]] = ( + _decrypted_sso_fields(environment_variables_param) + if isinstance(environment_variables_param, Mapping) + else MappingProxyType({}) + ) + stored_values: Final = {k: v for k, v in (before_sso_data or {}).items() if v is not None} + fallback_values: Final = {**legacy_env_values, **stored_values} + submitted: Final = sso_config.model_dump(include=sso_config.model_fields_set) + sso_data: Final = { + field_name: submitted[field_name] if field_name in submitted else fallback_values.get(field_name) + for field_name in SSOConfig.model_fields + } + for field_name, value in submitted.items(): if field_name in SSO_FIELD_ENV_VARS: env_var_name = SSO_FIELD_ENV_VARS[field_name] if value: @@ -1348,25 +1396,17 @@ async def update_sso_settings( ) ) - # Remove SSO-related env vars from config.environment_variables - try: - env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client, use_writer=True)).find_unique( - where={"param_name": "environment_variables"} + if isinstance(environment_variables_param, (PrismaError, ValueError)): + raise HTTPException( + status_code=500, + detail={"error": f"Error updating environment_variables: {environment_variables_param}"}, ) - # If no environment_variables entry exists, nothing to clean up - if env_var_entry is not None: - if env_var_entry.param_value is not None: - if isinstance(env_var_entry.param_value, str): - environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value) - else: - environment_variables = dict(env_var_entry.param_value) - else: - environment_variables = {} - + if environment_variables_param is not None: + try: env_vars_to_remove: Final = set(SSO_FIELD_ENV_VARS.values()) filtered_env_vars: Final = { - key: value for key, value in environment_variables.items() if key not in env_vars_to_remove + key: value for key, value in environment_variables_param.items() if key not in env_vars_to_remove } await _config_param_db(ConfigRepository(prisma_client)).update( @@ -1376,16 +1416,16 @@ async def update_sso_settings( }, ) await invalidate_config_param("environment_variables") - except Exception as e: - raise HTTPException( - status_code=500, - detail={"error": f"Error updating environment_variables: {e}"}, - ) + except Exception as e: + raise HTTPException( + status_code=500, + detail={"error": f"Error updating environment_variables: {e}"}, + ) return { "message": "SSO settings updated successfully", "status": "success", - "settings": sso_data, + "settings": mask_sensitive_keys(sso_data, set(SSO_SECRET_FIELDS)), } diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a1b36e5b7a1..837433fd8cc 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -68,6 +68,7 @@ from litellm.constants import ( SPEND_LOG_QUEUE_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, + VLLM_CHOICE_FIELDS, ) from litellm.cost_calculator import BaseTokenUsageProcessor from litellm.litellm_core_utils.bug_report import ( @@ -170,6 +171,7 @@ from litellm.litellm_core_utils.litellm_logging import ( deployment_pricing_model_info, ) from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.redact_messages import drop_token_ids from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -179,6 +181,7 @@ from litellm.litellm_core_utils.served_output_texts import ( ) from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model from litellm.llms import load_guardrail_translation_mappings +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, @@ -707,6 +710,56 @@ def _pipeline_step_streams_live(step: "PipelineStep") -> bool: return not unified_guardrail.resolve_streaming_flag(callback, "streaming_buffer_until_moderated", True) +def _drop_stream_chunk_token_ids(chunk: object) -> None: + for choice in stream_item_items(chunk, "choices"): + drop_token_ids(choice) + + +def _drop_built_stream_token_ids(request_data: Mapping[str, object]) -> None: + logging_obj: Final = request_data.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.drops_stream_token_ids = True + + +def _drop_token_ids_ahead_of_chunk_check( + callback: CustomLogger, chunk: object, request_data: Mapping[str, object] +) -> None: + if _overrides_hook(callback, "async_post_call_streaming_hook"): + _drop_stream_chunk_token_ids(chunk) + _drop_built_stream_token_ids(request_data) + + +_PLAIN_STREAM_FIELDS: Final = TypeAdapter(dict[str, object]) +_EMPTY_STREAM_VALUES: Final[tuple[object, ...]] = (None, "", [], {}) + + +def _holds_stream_payload(value: object) -> bool: + if isinstance(value, dict): + nested: Final = _PLAIN_STREAM_FIELDS.validate_python(value).values() + return any(item not in _EMPTY_STREAM_VALUES for item in nested) + return value not in _EMPTY_STREAM_VALUES + + +def _choice_carries_only_token_ids(choice: object) -> bool: + if stream_item_field(choice, "token_ids") is None or not isinstance(choice, BaseModel): + return False + fields: Final = _PLAIN_STREAM_FIELDS.validate_python(choice.model_dump()) + return not any(_holds_stream_payload(value) for key, value in fields.items() if key not in ("index", "token_ids")) + + +def _is_token_id_only_chunk(chunk: object) -> bool: + choices: Final = stream_item_items(chunk, "choices") + return bool(choices) and all(_choice_carries_only_token_ids(choice) for choice in choices) + + +async def _without_token_id_only_chunks(response: AsyncIterable[object]) -> AsyncGenerator[object, None]: + async for chunk in response: + if _is_token_id_only_chunk(chunk): + _drop_stream_chunk_token_ids(chunk) + continue + yield chunk + + def _pipelines_stream_live(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> bool: return all(_pipeline_step_streams_live(step) for step in _pipeline_steps(pipelines)) @@ -1707,7 +1760,10 @@ def _proxy_only_error_log_fields( def _failure_fields_to_lift( - request_data: Mapping[str, object], provider_billed: bool = False, dispatched_call_type: str | None = None + request_data: Mapping[str, object], + provider_billed: bool = False, + dispatched_call_type: str | None = None, + route_call_type: str | None = None, ) -> Mapping[str, object]: """Failure-path callbacks run after ``litellm_logging_obj`` is popped from request_data (it is not serialisable), so the caller merges these fields @@ -1717,7 +1773,8 @@ def _failure_fields_to_lift( attribution on failed-request spend logs.""" _logging_obj: Final = request_data.get("litellm_logging_obj") if _logging_obj is None: - return _EMPTY_LIFT + return _EMPTY_LIFT if dispatched_call_type is None else MappingProxyType({"call_type": dispatched_call_type}) + _estimate_call_type: Final = dispatched_call_type or route_call_type _model_call_details: Final = getattr(_logging_obj, "model_call_details", {}) _first_handoff: Final = _model_call_details.get("first_api_call_start_time") _litellm_params: Final = getattr(_logging_obj, "litellm_params", None) @@ -1732,10 +1789,10 @@ def _failure_fields_to_lift( provider_billed=_billed, deployment_id=_deployment_id, base_model=_base_model, - dispatched_call_type=dispatched_call_type, + dispatched_call_type=_estimate_call_type, ) _prompt_estimate: Final = ( - _estimated_prompt_usage(_model_call_details, request_data, dispatched_call_type) + _estimated_prompt_usage(_model_call_details, request_data, _estimate_call_type) if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") else None ) @@ -1755,7 +1812,7 @@ def _failure_fields_to_lift( _billed, _deployment_id, _base_model, - dispatched_call_type, + _estimate_call_type, ) if Logging.usage_capture_callback is not None and (_billed or _model_call_details.get("combined_usage_object") is not None) @@ -4051,13 +4108,15 @@ class ProxyLogging: route=route, original_exception=original_exception, error_type=error_type, + dispatched_call_type=dispatched_call_type, ) request_data.update( await offload_token_count(_failure_fields_to_lift)( request_data, provider_billed=_provider_billed_the_failure(original_exception), - dispatched_call_type=dispatched_call_type or _call_type_for_route(user_api_key_dict.request_route), + dispatched_call_type=dispatched_call_type, + route_call_type=_call_type_for_route(user_api_key_dict.request_route), ) ) @@ -4146,6 +4205,7 @@ class ProxyLogging: route: str | None = None, original_exception: Exception | None = None, error_type: ProxyErrorTypes | None = None, + dispatched_call_type: str | None = None, ): """ Handle logging for proxy only errors by calling `litellm_logging_obj.async_failure_handler` @@ -4245,7 +4305,7 @@ class ProxyLogging: litellm_logging_obj.model_call_details["input"] = input body_shape_call_type = CallTypes.aembedding.value resolved_call_type: Final = _call_type_for_route(route) or ( - None if set_up_by_request else body_shape_call_type + None if set_up_by_request else dispatched_call_type or body_shape_call_type ) if resolved_call_type is not None and litellm_logging_obj.call_type != CallTypes.pass_through.value: litellm_logging_obj.call_type = resolved_call_type @@ -4710,6 +4770,7 @@ class ProxyLogging: else: _callback = callback if _callback is not None and isinstance(_callback, CustomLogger): + _drop_token_ids_ahead_of_chunk_check(_callback, response, data) if str_so_far is not None: complete_response = str_so_far + response_str else: @@ -4731,7 +4792,7 @@ class ProxyLogging: self, response, user_api_key_dict: UserAPIKeyAuth, - request_data: dict, + request_data: dict[str, object], ): """ Allow user to modify outgoing streaming data -> Given a whole response iterator. @@ -4765,20 +4826,26 @@ class ProxyLogging: # Merge model-level guardrails before checking which guardrails to run request_data = check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router) - current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) + running_overrides: Final = tuple( + (callback, kind) + for callback, kind in caps.iterator_overrides + if not isinstance(callback, CustomGuardrail) + or ( + callback.guardrail_name not in pipeline_gated_names + and callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is True + ) + ) + drops_token_ids: Final = bool(post_call_pipelines) or any( + isinstance(callback, CustomGuardrail) for callback, _ in running_overrides + ) + if drops_token_ids: + _drop_built_stream_token_ids(request_data) + current_response = _without_token_id_only_chunks(response) if drops_token_ids else response guarded_layers: Final[list[AsyncGenerator[object, None]]] = [] # mutable-ok: closed on disconnect - for resolved_callback, kind in caps.iterator_overrides: - if isinstance(resolved_callback, CustomGuardrail): - if resolved_callback.guardrail_name in pipeline_gated_names: - continue - if ( - resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) - is not True - ): - continue + for resolved_callback, kind in running_overrides: effective_kind = ( "apply_guardrail" if ( @@ -4836,6 +4903,8 @@ class ProxyLogging: try: async for chunk in current_response: served_chunks.append(chunk) + if drops_token_ids: + _drop_stream_chunk_token_ids(served_chunks[-1]) yield chunk except (GeneratorExit, asyncio.CancelledError): await ProxyLogging._close_guarded_layers(guarded_layers) @@ -10203,7 +10272,7 @@ _PRESERVED_NONE_FIELDS: Final[list[tuple[str, str]]] = [ ("message", "refusal"), ("delta", "content"), # null in streaming chunks ] -_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs",) +_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs", *VLLM_CHOICE_FIELDS) _PRESERVED_NONE_RESPONSE_FIELDS: Final = ("system_fingerprint",) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index e03cb44efda..be1df4aa0f6 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -2,11 +2,13 @@ import asyncio import os +import traceback from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, cast +import httpx from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict, assert_never @@ -14,10 +16,14 @@ import litellm from litellm.constants import ( AZURE_OPENAI_AUDIO_PROVIDERS, REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, + REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, request_timeout, ) +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.xai.common_utils import XAIModelInfo @@ -36,18 +42,24 @@ from litellm.utils import ( ModelInfoLookupMiss, ProviderConfigManager, aresolve_api_key_secret_reference, + get_base_model_from_metadata, lookup_model_info_for_registration, ) -from ..litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params +from ..litellm_core_utils.core_helpers import normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from ..litellm_core_utils.realtime_streaming import client_websocket_scope +from ..litellm_core_utils.realtime_errors import realtime_error_event, websocket_close_reason +from ..litellm_core_utils.realtime_streaming import ( + TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE, + client_websocket_scope, + translation_guardrails_apply, +) from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client from ..llms.bedrock.realtime.handler import BedrockRealtime, raise_if_session_committed -from ..llms.custom_httpx.http_handler import realtime_ssl_for_url +from ..llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, realtime_ssl_for_url from ..llms.openai.common_utils import is_openai_backed_api_base from ..llms.openai.live.handler import OpenAILiveSessions from ..llms.openai.realtime.handler import ( @@ -66,8 +78,6 @@ from ..utils import client as wrapper_client if TYPE_CHECKING: from fastapi import WebSocket - from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig - azure_realtime: Final = AzureOpenAIRealtime() openai_realtime: Final = OpenAIRealtime() openai_live_sessions: Final = OpenAILiveSessions() @@ -77,7 +87,7 @@ vertex_llm_base: Final = VertexBase() base_llm_http_handler = BaseLLMHTTPHandler() _EMPTY_MODEL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) _EMPTY_AUTH_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) -_DEPLOYMENT_PARAMS: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) def _model_params_with_stored_credentials(model_params: Mapping[str, object]) -> Mapping[str, object]: @@ -90,26 +100,145 @@ def _model_params_with_stored_credentials(model_params: Mapping[str, object]) -> return MappingProxyType({**credential_values, **model_params}) -def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]: - if "model" not in session: - return session - return {**session, "model": model_name} +def _transcription_model_overrides(session: Mapping[str, object], model_name: str) -> Mapping[str, object]: + flat_transcription: Final = as_str_mapping(session.get("input_audio_transcription")) + audio: Final = as_str_mapping(session.get("audio")) or {} + audio_input: Final = as_str_mapping(audio.get("input")) or {} + nested_transcription: Final = as_str_mapping(audio_input.get("transcription")) + flat_part: Final = ( + {} if flat_transcription is None else {"input_audio_transcription": {**flat_transcription, "model": model_name}} + ) + nested_part: Final = ( + {} + if nested_transcription is None and flat_transcription is not None + else { + "audio": { + **audio, + "input": {**audio_input, "transcription": {**(nested_transcription or {}), "model": model_name}}, + } + } + ) + return MappingProxyType({**flat_part, **nested_part}) -def _resolved_request_session(request_data: Mapping[str, object], model_name: str) -> Mapping[str, object] | None: - session: Final = as_str_mapping(request_data.get("session")) - if not isinstance(session, dict): - return None - return _with_resolved_session_model(session, model_name) +def with_resolved_session_model(session: Mapping[str, object], model_name: str) -> dict[str, object]: + top_level_model: Final = {} if "model" not in session else {"model": model_name} + transcription_models: Final = ( + _transcription_model_overrides(session, model_name) if session.get("type") == "transcription" else {} + ) + return {**session, **top_level_model, **transcription_models} -def _build_litellm_metadata(kwargs: dict) -> dict[str, object]: - """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" - metadata: Final[dict[str, object]] = {**(kwargs.get("litellm_metadata") or {})} - guardrails: Final = (kwargs.get("metadata") or {}).get("guardrails") or kwargs.get("guardrails") or [] - if guardrails: - metadata["guardrails"] = guardrails - return metadata +def _without_top_level_nulls(fields: Mapping[str, object]) -> dict[str, object]: + return {key: value for key, value in fields.items() if value is not None} + + +_POLICY_PIPELINE_METADATA_KEYS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails") + + +def _build_litellm_metadata(kwargs: Mapping[str, object]) -> dict[str, object]: + request_metadata: Final = _OBJECT_MAPPING.validate_python(kwargs.get("metadata") or {}) + internal_metadata: Final = _OBJECT_MAPPING.validate_python(kwargs.get("litellm_metadata") or {}) + policy_pipeline_state: Final = { + key: request_metadata[key] for key in _POLICY_PIPELINE_METADATA_KEYS if key in request_metadata + } + guardrails: Final = request_metadata.get("guardrails") or kwargs.get("guardrails") + return {**policy_pipeline_state, **internal_metadata, **({"guardrails": guardrails} if guardrails else {})} + + +_GUARDRAIL_OPT_OUT_METADATA_KEYS: Final = ("user_api_key_metadata", "user_api_key_team_metadata") + + +def _translation_litellm_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: + request_metadata: Final = _OBJECT_MAPPING.validate_python(kwargs.get("metadata") or {}) + guardrail_opt_outs: Final = { + key: request_metadata[key] for key in _GUARDRAIL_OPT_OUT_METADATA_KEYS if key in request_metadata + } + return MappingProxyType({**guardrail_opt_outs, **_build_litellm_metadata(kwargs)}) + + +_HTTP_TRANSLATION_PROVIDERS: Final = frozenset({LlmProviders.OPENAI.value, LlmProviders.AZURE.value}) +_WEBSOCKET_TRANSLATION_PROVIDERS: Final = frozenset({LlmProviders.OPENAI.value, *AZURE_OPENAI_AUDIO_PROVIDERS}) + + +@dataclass(frozen=True, slots=True) +class _TranslationProviderUnsupported: + provider: str + supported_providers: frozenset[str] + + +@dataclass(frozen=True, slots=True) +class _TranslationGuardrailsUnsupported: + pass + + +_TranslationRejection = _TranslationProviderUnsupported | _TranslationGuardrailsUnsupported + + +def _translation_rejection( + custom_llm_provider: str, supported_providers: frozenset[str], litellm_metadata: Mapping[str, object] +) -> _TranslationRejection | None: + if custom_llm_provider not in supported_providers: + return _TranslationProviderUnsupported(provider=custom_llm_provider, supported_providers=supported_providers) + if translation_guardrails_apply({"litellm_metadata": dict(litellm_metadata)}): + return _TranslationGuardrailsUnsupported() + return None + + +def _translation_rejection_message(rejection: _TranslationRejection) -> str: + match rejection: + case _TranslationProviderUnsupported(provider=provider, supported_providers=supported_providers): + return ( + f"Realtime translation supports only the {', '.join(sorted(supported_providers))} providers, " + f"got {provider}" + ) + case _TranslationGuardrailsUnsupported(): + return TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE + case _: + assert_never(rejection) + + +def _translation_rejection_error_type(rejection: _TranslationRejection) -> str: + match rejection: + case _TranslationProviderUnsupported(): + return "invalid_request_error" + case _TranslationGuardrailsUnsupported(): + return "guardrail_error" + case _: + assert_never(rejection) + + +def _translation_rejection_error( + rejection: _TranslationRejection, model: str, custom_llm_provider: str +) -> BadRequestError: + return BadRequestError( + message=_translation_rejection_message(rejection), model=model, llm_provider=custom_llm_provider + ) + + +def _raise_if_http_translation_rejected(custom_llm_provider: str, model: str, kwargs: Mapping[str, object]) -> None: + rejection: Final = _translation_rejection( + custom_llm_provider, _HTTP_TRANSLATION_PROVIDERS, _translation_litellm_metadata(kwargs) + ) + if rejection is None: + return + raise _translation_rejection_error(rejection, model, custom_llm_provider) + + +async def _refuse_translation_session( + websocket: "WebSocket", + rejection: _TranslationRejection, + logging_obj: LiteLLMLogging, + model: str, + custom_llm_provider: str, +) -> None: + message: Final = _translation_rejection_message(rejection) + await websocket.send_text(realtime_error_event(message, error_type=_translation_rejection_error_type(rejection))) + await websocket.close(code=1008, reason=websocket_close_reason(message, fallback="Translation session refused")) + error: Final = _translation_rejection_error(rejection, model, custom_llm_provider) + await logging_obj.dispatch_failure_handlers( + error, "".join(traceback.format_exception(error)), prefer_async_handlers=True + ) def _get_realtime_http_provider_config( @@ -125,10 +254,6 @@ def _get_realtime_http_provider_config( Uses ProviderConfigManager so each provider keeps its credential-resolution and URL-construction logic in its own transformation class. """ - from litellm.llms.base_llm.realtime.http_transformation import ( - BaseRealtimeHTTPConfig, - ) - provider_config: BaseRealtimeHTTPConfig | None = None if custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_realtime_http_config( @@ -163,6 +288,32 @@ async def _realtime_http_bearer_token( return bearer_token or "" +async def _get_realtime_http_extra_headers( + provider_config: BaseRealtimeHTTPConfig | None, + litellm_params: GenericLiteLLMParams, + resolved_api_key: str, + extra_headers: Mapping[str, object] | None, +) -> Mapping[str, object] | None: + if provider_config is None: + return extra_headers + return await asyncio.to_thread( + provider_config.get_extra_headers, + litellm_params=litellm_params, + api_key=resolved_api_key, + extra_headers=extra_headers, + ) + + +def _client_secret_session_body(req: RealtimeClientSecretRequest, model_name: str) -> Mapping[str, object] | None: + if req.session is not None: + return with_resolved_session_model( + _without_top_level_nulls(req.session.model_dump(exclude_unset=True)), model_name + ) + if req.model is not None: + return {"type": "realtime", "model": model_name} + return None + + @wrapper_client async def acreate_realtime_client_secret( model: str | None = None, @@ -176,31 +327,42 @@ async def acreate_realtime_client_secret( session=RealtimeSessionConfig.model_validate(session) if session else None, expires_after=RealtimeExpiresAfter.model_validate(expires_after) if expires_after else None, ) - model_name = (req.session.model if req.session is not None else None) or req.model or "gpt-4o-realtime-preview" + transcription_model: Final = ( + req.session.audio.input.transcription.model + if req.session is not None + and req.session.type == "transcription" + and req.session.audio is not None + and req.session.audio.input is not None + and req.session.audio.input.transcription is not None + else None + ) + requested_model_name: Final = ( + transcription_model + or (req.session.model if req.session is not None else None) + or req.model + or "gpt-4o-realtime-preview" + ) litellm_logging_obj: Final[LiteLLMLogging] = kwargs.get("litellm_logging_obj") litellm_params: Final = GenericLiteLLMParams(**kwargs) - ( - model_name, - custom_llm_provider, - dynamic_api_key, - dynamic_api_base, - ) = get_llm_provider( - model=model_name, + model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( + model=requested_model_name, custom_llm_provider=litellm_params.custom_llm_provider, api_base=litellm_params.api_base, api_key=litellm_params.api_key, ) - ( - provider_config, - resolved_api_base, - resolved_api_key, - ) = _get_realtime_http_provider_config( + provider_config, resolved_api_base, resolved_api_key = _get_realtime_http_provider_config( custom_llm_provider=custom_llm_provider, dynamic_api_base=dynamic_api_base, dynamic_api_key=dynamic_api_key, litellm_params=litellm_params, ) + resolved_extra_headers: Final = await _get_realtime_http_extra_headers( + provider_config=provider_config, + litellm_params=litellm_params, + resolved_api_key=resolved_api_key, + extra_headers=kwargs.get("extra_headers"), + ) litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model_name, @@ -208,26 +370,124 @@ async def acreate_realtime_client_secret( litellm_params={"api_base": resolved_api_base}, custom_llm_provider=custom_llm_provider, ) - request_data: Final = req.model_dump(exclude_none=True, exclude={"model"}) - resolved_session: Final = _resolved_request_session(request_data, model_name) - if resolved_session is not None: - request_data["session"] = resolved_session + session_body: Final = _client_secret_session_body(req, model_name) return await base_llm_http_handler.async_realtime_client_secret_handler( api_base=resolved_api_base, api_key=await _realtime_http_bearer_token( custom_llm_provider, resolved_api_key, resolved_api_base, litellm_params ), - request_data=request_data, + request_data={ + **req.model_dump(exclude_none=True, exclude={"model", "session"}), + **({} if session_body is None else {"session": session_body}), + }, logging_obj=litellm_logging_obj, timeout=timeout or request_timeout, provider_config=provider_config, model=model_name, - extra_headers=kwargs.get("extra_headers"), + extra_headers=resolved_extra_headers, client=kwargs.get("client"), api_version=litellm_params.api_version, ) +@wrapper_client +async def acreate_realtime_translation_client_secret( + model: str | None = None, + session: Mapping[str, Any] | None = None, + expires_after: Mapping[str, Any] | None = None, + timeout: float | None = None, + **kwargs: object, # kwargs-ok: public API preserves provider-specific request options +) -> httpx.Response: + requested_model_name: Final = model or (session or {}).get("model") or "gpt-realtime-translate" + session_config: Final = RealtimeSessionConfig.model_validate( + { + **(session or {}), + "type": "translation", + "model": requested_model_name, + } + ) + req: Final = RealtimeClientSecretRequest( + model=requested_model_name, + session=session_config, + expires_after=RealtimeExpiresAfter(**expires_after) if expires_after else None, + ) + litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(litellm_logging_obj, LiteLLMLogging): + raise TypeError("litellm_logging_obj must be a LiteLLM Logging instance") + litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs) + + model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( + model=requested_model_name, + custom_llm_provider=litellm_params.custom_llm_provider, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + ) + _raise_if_http_translation_rejected(custom_llm_provider, model_name, kwargs) + provider_config, resolved_api_base, resolved_api_key = _get_realtime_http_provider_config( + custom_llm_provider=custom_llm_provider, + dynamic_api_base=dynamic_api_base, + dynamic_api_key=dynamic_api_key, + litellm_params=litellm_params, + ) + resolved_extra_headers: Final = await _get_realtime_http_extra_headers( + provider_config=provider_config, + litellm_params=litellm_params, + resolved_api_key=resolved_api_key, + extra_headers=as_str_mapping(kwargs.get("extra_headers")), + ) + litellm_logging_obj.update_from_kwargs( + kwargs=kwargs, + model=model_name, + optional_params={ + "expires_after": expires_after, + "session": session, + }, + litellm_params={"api_base": resolved_api_base}, + custom_llm_provider=custom_llm_provider, + ) + translation_session: Final = with_resolved_session_model( + _without_top_level_nulls(session_config.model_dump(exclude_unset=True)), model_name + ) + return await base_llm_http_handler.async_realtime_translation_client_secret_handler( + api_base=resolved_api_base, + api_key=await _realtime_http_bearer_token( + custom_llm_provider, resolved_api_key, resolved_api_base, litellm_params + ), + request_data={ + **req.model_dump(exclude_none=True, exclude={"model", "session"}), + "session": {key: value for key, value in translation_session.items() if key != "type"}, + }, + logging_obj=litellm_logging_obj, + timeout=timeout or request_timeout, + provider_config=provider_config, + model=model_name, + extra_headers=resolved_extra_headers, + client=client if isinstance(client := kwargs.get("client"), (HTTPHandler, AsyncHTTPHandler)) else None, + api_version=litellm_params.api_version, + ) + + +def _with_upstream_transcription_model(request_data: Mapping[str, object], model_name: str) -> dict[str, object]: + flat: Final = as_str_mapping(request_data.get("input_audio_transcription")) + audio: Final = as_str_mapping(request_data.get("audio")) or {} + audio_input: Final = as_str_mapping(audio.get("input")) or {} + nested: Final = as_str_mapping(audio_input.get("transcription")) + return { + **request_data, + **({} if flat is None else {"input_audio_transcription": {**flat, "model": model_name}}), + **( + {} + if nested is None + else {"audio": {**audio, "input": {**audio_input, "transcription": {**nested, "model": model_name}}}} + ), + } + + +def _deployment_base_model(kwargs: Mapping[str, object], custom_llm_provider: str) -> str | None: + base_model: Final = get_base_model_from_metadata(model_call_details={"litellm_params": kwargs}) + return None if base_model is None else base_model.removeprefix(f"{custom_llm_provider}/") + + @wrapper_client async def acreate_realtime_transcription_session( model: str | None = None, @@ -273,6 +533,12 @@ async def acreate_realtime_transcription_session( dynamic_api_key=dynamic_api_key, litellm_params=litellm_params, ) + resolved_extra_headers: Final = await _get_realtime_http_extra_headers( + provider_config=provider_config, + litellm_params=litellm_params, + resolved_api_key=resolved_api_key, + extra_headers=kwargs.get("extra_headers"), + ) litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model_name, @@ -280,13 +546,11 @@ async def acreate_realtime_transcription_session( litellm_params={"api_base": resolved_api_base}, custom_llm_provider=custom_llm_provider, ) - request_data: Final = req.model_dump(exclude_none=True, exclude={"model"}) - # Ensure the upstream body's input_audio_transcription.model matches the - # authorized routing model. This prevents a caller from supplying an allowed - # top-level model for auth while sneaking a different model into the nested - # transcription config that gets forwarded to the provider. - if isinstance(request_data.get("input_audio_transcription"), dict): - request_data["input_audio_transcription"]["model"] = model_name + # The upstream transcription model must match the authorized routing model, so a + # caller cannot pass auth with an allowed top-level model and forward a different one. + request_data: Final = _with_upstream_transcription_model( + _without_top_level_nulls(req.model_dump(exclude_unset=True, exclude={"model"})), model_name + ) return await base_llm_http_handler.async_realtime_transcription_session_handler( api_base=resolved_api_base, api_key=await _realtime_http_bearer_token( @@ -297,7 +561,7 @@ async def acreate_realtime_transcription_session( timeout=timeout or request_timeout, provider_config=provider_config, model=model_name, - extra_headers=kwargs.get("extra_headers"), + extra_headers=resolved_extra_headers, client=kwargs.get("client"), api_version=litellm_params.api_version, ) @@ -334,7 +598,7 @@ async def arealtime_calls( litellm_params=litellm_params, ) if session is not None: - session = _with_resolved_session_model(session, model_name) + session = with_resolved_session_model(session, model_name) litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model_name, @@ -357,6 +621,68 @@ async def arealtime_calls( ) +@wrapper_client +async def arealtime_translation_calls( + openai_ephemeral_key: str, + sdp_body: bytes, + model: str | None = None, + session: Mapping[str, Any] | None = None, + timeout: float | None = None, + **kwargs: object, # kwargs-ok: public API preserves provider-specific request options +) -> httpx.Response: + requested_model_name: Final = model or "gpt-realtime-translate" + litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(litellm_logging_obj, LiteLLMLogging): + raise TypeError("litellm_logging_obj must be a LiteLLM Logging instance") + litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs) + + model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( + model=requested_model_name, + custom_llm_provider=litellm_params.custom_llm_provider, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + ) + _raise_if_http_translation_rejected(custom_llm_provider, model_name, kwargs) + provider_config, resolved_api_base, _ = _get_realtime_http_provider_config( + custom_llm_provider=custom_llm_provider, + dynamic_api_base=dynamic_api_base, + dynamic_api_key=dynamic_api_key, + litellm_params=litellm_params, + ) + session_config: Final = with_resolved_session_model( + { + **(session or {}), + "type": "translation", + "model": model_name, + }, + model_name, + ) + litellm_logging_obj.update_from_kwargs( + kwargs=kwargs, + model=model_name, + optional_params={ + "realtime_translation_calls": True, + "session": session_config, + }, + litellm_params={"api_base": resolved_api_base}, + custom_llm_provider=custom_llm_provider, + ) + return await base_llm_http_handler.async_realtime_calls_handler( + api_base=resolved_api_base, + openai_ephemeral_key=openai_ephemeral_key, + sdp_body=sdp_body, + logging_obj=litellm_logging_obj, + timeout=timeout or request_timeout, + provider_config=provider_config, + model=model_name, + session_config=session_config, + extra_headers=as_str_mapping(kwargs.get("extra_headers")), + client=client if isinstance(client := kwargs.get("client"), (HTTPHandler, AsyncHTTPHandler)) else None, + api_version=litellm_params.api_version, + translation=True, + ) + + async def vertex_access_token_resolver( credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, @@ -486,6 +812,7 @@ async def _arealtime( timeout: float | None = None, query_params: RealtimeQueryParams | None = None, live_session_start: Mapping[str, object] | None = None, + realtime_mode: str = "realtime", **kwargs, ) -> None: """ @@ -530,7 +857,24 @@ async def _arealtime( custom_llm_provider=_custom_llm_provider, ) - deployment_params: Final = _DEPLOYMENT_PARAMS.validate_python(kwargs) + session_litellm_metadata: Final = ( + dict(_translation_litellm_metadata(kwargs)) + if realtime_mode == "translation" + else _build_litellm_metadata(kwargs) + ) + translation_rejection: Final = ( + _translation_rejection(_custom_llm_provider, _WEBSOCKET_TRANSLATION_PROVIDERS, session_litellm_metadata) + if realtime_mode == "translation" + else None + ) + if translation_rejection is not None: + await _refuse_translation_session( + websocket, translation_rejection, litellm_logging_obj, model, _custom_llm_provider + ) + litellm_logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True + return + + deployment_params: Final = _OBJECT_MAPPING.validate_python(kwargs) if live_session_start is not None: await _alive_session( model=model, @@ -540,7 +884,7 @@ async def _arealtime( session_start=live_session_start, api_base=dynamic_api_base or litellm_params.api_base, api_key=dynamic_api_key, - litellm_metadata=_build_litellm_metadata(kwargs), + litellm_metadata=session_litellm_metadata, litellm_params=deployment_params, ) return @@ -566,7 +910,7 @@ async def _arealtime( timeout=timeout, headers=headers, user_api_key_dict=kwargs.get("user_api_key_dict"), - litellm_metadata=_build_litellm_metadata(kwargs), + litellm_metadata=session_litellm_metadata, query_params=query_params, ) elif _custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: @@ -581,9 +925,6 @@ async def _arealtime( or litellm_params.get("realtime_protocol") or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") ) - realtime_protocol: Final = azure_realtime_protocol_for_client( - configured_realtime_protocol, query_params=query_params, websocket=websocket - ) resolved_azure_ad_token: Final = ( None if api_key @@ -592,6 +933,14 @@ async def _arealtime( GenericLiteLLMParams(**kwargs, azure_ad_token=azure_ad_token), ) ) + realtime_protocol: Final = azure_realtime_protocol_for_client( + configured_realtime_protocol, + model=model, + base_model=_deployment_base_model(kwargs, _custom_llm_provider), + query_params=query_params, + realtime_mode=realtime_mode, + websocket=websocket, + ) await azure_realtime.async_realtime( model=model, websocket=websocket, @@ -602,10 +951,13 @@ async def _arealtime( client=None, timeout=timeout, logging_obj=litellm_logging_obj, - realtime_protocol=realtime_protocol, + realtime_protocol=realtime_protocol.name, query_params=query_params, + realtime_mode=realtime_mode, user_api_key_dict=kwargs.get("user_api_key_dict"), - litellm_metadata=_build_litellm_metadata(kwargs), + litellm_metadata=session_litellm_metadata, + bridges_beta_client_to_ga=realtime_protocol.bridges_beta_client_to_ga, + drop_params=drop_params, ) elif _custom_llm_provider == "openai": api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or "https://api.openai.com/" @@ -622,8 +974,10 @@ async def _arealtime( client=None, timeout=timeout, query_params=query_params, + realtime_mode=realtime_mode, user_api_key_dict=kwargs.get("user_api_key_dict"), - litellm_metadata=_build_litellm_metadata(kwargs), + litellm_metadata=session_litellm_metadata, + drop_params=drop_params, ) elif _custom_llm_provider == "bedrock": # Extract AWS parameters from kwargs @@ -678,7 +1032,7 @@ async def _arealtime( timeout=timeout, query_params=query_params, user_api_key_dict=kwargs.get("user_api_key_dict"), - litellm_metadata=_build_litellm_metadata(kwargs), + litellm_metadata=session_litellm_metadata, ) elif _custom_llm_provider == "vertex_ai": vertex_credentials: Final = ( @@ -738,7 +1092,7 @@ async def _arealtime( timeout=timeout, headers=headers, user_api_key_dict=kwargs.get("user_api_key_dict"), - litellm_metadata=_build_litellm_metadata(kwargs), + litellm_metadata=session_litellm_metadata, query_params=query_params, ) else: diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 7cf163e7ce8..7e6ee84e2b2 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,14 +1,16 @@ import asyncio import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from itertools import chain +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.responses.litellm_completion_transformation.client_tool_search import ( bridge_client_tool_search, client_tool_search_loaded_tools, @@ -69,6 +71,90 @@ def _is_a_conversation_turn(spend_log: "SpendLogsPayload") -> bool: return spend_log.get("call_type") not in _MCP_TOOL_CALL_TYPES and not _is_a_failed_turn(spend_log) +_RESPONSE_READ_CALL_TYPES: Final = frozenset(("aget_responses", "acancel_responses")) + + +def _output_item_id(item: object) -> str | None: + match item: + case {"id": str() as item_id}: + return item_id + case _: + return None + + +def _is_a_response_read(spend_log: "SpendLogsPayload") -> bool: + return "call_type" in spend_log and spend_log["call_type"] in _RESPONSE_READ_CALL_TYPES + + +def _stored_response(spend_log: "SpendLogsPayload") -> object: + return spend_log.get("response") + + +def _output_items(spend_log: "SpendLogsPayload") -> Iterator[tuple[str | None, object]]: + output: Final = _stored_responses_api_output(_stored_response(spend_log)) or [] + return zip(map(_output_item_id, output), output) + + +def _identified_output_items(spend_log: "SpendLogsPayload") -> Iterator[tuple[str, object]]: + return ((item_id, item) for item_id, item in _output_items(spend_log) if item_id is not None) + + +def _output_item_ids(spend_log: "SpendLogsPayload") -> frozenset[str]: + return frozenset(item_id for item_id, _ in _identified_output_items(spend_log)) + + +def _first_index_holding_each_item(indexed_turns: Iterable[tuple[int, "SpendLogsPayload"]]) -> Mapping[str, int]: + return dict( + chain.from_iterable( + ((item_id, index) for item_id in _output_item_ids(turn)) for index, turn in reversed(tuple(indexed_turns)) + ) + ) + + +def _turn_replaying_each_item(turns: Sequence["SpendLogsPayload"]) -> Mapping[str, int]: + indexed_turns: Final = tuple(enumerate(turns)) + first_read_holding: Final = _first_index_holding_each_item( + (index, turn) for index, turn in indexed_turns if _is_a_response_read(turn) + ) + first_other_turn_holding: Final = _first_index_holding_each_item( + (index, turn) for index, turn in indexed_turns if not _is_a_response_read(turn) + ) + return MappingProxyType({**first_read_holding, **first_other_turn_holding}) + + +def _latest_read_copy_of_each_item(turns: Sequence["SpendLogsPayload"]) -> Mapping[str, object]: + return MappingProxyType( + dict(chain.from_iterable(_identified_output_items(turn) for turn in turns if _is_a_response_read(turn))) + ) + + +def _with_only_new_output( + index: int, turn: "SpendLogsPayload", turn_replaying: Mapping[str, int], latest_copy: Mapping[str, object] +) -> "SpendLogsPayload | None": + if not _is_a_response_read(turn): + return turn + stored_response: Final = as_str_mapping(_stored_response(turn)) + new_output: Final = [ + item if item_id is None else latest_copy[item_id] + for item_id, item in _output_items(turn) + if item_id is None or turn_replaying.get(item_id) == index + ] + if stored_response is None or not new_output: + return None + replayed_turn: Final[SpendLogsPayload] = {**turn, "response": {**stored_response, "output": new_output}} + return replayed_turn + + +def _conversation_turns(spend_logs: Sequence["SpendLogsPayload"]) -> tuple["SpendLogsPayload", ...]: + turns: Final = tuple(filter(_is_a_conversation_turn, spend_logs)) + turn_replaying: Final = _turn_replaying_each_item(turns) + latest_copy: Final = _latest_read_copy_of_each_item(turns) + replayed: Final = ( + _with_only_new_output(index, turn, turn_replaying, latest_copy) for index, turn in enumerate(turns) + ) + return tuple(turn for turn in replayed if turn is not None) + + _MCP_AUTO_EXECUTE_OUTPUT_TYPES: Final = frozenset(("mcp_tools_fetched", "tool_execution_results")) @@ -101,7 +187,7 @@ class StoredTurns: def _compaction_point( index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None ) -> _ReplayStart | None: - if response_holds_compaction_block(spend_log.get("response")): + if response_holds_compaction_block(_stored_response(spend_log)): return _ReplayStart(index=index, at_output=True) if proxy_server_request is not None and holds_compaction_summary(proxy_server_request.get("input")): return _ReplayStart(index=index, at_output=False) @@ -271,12 +357,10 @@ class ResponsesSessionHandler: previous_response_id: str, client_response_ids: tuple[str, ...] = (), ) -> StoredTurns: - spend_logs: Final = tuple( - spend_log - for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + spend_logs: Final = _conversation_turns( + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( previous_response_id, client_response_ids ) - if _is_a_conversation_turn(spend_log) ) return StoredTurns( spend_logs=spend_logs, @@ -378,7 +462,7 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final[object] = spend_log.get("response", "{}") + _response_output: Final = _stored_response(spend_log) if (responses_api_output := _stored_responses_api_output(_response_output)) is not None: chat_completion_message_history.extend( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index ddcdc66ff20..521a698b276 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -27,7 +27,11 @@ from litellm.responses.litellm_completion_transformation.tool_search import ( from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) -from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost +from litellm.responses.streaming_iterator import ( + SEQUENCE_NUMBER_STAMPED_ON_EMISSION, + ResponsesAPIStreamingIterator, + stamp_responses_usage_cost, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, @@ -239,14 +243,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return execution def _next_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> OutputItemAddedEvent: - self._sequence_number += 1 - added: Final = OutputItemAddedEvent( + return OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, item=item, ) - added.__dict__["sequence_number"] = self._sequence_number - return added def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: self._pending_tool_events.append(self._next_output_item_added(output_index, item)) @@ -268,12 +269,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): item: Final = build_tool_search_call( call_id, final_args or self._tool_args_by_call_id.get(call_id, ""), "completed", execution ) - self._sequence_number += 1 self._pending_tool_events.append( OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, - sequence_number=self._sequence_number, item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()), ) ) @@ -346,12 +345,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if streamed_item is None: return self._sent_compaction_item_done = True - self._sequence_number += 1 self._pending_response_events.append( OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=self._compaction_output_index, - sequence_number=self._sequence_number, item=streamed_item, ) ) @@ -494,15 +491,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chunk_size = 10 # Match typical OpenAI delta size for i in range(0, len(fn_args_delta), chunk_size): delta_chunk = fn_args_delta[i : i + chunk_size] - self._sequence_number += 1 delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, delta=delta_chunk, ) - # Add sequence_number as extra field (BaseLiteLLMOpenAIResponseObject allows extra fields) - delta_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(delta_event) def _queue_final_tool_call_done_events(self, litellm_complete_object: ModelResponse) -> None: @@ -567,33 +561,27 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chunk_size = 10 # Match typical OpenAI delta size for i in range(0, len(remaining_args), chunk_size): delta_chunk = remaining_args[i : i + chunk_size] - self._sequence_number += 1 delta_event = FunctionCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, delta=delta_chunk, ) - delta_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(delta_event) - self._sequence_number += 1 done_event = FunctionCallArgumentsDoneEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, arguments=final_args, ) - done_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(done_event) - self._sequence_number += 1 item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, final_args, "completed") item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"]) item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, - sequence_number=self._sequence_number, item=BaseLiteLLMOpenAIResponseObject( **item_kwargs, **(_tool_call_provider_specific_fields(tc) if self.tool_call_provider_fields else {}), @@ -625,16 +613,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): _WEB_SEARCH_PROGRESS_EVENTS[:1] if item.status == "in_progress" else _WEB_SEARCH_PROGRESS_EVENTS ) for event_type, event_class in progress_events: - self._sequence_number += 1 - event = event_class(type=event_type, output_index=output_index, item_id=item.id) - event.__dict__["sequence_number"] = self._sequence_number - self._pending_tool_events.append(event) - self._sequence_number += 1 + self._pending_tool_events.append(event_class(type=event_type, output_index=output_index, item_id=item.id)) self._pending_tool_events.append( OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, - sequence_number=self._sequence_number, item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()), ) ) @@ -770,31 +753,24 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): """ response_created_event_data: Final = self._default_response_created_event_data() - self._sequence_number += 1 - event: Final = ResponseCreatedEvent( + return ResponseCreatedEvent( type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=self._with_encoded_response_id(ResponsesAPIResponse(**response_created_event_data)), ) - event.__dict__["sequence_number"] = self._sequence_number - return event def create_response_in_progress_event(self) -> ResponseInProgressEvent: response_in_progress_event_data: Final = self._default_response_created_event_data() response_in_progress_event_data["status"] = "in_progress" - self._sequence_number += 1 - event: Final = ResponseInProgressEvent( + return ResponseInProgressEvent( type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, response=self._with_encoded_response_id(ResponsesAPIResponse(**response_in_progress_event_data)), ) - event.__dict__["sequence_number"] = self._sequence_number - return event def create_output_item_added_event(self) -> OutputItemAddedEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{uuid.uuid4()}" - self._sequence_number += 1 - event: Final = OutputItemAddedEvent( + return OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=self._message_index(), item=BaseLiteLLMOpenAIResponseObject( @@ -807,23 +783,18 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__["sequence_number"] = self._sequence_number - return event def create_content_part_added_event(self) -> ContentPartAddedEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{uuid.uuid4()}" - self._sequence_number += 1 - event: Final = ContentPartAddedEvent( + return ContentPartAddedEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, item_id=self._cached_item_id, output_index=self._message_index(), content_index=0, part=BaseLiteLLMOpenAIResponseObject(**{"type": "output_text", "text": "", "annotations": []}), ) - event.__dict__["sequence_number"] = self._sequence_number - return event def _queue_message_item_added_events(self) -> None: if not self.sent_message_item_added_event: @@ -834,7 +805,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _queue_message_item_added_event(self) -> None: item_id: Final = self._message_item_id() self.sent_message_item_added_event = True - self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=self._message_index(), @@ -848,7 +818,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__["sequence_number"] = self._sequence_number self._pending_response_events.append(event) def _queue_refusal_events(self, refusal: str) -> None: @@ -885,12 +854,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): part=ContentPartDonePartRefusal(type="refusal", refusal=refusal), ), ) - first_sequence_number: Final = self._sequence_number + 1 - self._sequence_number += len(events) - self._pending_response_events.extend( - event.model_copy(update={"sequence_number": sequence_number}) - for sequence_number, event in enumerate(events, start=first_sequence_number) - ) + self._pending_response_events.extend(events) def _merge_provider_specific_fields(self, src: dict) -> None: """Merge provider_specific_fields using last-value-wins for lists. @@ -944,7 +908,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self, reasoning_item_id: str, reasoning_content: str, - sequence_number: int, ) -> ReasoningSummaryTextDoneEvent: """ Create response.reasoning_summary_text.done event. @@ -954,7 +917,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "type": "response.reasoning_summary_text.done", "item_id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c", "output_index": 0, - "sequence_number": 97, "summary_index": 0, "text": "**Clarifying the first humans**\n\nThe I'm addressing the user's specific interest." } @@ -963,7 +925,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, item_id=reasoning_item_id, output_index=self._reasoning_index(), - sequence_number=sequence_number, + sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION, summary_index=0, text=reasoning_content, ) @@ -972,7 +934,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self, reasoning_item_id: str, reasoning_content: str, - sequence_number: int, ) -> ReasoningSummaryPartDoneEvent: """ Create response.reasoning_summary_part.done event. @@ -986,7 +947,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "type": "summary_text", "text": "**Clarifying the first humans**\n\nThe earlier hominins. It feels important to ensure I'm addressing the user's specific interest." }, - "sequence_number": 98, "summary_index": 0 } """ @@ -994,7 +954,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): type=ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, item_id=reasoning_item_id, output_index=self._reasoning_index(), - sequence_number=sequence_number, + sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION, summary_index=0, part=BaseLiteLLMOpenAIResponseObject( **{ @@ -1066,7 +1026,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=self._message_index(), - sequence_number=1, item=BaseLiteLLMOpenAIResponseObject( **{ "id": self._cached_item_id, @@ -1082,7 +1041,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self, reasoning_item_id: str, reasoning_content: str, - sequence_number: int, ) -> OutputItemDoneEvent: """ Create response.output_item.done event for reasoning items. @@ -1091,7 +1049,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): { "type": "response.output_item.done", "output_index": 0, - "sequence_number": 99, "item": { "id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c", "type": "reasoning", @@ -1107,7 +1064,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=self._reasoning_index(), - sequence_number=sequence_number, item=BaseLiteLLMOpenAIResponseObject( **{ "id": reasoning_item_id, @@ -1223,7 +1179,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if chunk.choices[0].finish_reason is None and not _delta_carries_output(delta): return - self._sequence_number += 1 self.sent_output_item_added_event = True # Reasoning-first @@ -1233,20 +1188,20 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_reasoning_item_id = mint_reasoning_item_id() self._reasoning_item_id = self._cached_reasoning_item_id - event = OutputItemAddedEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=self._reasoning_index(), - item=BaseLiteLLMOpenAIResponseObject( - **{ - "id": self._cached_reasoning_item_id, - "type": "reasoning", - "status": "in_progress", - "summary": None, - } - ), + self._pending_response_events.append( + OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=self._reasoning_index(), + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": self._cached_reasoning_item_id, + "type": "reasoning", + "status": "in_progress", + "summary": None, + } + ), + ) ) - event.__dict__["sequence_number"] = self._sequence_number - self._pending_response_events.append(event) return # Tool-first @@ -1279,30 +1234,34 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) reasoning_item_id: Final = self._cached_reasoning_item_id self._queue_compaction_item_done() - self._sequence_number += 1 text_done_event: Final = self.create_reasoning_summary_text_done_event( reasoning_item_id=reasoning_item_id, reasoning_content=reasoning_content, - sequence_number=self._sequence_number, ) - self._sequence_number += 1 part_done_event: Final = self.create_reasoning_summary_part_done_event( reasoning_item_id=reasoning_item_id, reasoning_content=reasoning_content, - sequence_number=self._sequence_number, ) - self._sequence_number += 1 reasoning_output_item_done_event: Final = self.create_reasoning_output_item_done_event( reasoning_item_id=reasoning_item_id, reasoning_content=reasoning_content, - sequence_number=self._sequence_number, ) self._pending_response_events.extend([text_done_event, part_done_event, reasoning_output_item_done_event]) self._reasoning_done_emitted = True self._reasoning_active = False + def _stamp_sequence_number(self, event: BaseLiteLLMOpenAIResponseObject) -> BaseLiteLLMOpenAIResponseObject: + stamped: Final = event.model_copy(update={"sequence_number": self._sequence_number}) + self._sequence_number += 1 + return stamped + async def __anext__( self, + ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: + return self._stamp_sequence_number(await self._anext_unstamped_event()) + + async def _anext_unstamped_event( + self, ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: try: while True: @@ -1363,6 +1322,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def __next__( self, + ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: + return self._stamp_sequence_number(self._next_unstamped_event()) + + def _next_unstamped_event( + self, ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: try: while True: @@ -1504,16 +1468,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if not self.sent_content_part_added_event: self._queue_message_item_added_events() self._hold_tool_calls_behind_text(getattr(chunk.choices[0].delta, "tool_calls", None)) - self._sequence_number += 1 - text_delta_event: Final = OutputTextDeltaEvent( + return OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id=item_id, output_index=self._message_index(), content_index=0, delta=delta_content, ) - text_delta_event.__dict__["sequence_number"] = self._sequence_number - return text_delta_event # Priority 3: Handle tool call deltas (if any) -> queue events and emit them # For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 0f4ed0611be..3e605e926f7 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -52,6 +52,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( cache_control_count, dropped_cache_control, newest_marks_within_cache_control_limit, + routed_model_info, whole_messages_within_cache_control_limit, with_cache_control, ) @@ -782,13 +783,15 @@ def _replays_compaction_block(input: str | ResponseInputParam, context_managemen _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai")) -def _bedrock_target_keeps_item_marks(model: str, drop_params: bool) -> bool: - declared: Final = bedrock_declared_cache_point_support(model) +def _bedrock_target_keeps_item_marks(model: str, drop_params: bool, model_info: Mapping[str, object] | None) -> bool: + declared: Final = bedrock_declared_cache_point_support(model, model_info) return declared is True or (declared is None and not drop_params) -def _target_keeps_item_marks(model: str, custom_llm_provider: str | None, drop_params: bool) -> bool: - if supports_anthropic_cache_control(model, custom_llm_provider): +def _target_keeps_item_marks( + model: str, custom_llm_provider: str | None, drop_params: bool, model_info: Mapping[str, object] | None +) -> bool: + if supports_anthropic_cache_control(model, custom_llm_provider, model_info): return True if custom_llm_provider in _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: return True @@ -798,7 +801,7 @@ def _target_keeps_item_marks(model: str, custom_llm_provider: str | None, drop_p return is_snowflake_claude_model(model) if custom_llm_provider == "databricks": return "claude" in model.lower() - return custom_llm_provider == "bedrock" and _bedrock_target_keeps_item_marks(model, drop_params) + return custom_llm_provider == "bedrock" and _bedrock_target_keeps_item_marks(model, drop_params, model_info) def _refuses_marks_on_tool_blocks(model: str, custom_llm_provider: str | None) -> bool: @@ -1234,7 +1237,10 @@ class LiteLLMCompletionResponsesConfig: request_kwargs: Final = _STR_KEY_DICT_ADAPTER.validate_python({**kwargs, "drop_params": drop_params}) context_management: Final = responses_api_request.get("context_management") keeps_item_marks: Final = _target_keeps_item_marks( - model, custom_llm_provider, drops_unsupported_params(request_kwargs.get("drop_params")) + model, + custom_llm_provider, + drops_unsupported_params(request_kwargs.get("drop_params")), + routed_model_info(request_kwargs.get("model_info")), ) request_input: Final = with_compaction_items_as_params(input) native_compaction: Final = _replays_compaction_block( @@ -1456,7 +1462,10 @@ class LiteLLMCompletionResponsesConfig: model: Final = str(request_fields.get("model") or "") provider: Final = requested_provider if isinstance(requested_provider, str) else None keeps_item_marks: Final = _target_keeps_item_marks( - model, provider, drops_unsupported_params(request_kwargs.get("drop_params")) + model, + provider, + drops_unsupported_params(request_kwargs.get("drop_params")), + routed_model_info(request_kwargs.get("model_info")), ) lookup: Final = session_lookup or partial( ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index ad1d35bbe17..3de63538f9a 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -27,6 +27,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( CARRY_UNMATCHED_MESSAGE_POINTS, AnthropicCacheControlHook, configured_injection_points, + routed_model_info, ) from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind @@ -882,6 +883,7 @@ def _stamp_injection_points_for_the_hook( custom_llm_provider, kwargs.get("api_base") or kwargs.get("base_url"), kwargs.get("prompt_cache_options"), + routed_model_info(kwargs.get("model_info")), ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 09ee6828157..da8ff2cd207 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -379,6 +379,7 @@ def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: Final = frozenset({"response.created", "response.in_progress", "response.queued"}) +SEQUENCE_NUMBER_STAMPED_ON_EMISSION: Final = -1 class BaseResponsesAPIStreamingIterator: @@ -1933,7 +1934,6 @@ def build_synthetic_response_events( _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), ] - sequence_number = 0 output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or [] for output_index, output_item in enumerate(output_items): output_item_payload = _dump_response_object(output_item) @@ -2043,35 +2043,31 @@ def build_synthetic_response_events( delta=summary_text[i : i + chunk_size], ) ) - sequence_number += 1 events.append( openai_types.ReasoningSummaryTextDoneEvent( type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, item_id=item_id, output_index=output_index, - sequence_number=sequence_number, + sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION, summary_index=summary_index, text=summary_text, ) ) - sequence_number += 1 events.append( openai_types.ReasoningSummaryPartDoneEvent( type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, item_id=item_id, output_index=output_index, - sequence_number=sequence_number, + sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION, summary_index=summary_index, part=openai_types.BaseLiteLLMOpenAIResponseObject(**summary_payload), ) ) - sequence_number += 1 events.append( openai_types.OutputItemDoneEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, - sequence_number=sequence_number, item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload), ) ) @@ -2082,7 +2078,7 @@ def build_synthetic_response_events( response=transformed, ) ) - return events + return [event.model_copy(update={"sequence_number": number}) for number, event in enumerate(events)] # --------------------------------------------------------------------------- diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 1a97c1cb48c..43ed9a3b96f 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -678,20 +678,25 @@ class ResponsesAPIRequestUtils: return items: Final = cast(list[object], request_input) # cast-ok: untyped client json stripped: Final = tuple( - ResponsesAPIRequestUtils._without_encrypted_reasoning(item) - if should_strip is None or (isinstance(item, Mapping) and should_strip(cast(Mapping[str, object], item))) + ResponsesAPIRequestUtils._without_encrypted_reasoning(reasoning) + if (reasoning := ResponsesAPIRequestUtils._encrypted_reasoning_item(item)) is not None + and (should_strip is None or should_strip(reasoning)) else item for item in items ) items[:] = (item for item in stripped if item is not None) @staticmethod - def _without_encrypted_reasoning(item: object) -> object | None: + def _encrypted_reasoning_item(item: object) -> Mapping[str, object] | None: if not isinstance(item, dict): - return item + return None reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"): - return reasoning + return None + return reasoning + + @staticmethod + def _without_encrypted_reasoning(reasoning: Mapping[str, object]) -> object | None: readable: Final = any( ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content") ) diff --git a/litellm/router_api.py b/litellm/router_api.py index c013a99573f..3007b5345c7 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -155,6 +155,7 @@ GroupField: TypeAlias = Literal[ "display_name", "mode", "service_tiers", + "max_input_tokens", "token_limits", "listing_info", ] @@ -855,6 +856,16 @@ class RouterRegistry: team_id: str | None = None, ) -> tuple[object, ...]: ... + @overload + def group( + self, + name: str, + *, + field: Literal["max_input_tokens"], + include_hidden: bool = False, + team_id: str | None = None, + ) -> tuple[int | None, ...]: ... + @overload def group( self, @@ -882,9 +893,7 @@ class RouterRegistry: field: GroupField = "info", include_hidden: bool = False, team_id: str | None = None, - ) -> ( - "ModelGroupInfo | DeploymentModelListingInfo | tuple[int | None, int | None] | tuple[object, ...] | str | None" - ): + ) -> "ModelGroupInfo | DeploymentModelListingInfo | tuple[int | None, ...] | tuple[object, ...] | str | None": """ What the router knows about a model group. ``field``: @@ -894,6 +903,7 @@ class RouterRegistry: (were ``get_configured_display_name``, ``get_configured_mode``, ``get_configured_token_limits``, ``get_model_listing_info``) - ``service_tiers``: (was ``get_configured_service_tiers``) + - ``max_input_tokens``: each routable deployment's own ``model_info.max_input_tokens`` """ r: Final = self._router if field == "info": @@ -906,6 +916,8 @@ class RouterRegistry: return r._core_get_configured_mode(name) if field == "service_tiers": return r._core_get_configured_service_tiers(name, team_id=team_id) + if field == "max_input_tokens": + return r._core_get_configured_max_input_tokens(name, team_id=team_id) if field == "token_limits": return r._core_get_configured_token_limits(name) if field == "listing_info": diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 30b05456987..7ccb23ebfa8 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -84,6 +84,8 @@ from litellm.constants import ( MID_STREAM_FAILURE_LOG_WAIT_SECONDS, OUTPUT_TOKEN_CEILING_PARAMS, PROMPT_CACHE_PIN_CACHE_MAX_SIZE, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, RUNTIME_UPDATABLE_ROUTER_SETTINGS, @@ -141,6 +143,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import ( is_registered_custom_provider, resolve_model_and_provider_for_metadata, ) +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, ) @@ -475,6 +478,7 @@ from litellm.types.router import ( from litellm.types.services import ServiceTypes from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + COST_MAP_LOOKUP_KEY, PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CustomPricingLiteLLMParams, GenericBudgetConfigType, @@ -486,6 +490,7 @@ from litellm.types.utils import ( StandardLoggingRoutingDecision, Usage, all_litellm_params, + echoed_cost_map_fields, is_server_derived_pricing_key, shared_backend_model_info, ) @@ -544,6 +549,7 @@ if TYPE_CHECKING: ResponseAPIUsage, ResponseIncompleteEvent, ResponsesAPIResponse, + ResponsesAPIStreamingResponse, ) Span = _Span @@ -668,6 +674,7 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType _NO_DEFAULT_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_CALL_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) _RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None) @@ -771,6 +778,17 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None: + model_info: Final = _configured_model_info(deployment) + catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) + catalog_entry: Final = ( + GetModelCostMap.loaded_model_cost_map().get(catalog_key) if isinstance(catalog_key, str) else None + ) + if catalog_entry is not None and "max_input_tokens" in echoed_cost_map_fields(model_info, catalog_entry): + return None + return coerce_token_limit(model_info.get("max_input_tokens")) + + def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request @@ -832,9 +850,11 @@ def _with_router_resolved_session_model(session: object, model_name: str) -> Map typed_session: Final = _SESSION_ADAPTER.validate_python(session) except ValidationError: return _NO_SESSION_KWARGS - if "model" not in typed_session: + if "model" not in typed_session and typed_session.get("type") != "transcription": return _NO_SESSION_KWARGS - return MappingProxyType({"session": {**typed_session, "model": model_name}}) + from litellm.realtime_api.main import with_resolved_session_model + + return MappingProxyType({"session": with_resolved_session_model(dict(typed_session), model_name)}) # Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks @@ -2850,6 +2870,14 @@ class PythonCore(DeprecatedRouterMembers): self._host.invoke("acreate_realtime_transcription_session"), call_type="acreate_realtime_transcription_session", ) + self.acreate_realtime_translation_client_secret = self.factory_function( + self._host.invoke("acreate_realtime_translation_client_secret"), + call_type="acreate_realtime_translation_client_secret", + ) + self.arealtime_translation_calls = self.factory_function( + self._host.invoke("arealtime_translation_calls"), + call_type="arealtime_translation_calls", + ) self._aresponses_websocket = self.factory_function( self._host.invoke("_aresponses_websocket"), call_type="_aresponses_websocket" ) @@ -4385,12 +4413,13 @@ class PythonCore(DeprecatedRouterMembers): fallback_headers_adopted: bool = False - def __init__(self, async_generator: AsyncGenerator): + def __init__(self, async_generator: "AsyncGenerator[ResponsesAPIStreamingResponse, None]"): import time from datetime import datetime self._async_generator = async_generator self._billing_source: object = source_iterator + self._last_sequence_number = -1 # Mirror every attribute BaseResponsesAPIStreamingIterator.__init__ # would have set. The wrapper bypasses super().__init__ (it has no # httpx.Response of its own and no provider config to drive), so @@ -4497,7 +4526,16 @@ class PythonCore(DeprecatedRouterMembers): if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: self.completed_response = chunk self._yielded_first_chunk = True - return chunk + sequence_number: Final = getattr(chunk, "sequence_number", None) + if not isinstance(sequence_number, int): + return chunk + if not self.fallback_headers_adopted: + self._last_sequence_number = max(self._last_sequence_number, sequence_number) + return chunk + self._last_sequence_number = max(self._last_sequence_number + 1, sequence_number) + if self._last_sequence_number == sequence_number: + return chunk + return chunk.model_copy(update={"sequence_number": self._last_sequence_number}) async def aclose(self): # async generators always expose aclose — no defensive check needed. @@ -4559,7 +4597,9 @@ class PythonCore(DeprecatedRouterMembers): exc, ) - wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) + wrapper: Final = FallbackResponsesStreamWrapper( + stream_with_fallbacks() # pyright: ignore[reportArgumentType] # relays response events typed as object + ) return wrapper def _aresponses_background_stream_can_release( @@ -6865,6 +6905,22 @@ class PythonCore(DeprecatedRouterMembers): raise safeguard_refusal_error(model=model, stop_details=refusal_details) self.success_calls[model_name] += 1 + realtime_session_logging: Final = ( + _CALL_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") + if generic_function_name == "_arealtime" + else None + ) + realtime_session_deployment_id: Final = ( + _MODEL_INFO_ADAPTER.validate_python(kwargs["model_info"]).get("id") + if isinstance(realtime_session_logging, LiteLLMLogging) + and realtime_session_logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + and realtime_session_logging.model_call_details.get(REALTIME_SESSION_ENDED_ABNORMALLY_KEY) is not True + else None + ) + if realtime_session_deployment_id is not None: + increment_deployment_successes_for_current_minute( + ctx=self._core_ctx, deployment_id=str(realtime_session_deployment_id) + ) verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) return response @@ -8406,6 +8462,8 @@ class PythonCore(DeprecatedRouterMembers): "acreate_realtime_client_secret", "arealtime_calls", "acreate_realtime_transcription_session", + "acreate_realtime_translation_client_secret", + "arealtime_translation_calls", "_aresponses_websocket", "acreate_fine_tuning_job", "acancel_fine_tuning_job", @@ -8676,6 +8734,8 @@ class PythonCore(DeprecatedRouterMembers): "acreate_realtime_client_secret", "arealtime_calls", "acreate_realtime_transcription_session", + "acreate_realtime_translation_client_secret", + "arealtime_translation_calls", ): return await self._ageneric_api_call_with_fallbacks( original_function=original_function, @@ -9125,11 +9185,17 @@ class PythonCore(DeprecatedRouterMembers): if include_fallback_errors: input_kwargs["include_fallback_errors"] = True - # ORDER-BASED FALLBACKS: prepend higher order levels to the fallback list - # Skip for error types that have their own dedicated fallback handlers - _skip_order_fallback: Final = isinstance( - e, - (litellm_exceptions.ContextWindowExceededError, litellm_exceptions.ContentPolicyViolationError), + context_window_fallback_model_group: Final = ( + self._get_fallback_model_group_for_lookup_groups( + fallbacks=context_window_fallbacks, + lookup_groups=lookup_groups, + ) + if isinstance(e, litellm_exceptions.ContextWindowExceededError) and context_window_fallbacks is not None + else None + ) + _skip_order_fallback: Final = ( + isinstance(e, litellm_exceptions.ContentPolicyViolationError) + or context_window_fallback_model_group is not None ) _request_team_id: Final[str | None] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id") # Use wildcard-aware lookup so order-based fallback also works for model @@ -9175,14 +9241,23 @@ class PythonCore(DeprecatedRouterMembers): "original_model_group": original_model_group, } ) - response = await run_async_fallback( - *args, - **input_kwargs, - ) - return response + try: + return await run_async_fallback( + *args, + **input_kwargs, + ) + except Exception: + if isinstance(original_exception, litellm_exceptions.ContextWindowExceededError): + raise original_exception + raise # Weighted intra-group failover (simple-shuffle only); see _maybe_run_weighted_failover. - if self.enable_weighted_failover and not _skip_order_fallback and original_model_group is not None: + if ( + self.enable_weighted_failover + and not _skip_order_fallback + and not isinstance(e, litellm_exceptions.ContextWindowExceededError) + and original_model_group is not None + ): response = await self._maybe_run_weighted_failover( exception=e, original_model_group=original_model_group, @@ -9200,7 +9275,7 @@ class PythonCore(DeprecatedRouterMembers): # check if client-side fallbacks are used (e.g. fallbacks = ["gpt-3.5-turbo", "claude-3-haiku"] or fallbacks=[{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hey, how's it going?"}]}] is_non_standard_fallback_format: Final = check_non_standard_fallback_format(fallbacks=fallbacks) - if is_non_standard_fallback_format: + if is_non_standard_fallback_format and context_window_fallback_model_group is None: input_kwargs.update( { "fallback_model_group": fallbacks, @@ -9216,16 +9291,7 @@ class PythonCore(DeprecatedRouterMembers): return response if isinstance(e, litellm_exceptions.ContextWindowExceededError): - if context_window_fallbacks is not None: - context_window_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_for_lookup_groups( - fallbacks=context_window_fallbacks, - lookup_groups=lookup_groups, - ) - ) - if context_window_fallback_model_group is None: - raise original_exception - + if context_window_fallback_model_group is not None: input_kwargs.update( { "fallback_model_group": context_window_fallback_model_group, @@ -9239,7 +9305,7 @@ class PythonCore(DeprecatedRouterMembers): ) return response - else: + elif context_window_fallbacks is None: error_message = f"model={model_group}. context_window_fallbacks={mask_sensitive_structure(context_window_fallbacks)}. fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" verbose_router_logger.info( msg=f"Got 'ContextWindowExceededError'. No context_window_fallback set. Defaulting \ @@ -12680,6 +12746,24 @@ class PythonCore(DeprecatedRouterMembers): for deployment in self._routable_deployments(model_name, team_id) ) + def _core_get_configured_max_input_tokens( + self, model_name: str, team_id: str | None = None + ) -> tuple[int | None, ...]: + """ + Return the max_input_tokens each deployment a request from team_id to + model_name can route to (see _routable_deployments) sets in its own + model_info, in model_list order, None for a deployment that sets none + or a malformed one. Unlike get_model_listing_info, neither the cost map + nor a discovered limit is read, so only an operator's own value counts, + and a value a saved /model/info response copied from the cost map is + skipped too. + + Returns an empty tuple for wildcard-expanded or unknown names. + """ + return tuple( + _operator_max_input_tokens(deployment) for deployment in self._routable_deployments(model_name, team_id) + ) + def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: """ Whether each deployment a request from team_id to model_name can route @@ -13202,6 +13286,7 @@ class PythonCore(DeprecatedRouterMembers): "providers": [llm_provider], "description": model_group_description, **model_info, + "supports_continuation": model_info.get("supports_continuation") is True, "supports_fast_mode": True, "supported_reasoning_efforts": None, "supported_endpoints": None, @@ -13269,6 +13354,8 @@ class PythonCore(DeprecatedRouterMembers): and model_info["supports_url_context"] is True ): model_group_info.supports_url_context = True + if model_info.get("supports_continuation") is True: + model_group_info.supports_continuation = True if model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True: model_group_info.supports_reasoning = True diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 4f4ed76c170..05dd4706b62 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1219,9 +1219,11 @@ def _with_classifier_forecast( forecasted_decision: Final[StandardLoggingRoutingDecision] = { **with_input, "classifier_probabilities": outcome.jev_verdict.probabilities, - "classifier_confidence": outcome.jev_verdict.confidence, } - return forecasted_decision + confidence: Final = outcome.jev_verdict.confidence + return ( + forecasted_decision if confidence is None else {**forecasted_decision, "classifier_confidence": confidence} + ) if outcome.llm_v2_forecast is not None: return _with_llm_v2_forecast(with_input, outcome.llm_v2_forecast) forecast: Final = outcome.capability_forecast @@ -2468,7 +2470,7 @@ class ComplexityRouter(CustomLogger): score=None, signals=( f"{config.provider}-classifier:{tier_name}", - f"{config.provider}-confidence={answer.confidence:.6f}", + *(() if answer.confidence is None else (f"{config.provider}-confidence={answer.confidence:.6f}",)), *( f"tier-probability:{label}={probability:.6f}" for label, probability in answer.probabilities.items() diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index bb7d2e91cd3..5cd4f00a941 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -210,7 +210,7 @@ class JevChoiceAnswer(LiteLLMBaseModel): type: Literal["choice"] choice: str probabilities: Mapping[str, JevProbability] - confidence: JevProbability + confidence: JevProbability | None = None class JevUsage(LiteLLMBaseModel): @@ -425,7 +425,7 @@ class HttpJevClassifierClient: class JevVerdict(NamedTuple): label: str probabilities: Mapping[str, float] - confidence: float + confidence: float | None model: str cost: float | None provider: str = "typesafe" diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 6f3f0794371..f3425596afa 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -36,31 +36,75 @@ Safe to enable globally: - No cache required. """ -from collections.abc import Iterator, Mapping, Sequence -from functools import cache +import hashlib +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, field +from functools import cache, lru_cache, partial +from time import monotonic from typing import TYPE_CHECKING, Final, Optional, cast -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError +import litellm from litellm._logging import verbose_router_logger +from litellm.exceptions import AuthenticationError, BadRequestError from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider, get_llm_provider +from litellm.litellm_core_utils.litellm_logging import has_sensitive_keyword from litellm.litellm_core_utils.prompt_templates.common_utils import ( anthropic_content_lists, encrypted_content_of_block, strip_encrypted_reasoning_from_messages, ) from litellm.litellm_core_utils.routing_common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param +from litellm.secret_managers.main import secret_manager_would_be_consulted from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import Deployment -from litellm.utils import get_order_filtered_deployments +from litellm.types.router import Deployment, GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager, get_order_filtered_deployments, with_named_credential if TYPE_CHECKING: from litellm.router_api import RoutingContext _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[object, object]) +_ACCOUNT_SCOPE_HEADERS: Final = frozenset({"openai-organization", "openai-project", "cookie"}) +_SINGLE_VENDOR_PROVIDERS: Final = frozenset({LlmProviders.OPENAI.value}) +_DEFAULT_ROUTE_SECRETS: Final = ("USE_LITELLM_PROXY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_API_BASE") +_SECRET_MANAGER_REFRESH_SECONDS: Final = 60 + + +@dataclass(frozen=True, slots=True) +class _ProviderRoute: + model: str + provider: str + api_base: str | None + + +@dataclass(frozen=True, slots=True) +class _EncryptionBoundary: + api_base: object + api_key: object + model: Callable[[], str | None] = field(compare=False) + account_scope: frozenset[tuple[str, object]] + explicit_pair: bool + + def is_shared_with(self, other: "_EncryptionBoundary") -> bool: + if (self.api_base, self.api_key) != (other.api_base, other.api_key): + return False + if self.explicit_pair and other.explicit_pair: + return True + return self.account_scope == other.account_scope and self.model() == other.model() + + +@dataclass(frozen=True, slots=True) +class _ParamsInRefreshWindow: + digest: str + window: int + params: Mapping[str, object] = field(compare=False) class EncryptedContentAffinityCheck(CustomLogger): @@ -83,9 +127,11 @@ class EncryptedContentAffinityCheck(CustomLogger): ctx: Optional["RoutingContext"] = None, enable_global_affinity: bool = True, model_group_affinity_config: dict[str, list[str]] | None = None, + clock: Callable[[], float] = monotonic, ) -> None: super().__init__() self.ctx = ctx + self._clock = clock self.enable_global_affinity = enable_global_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} @@ -209,44 +255,156 @@ class EncryptedContentAffinityCheck(CustomLogger): @staticmethod def _encryption_boundary_key( litellm_params: object, - ) -> tuple[object, object] | None: + clock: Callable[[], float] = monotonic, + ) -> _EncryptionBoundary | None: """ - ``(api_base, api_key)`` identifies an upstream encryption boundary. - The values are resolved from the deployment and its named credential - without modifying the deployment. + The upstream encryption boundary a deployment's encrypted reasoning belongs to. The values + are resolved the way the request resolves them, the deployment's own values winning over its + named credential's, without modifying the deployment. Two explicit ``(api_base, api_key)`` + pairs match on those values alone. On a provider that serves one vendor's models, a deployment + with no ``api_key`` takes the provider's default key, one with no ``api_base`` takes the endpoint + its Responses URL builder resolves, and a match involving either also needs the same model and + the same headers that pick an account. - Accepts any object exposing dict-style ``.get(key, default)``: plain - dicts (the common case in ``healthy_deployments``) as well as - ``LiteLLM_Params``-style Pydantic instances, which define a custom - ``.get()``. A stricter ``isinstance(dict)`` guard would silently drop - the latter from boundary matching and fall back to the full pool — - i.e. trigger the exact ``invalid_encrypted_content`` failure this - check exists to prevent. + Accepts plain dicts (the common case in ``healthy_deployments``) as + well as ``LiteLLM_Params`` Pydantic instances. A stricter + ``isinstance(dict)`` guard would silently drop the latter from + boundary matching and fall back to the full pool — i.e. trigger the + exact ``invalid_encrypted_content`` failure this check exists to prevent. + + When a secret manager would be asked for the provider defaults, the boundary of one set of + deployment and named credential values is reused for up to a minute, so a run of pin misses + does not ask it again for every candidate, while a changed credential resolves at once. """ - getter: Final = getattr(litellm_params, "get", None) - if not callable(getter): + params: Final = EncryptedContentAffinityCheck._params_mapping(litellm_params) + if params is None: return None - api_base: Final = getter("api_base") - api_key: Final = getter("api_key") - credential_name: Final = getter("litellm_credential_name") - credential_values: Final[Mapping[str, object] | None] = ( - CredentialAccessor.get_credential_values(credential_name) - if isinstance(credential_name, str) and credential_name - else None + effective_params: Final = with_named_credential(params) + if not any(secret_manager_would_be_consulted(name) for name in _DEFAULT_ROUTE_SECRETS): + return EncryptedContentAffinityCheck._resolve_encryption_boundary(effective_params) + return EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window( + _ParamsInRefreshWindow( + hashlib.sha256(repr(effective_params).encode()).hexdigest(), + int(clock() // _SECRET_MANAGER_REFRESH_SECONDS), + effective_params, + ) ) - effective_api_base: Final = ( - credential_values.get("api_base") - if credential_values is not None and "api_base" in credential_values - else api_base - ) - effective_api_key: Final = ( - credential_values.get("api_key") - if credential_values is not None and "api_key" in credential_values - else api_key - ) - if not effective_api_base or not effective_api_key: + + @staticmethod + @lru_cache(maxsize=1024) + def _encryption_boundary_in_refresh_window(key: _ParamsInRefreshWindow) -> _EncryptionBoundary | None: + return EncryptedContentAffinityCheck._resolve_encryption_boundary(key.params) + + @staticmethod + def _resolve_encryption_boundary(effective_params: Mapping[str, object]) -> _EncryptionBoundary | None: + explicit_api_key: Final = effective_params.get("api_key") + explicit_api_base: Final = effective_params.get("api_base") + account_scope: Final = EncryptedContentAffinityCheck._account_scope(effective_params) + if explicit_api_key and explicit_api_base: + return _EncryptionBoundary( + explicit_api_base, + explicit_api_key, + cache(partial(EncryptedContentAffinityCheck._single_vendor_model, effective_params)), + account_scope, + explicit_pair=True, + ) + route: Final = EncryptedContentAffinityCheck._single_vendor_route(effective_params) + if route is None: + return None + api_key: Final = explicit_api_key or EncryptedContentAffinityCheck._default_api_key(route) + if not api_key: + return None + api_base: Final = explicit_api_base or EncryptedContentAffinityCheck._default_responses_endpoint( + route, effective_params + ) + if not api_base: + return None + return _EncryptionBoundary(api_base, api_key, lambda: route.model, account_scope, explicit_pair=False) + + @staticmethod + def _params_mapping(litellm_params: object) -> Mapping[str, object] | None: + if isinstance(litellm_params, BaseModel): + return litellm_params.model_dump(exclude_none=True) + if isinstance(litellm_params, Mapping): + return _REQUEST_KWARGS_ADAPTER.validate_python(litellm_params) + return None + + @staticmethod + def _account_scope(litellm_params: Mapping[str, object]) -> frozenset[tuple[str, object]]: + return frozenset( + (name.lower(), str(value)) + for name, value in EncryptedContentAffinityCheck._header_items(litellm_params) + if name.lower() in _ACCOUNT_SCOPE_HEADERS or has_sensitive_keyword(name) + ) + + @staticmethod + def _header_items(litellm_params: Mapping[str, object]) -> Iterator[tuple[str, object]]: + headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("headers")) + extra_headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("extra_headers")) + overridden: Final = frozenset(name.lower() for name in extra_headers) + yield from ((name, value) for name, value in headers.items() if name.lower() not in overridden) + yield from extra_headers.items() + + @staticmethod + def _header_mapping(headers: object) -> Mapping[str, object]: + if isinstance(headers, Mapping): + return {str(name): value for name, value in _HEADERS_ADAPTER.validate_python(headers).items()} + return {} + + @staticmethod + def _single_vendor_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None: + route: Final = EncryptedContentAffinityCheck._provider_route(litellm_params) + return route if route is not None and route.provider in _SINGLE_VENDOR_PROVIDERS else None + + @staticmethod + def _single_vendor_model(litellm_params: Mapping[str, object]) -> str | None: + route: Final = EncryptedContentAffinityCheck._single_vendor_route(litellm_params) + return route.model if route is not None else None + + @staticmethod + def _provider_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None: + model: Final = litellm_params.get("model") + if not isinstance(model, str): + return None + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") + declared: Final = declared_provider( + model, custom_llm_provider if isinstance(custom_llm_provider, str) else None + ) + if declared in litellm.provider_list and declared not in _SINGLE_VENDOR_PROVIDERS: + return None + try: + provider_model, provider, _, provider_api_base = get_llm_provider( + model=model, litellm_params=GenericLiteLLMParams.model_validate(dict(litellm_params)) + ) + except (BadRequestError, ValidationError): + return None + return _ProviderRoute(provider_model, provider, provider_api_base) + + @staticmethod + def _default_api_key(route: _ProviderRoute) -> str | None: + model_info: Final = ProviderConfigManager.get_provider_model_info( + model=route.model, provider=LlmProviders(route.provider) + ) + if model_info is None: + return None + try: + return model_info.get_api_key() + except NotImplementedError: + return None + + @staticmethod + def _default_responses_endpoint(route: _ProviderRoute, litellm_params: Mapping[str, object]) -> str | None: + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config( + provider=route.provider, model=route.model + ) + if responses_config is None: + return None + try: + return responses_config.get_complete_url( # pyright: ignore[reportUnknownMemberType] # the base types litellm_params as a bare dict + api_base=route.api_base, litellm_params=dict(litellm_params) + ) + except (ValueError, BaseLLMException, AuthenticationError): return None - return (effective_api_base, effective_api_key) def _find_deployments_on_same_encryption_boundary( self, @@ -255,7 +413,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ) -> tuple[list[dict], Deployment | None]: """ Deployments in ``healthy_deployments`` sharing the originating - deployment's ``(api_base, api_key)``, alongside the originating + deployment's encryption boundary, alongside the originating deployment object (or ``None`` if it was removed / router unavailable). """ if self.ctx is None: @@ -263,11 +421,16 @@ class EncryptedContentAffinityCheck(CustomLogger): originating: Final = self.ctx.registry().deployment(model_id) if originating is None: return [], None - boundary: Final = self._encryption_boundary_key(originating.litellm_params.model_dump(exclude_none=True)) + boundary: Final = self._encryption_boundary_key( + originating.litellm_params.model_dump(exclude_none=True), self._clock + ) if boundary is None: return [], originating matches: Final = [ - d for d in healthy_deployments if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary + d + for d in healthy_deployments + if (candidate := self._encryption_boundary_key(d.get("litellm_params", {}), self._clock)) is not None + and boundary.is_shared_with(candidate) ] return matches, originating @@ -280,6 +443,7 @@ class EncryptedContentAffinityCheck(CustomLogger): target_deployments: Sequence[Mapping[str, object]], *, unmarked_origin: str | None, + clock: Callable[[], float] = monotonic, ) -> None: """ Drop the encrypted reasoning that none of ``target_deployments`` minted or shares an @@ -293,12 +457,19 @@ class EncryptedContentAffinityCheck(CustomLogger): for target in target_deployments if isinstance((model_info := target.get("model_info")), Mapping) and model_info.get("id") is not None ) - target_boundaries: Final = frozenset( - boundary - for target in target_deployments - if (boundary := EncryptedContentAffinityCheck._encryption_boundary_key(target.get("litellm_params"))) - is not None - ) + + @cache + def target_boundaries() -> tuple[_EncryptionBoundary, ...]: + return tuple( + boundary + for target in target_deployments + if ( + boundary := EncryptedContentAffinityCheck._encryption_boundary_key( + target.get("litellm_params"), clock + ) + ) + is not None + ) @cache def target_can_decrypt(marked_origin: str | None) -> bool: @@ -312,12 +483,14 @@ class EncryptedContentAffinityCheck(CustomLogger): origin: Final = ctx.registry().deployment(origin_model_id) origin_boundary: Final = ( EncryptedContentAffinityCheck._encryption_boundary_key( - origin.litellm_params.model_dump(exclude_none=True) + origin.litellm_params.model_dump(exclude_none=True), clock ) if origin is not None else None ) - return origin_boundary is not None and origin_boundary in target_boundaries + return origin_boundary is not None and any( + origin_boundary.is_shared_with(target_boundary) for target_boundary in target_boundaries() + ) def should_strip_input_item(item: Mapping[str, object]) -> bool: return not target_can_decrypt(EncryptedContentAffinityCheck._model_id_of_input_item(item)) @@ -424,7 +597,12 @@ class EncryptedContentAffinityCheck(CustomLogger): ) request_kwargs["_encrypted_content_affinity_pinned"] = True self.strip_reasoning_the_targets_cannot_decrypt( - self.ctx, request_input, anthropic_messages, (deployment,), unmarked_origin=model_id + self.ctx, + request_input, + anthropic_messages, + (deployment,), + unmarked_origin=model_id, + clock=self._clock, ) return [deployment] @@ -442,7 +620,12 @@ class EncryptedContentAffinityCheck(CustomLogger): ) request_kwargs["_encrypted_content_affinity_pinned"] = True self.strip_reasoning_the_targets_cannot_decrypt( - self.ctx, request_input, anthropic_messages, boundary_matches, unmarked_origin=model_id + self.ctx, + request_input, + anthropic_messages, + boundary_matches, + unmarked_origin=model_id, + clock=self._clock, ) return boundary_matches diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index f78b5982186..1d91255e28c 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -55,6 +55,17 @@ def _sorted_pairs(pairs: Iterable[tuple[str, JsonValue]]) -> tuple[tuple[str, Js return tuple(sorted(pairs, key=lambda pair: pair[0])) +_CACHE_MARKER_KEYS: Final = frozenset(("cache_control", "prompt_cache_breakpoint")) + + +def _is_cache_marked(block: Mapping[str, object]) -> bool: + match block: + case {"cache_control": {"type": "ephemeral"}}: + return True + case _: + return block.get("prompt_cache_breakpoint") is not None + + def _canonical_bytes(value: object) -> bytes: return json.dumps(value, sort_keys=True, separators=(",", ":")).encode() @@ -66,7 +77,7 @@ def _block_unit( return _canonical_bytes((envelope, block)), message_run_type block_type: Final = block.get("type") block_run_type: Final = block_type if isinstance(block_type, str) and block_type in _TOOL_RUN_BLOCK_TYPES else None - stripped: Final = _sorted_pairs(item for item in block.items() if item[0] != "cache_control") + stripped: Final = _sorted_pairs(item for item in block.items() if item[0] not in _CACHE_MARKER_KEYS) return _canonical_bytes((envelope, stripped)), message_run_type or block_run_type @@ -264,15 +275,9 @@ class PromptCachingCache: continue for content_idx, content_block in enumerate(content): - if isinstance(content_block, dict): - cache_control = content_block.get("cache_control") - if ( - cache_control is not None - and isinstance(cache_control, dict) - and cache_control.get("type") == "ephemeral" - ): - last_cacheable_message_idx = msg_idx - last_cacheable_content_idx = content_idx + if isinstance(content_block, dict) and _is_cache_marked(content_block): + last_cacheable_message_idx = msg_idx + last_cacheable_content_idx = content_idx # If no cacheable block found, return empty list (no cacheable prefix) if last_cacheable_message_idx is None: diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 629d9098970..9ba8cbc35a7 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -66,6 +66,8 @@ class Feature(str, Enum): """Extra body fields that Python transforms and merges into the provider request.""" PROVIDER_AFFINITY = "provider_affinity" """A provider affinity header that Python validates and populates from the session id.""" + FOUNDRY_BETA_ROW = "foundry_beta_row" + """A Foundry api_base, where Python drops the anthropic-beta values Foundry rejects.""" PRE_REQUEST_HOOK = "pre_request_hook" """A logger overriding async_pre_request_hook, which may rewrite the request.""" WEBSEARCH_INTERCEPTION = "websearch_interception" @@ -170,6 +172,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.PROVIDER_AFFINITY, + Feature.FOUNDRY_BETA_ROW, Feature.SAFEGUARDS_CLASSIFIER, Feature.PRE_REQUEST_HOOK, Feature.WEBSEARCH_INTERCEPTION, diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index e155680bd30..92590531cdc 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -12,7 +12,7 @@ import datetime import traceback import uuid import weakref -from collections.abc import Awaitable, Coroutine, Mapping +from collections.abc import Awaitable, Coroutine, Mapping, Sequence from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -22,6 +22,7 @@ from typing import ( cast, # noqa: TID251 # bounded compatibility calls into legacy Python integrations ) +import httpx from pydantic import TypeAdapter if TYPE_CHECKING: @@ -243,6 +244,27 @@ def post_call( logger.post_call(original_response=original_response, api_key=api_key, additional_args=additional_args) +def record_response(logger: Logging, status: int, headers: Sequence[tuple[str, str]], body: str) -> None: + """Stores the provider response the way Python's Messages handler does, for logging to replay.""" + response: Final = httpx.Response(status_code=status, headers=list(headers), content=body.encode()) + details: Final = logger.model_call_details # the call's own logger: recording on it is the point + details["httpx_response"] = response + details["response_headers"] = dict(response.headers) + + +def record_api_window(logger: Logging, start: float, end: float) -> None: + """Records the provider exchange's window the way Python's timed `AsyncHTTPHandler.post` does.""" + from litellm.litellm_core_utils.logging_utils import ( + _set_duration_in_model_call_details, # pyright: ignore[reportPrivateUsage] # the timed post's own recorder + ) + + _set_duration_in_model_call_details( + logging_obj=logger, + start_time=datetime.datetime.fromtimestamp(start), # noqa: DTZ006 # the timed post records naive local times + end_time=datetime.datetime.fromtimestamp(end), # noqa: DTZ006 # the timed post records naive local times + ) + + def defers_async_logging(logger: LoggingSurface) -> bool: return bool(getattr(logger, "defer_async_logging", False)) @@ -359,6 +381,27 @@ def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_t _STREAM_HEADERS: Final = TypeAdapter(Mapping[str, str]) +class _HiddenParamsCarrier: + """Stands in for the stream `@client` stamps, so the head becomes its `_hidden_params`.""" + + def __init__(self, hidden_params: dict[str, object]) -> None: # mutable-ok: the stamps land in this dict + self._hidden_params: object = hidden_params + + +def stamp_stream_head( + head: dict[str, object], # mutable-ok: the stream's hidden params, stamped in place + logger: Logging, + kwargs: dict[str, object], # mutable-ok: `finalize` takes the call kwargs as a dict + start_time: datetime.datetime, + end_time: datetime.datetime, +) -> None: + carrier: Final = _HiddenParamsCarrier(head) + finalize(carrier, logger, kwargs, start_time, end_time) + stamped: Final = carrier._hidden_params # pyright: ignore[reportPrivateUsage] # the carrier's own field + if stamped is not head and isinstance(stamped, dict): + head.update(cast("dict[str, object]", stamped)) # cast-ok: update_response_metadata writes a str-keyed dict + + def stream_opened(logger: Logging, head: Mapping[str, object]) -> None: logger.stream = True logger.model_call_details["stream"] = True diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index 017ff65e2a7..c8b53097f46 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -1,21 +1,35 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +import datetime +from collections.abc import Callable, Mapping, Sequence from dataclasses import asdict, dataclass -from typing import Final, cast # noqa: TID251 # narrows the normalized native payload to the public TypedDict +from typing import Final, TypeAlias, cast # noqa: TID251 # narrows native payloads to public TypedDicts import httpx from pydantic import TypeAdapter, ValidationError import litellm +from litellm.litellm_core_utils import exception_mapping_utils from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.reasoning_effort_capability import deployment_reasoning_effort_levels from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout +from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.anthropic.messages.native import handler +from litellm.llms.anthropic.messages.native.handler import ( + _native_messages_config as native_messages_config, # pyright: ignore[reportPrivateUsage] # the handler's own config lookup +) +from litellm.llms.anthropic.messages.native.handler import ( + _with_deployment_output_effort as _deployment_output_effort, # pyright: ignore[reportPrivateUsage] # the handler's own effort rounding +) from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled from litellm.rust_bridge import failures +from litellm.rust_bridge.host import callbacks_legacy_python from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse _DROP_PATHS: Final = TypeAdapter(list[object]) +_METADATA: Final = TypeAdapter(dict[str, object]) _BODY_FIELDS: Final = frozenset( { "max_tokens", @@ -40,6 +54,7 @@ _BODY_FIELDS: Final = frozenset( "cache_control", "reasoning_effort", "safeguards", + "service_tier", } ) @@ -59,8 +74,40 @@ def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: ) +def inferred_provider(model: str, api_base: str | None) -> str | None: + """The provider Python's handler resolves for a call that names none, or None when the model + carries no provider `get_llm_provider` knows.""" + try: + _, provider, _, _ = litellm.get_llm_provider(model=model, api_base=api_base) + except litellm.BadRequestError: + return None + return provider + + def body_params(request: Mapping[str, object]) -> Mapping[str, object]: - return known_params(_BODY_FIELDS, request) + params: Final = known_params(_BODY_FIELDS, request) + if deployment_reasoning_effort_levels(request.get("model_info")) is None: + return params + return _with_deployment_effort(params, request) + + +def _with_deployment_effort(params: Mapping[str, object], request: Mapping[str, object]) -> Mapping[str, object]: + """Python's handler rounds the requested effort to the levels the deployment declares in + `model_info.reasoning_effort_levels`, the one place a deployment, not the model, decides.""" + explicit_provider: Final = optional_str(request.get("custom_llm_provider")) + try: + model, provider, _, _ = litellm.get_llm_provider( + model=str(request["model"]), custom_llm_provider=explicit_provider + ) + except litellm.BadRequestError: + return params + return _deployment_output_effort( + params, + model=model, + custom_llm_provider=provider, + kwargs=request, + drop_params=litellm.drop_params is True or request.get("drop_params") is True, + ) def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]: @@ -71,18 +118,147 @@ def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, obj return anthropic_messages_stream_hidden_params(httpx.Headers(list(headers))) -def map_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception: - if getattr(error, "messages_request_error", False): - return litellm.BadRequestError( - message=str(error), - model=str(request["model"]).removeprefix(f"{request_provider}/"), - llm_provider=request_provider, - ) - return failures.map_native_failure( - error, str(request["model"]), request_provider, request, optional_str(request.get("api_base")) +_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) +_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) +_ERROR_CLASS: TypeAlias = Callable[..., Exception] +_validate_metadata: Final = cast( # cast-ok: the handler's untyped metadata validator + Callable[[dict[str, object]], object], handler.validate_anthropic_api_metadata +) +_exception_type: Final = cast( # cast-ok: the legacy public exception mapper is untyped + Callable[..., Exception], exception_mapping_utils.exception_type +) +_BODY_READ_TIMEOUT: Final = "Timeout on reading data from socket" # aiohttp's SocketTimeoutError text +_MISSING_KEY: Final = ( + "Missing Anthropic API Key - A call is being made to anthropic but no key is set " + "either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` " + "or `ANTHROPIC_AUTH_TOKEN` in your environment vars" +) + + +def _rejection(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception: + """The error Python's handler raises for a request it refuses before sending it.""" + metadata: Final = request.get("metadata") + if isinstance(metadata, dict): + try: + _validate_metadata(_METADATA.validate_python(metadata)) + except ValidationError as invalid: + return invalid + if getattr(error, "messages_unsupported_param", False): + return litellm.UnsupportedParamsError(message=str(error)) + return litellm.BadRequestError( + message=str(error), + model=str(request["model"]).removeprefix(f"{request_provider}/"), + llm_provider=request_provider, ) +def _timeout_text(phase: str, request: Mapping[str, object]) -> str: + """`litellm.Timeout`'s text as Python's HTTP client raises it for the same timeout.""" + seconds: Final = request.get("timeout") + logging_obj: Final = request.get("litellm_logging_obj") + start_time: Final = getattr(logging_obj, "start_time", None) + time_taken: Final = ( + round( + (datetime.datetime.now() - start_time).total_seconds(), # noqa: DTZ005 # the logger's start_time is naive local time + 3, + ) + if isinstance(start_time, datetime.datetime) + else None + ) + return str( + litellm.Timeout( + message=f"{phase} timed out. Timeout passed={seconds}, time taken={time_taken} seconds", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ) + ) + + +def _upstream_failure( + error: Exception, request: Mapping[str, object] +) -> tuple[int, str, Sequence[tuple[str, str]]] | None: + """The status, text and headers Python's HTTP client reports for the same failure.""" + timeout_phase: Final = getattr(error, "messages_timeout", None) + if isinstance(timeout_phase, str): + return 408, _timeout_text(timeout_phase, request), () + try: + status, text = _UPSTREAM_ARGS.validate_python(error.args) + return status, text, _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) + except ValidationError: + return None + + +def _provider_error(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception | None: + """The provider exception Python's `_handle_error` builds for the same failure.""" + failure: Final = _upstream_failure(error, request) + if failure is None: + return None + status, text, headers = failure + api_base: Final = request.get("api_base") + config: Final = native_messages_config( + str(request["model"]).removeprefix(f"{request_provider}/"), + request_provider, + request.get("model_info"), + api_base if isinstance(api_base, str) else None, + ) + return ( + cast(_ERROR_CLASS, config.get_error_class)( # cast-ok: the config's untyped headers parameter + error_message=text, status_code=status, headers=dict(headers) + ) + if config is not None + else AnthropicError(status_code=status, message=text, headers=httpx.Headers(headers)) + ) + + +def _log_resolved_call(request: Mapping[str, object], request_provider: str) -> None: + """Tells the logger the resolved call, as Python's handler does before its provider config + transforms the request.""" + logger: Final = request.get("litellm_logging_obj") + if not isinstance(logger, LiteLLMLoggingObj): + return + callbacks_legacy_python.update_logging( + logger, + kwargs={name: value for name, value in request.items() if name != "metadata"}, + model=str(request["model"]).removeprefix(f"{request_provider}/"), + optional_params=dict(body_params(request)), + litellm_params={"model_info": request.get("model_info"), **body_params(request)}, + custom_llm_provider=request_provider, + ) + + +def map_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception: + """Raises what Python's Messages handler raises for the same failure.""" + if getattr(error, "messages_timeout", None) == "Body": + # Python reads a stream's body outside its HTTP client's error mapping, so the transport's + # own timeout reaches the caller unmapped. + return httpx.ReadTimeout(_BODY_READ_TIMEOUT) + if getattr(error, "messages_late_rejection", False): + _log_resolved_call(request, request_provider) + if getattr(error, "messages_request_error", False): + return _rejection(error, request, request_provider) + if getattr(error, "messages_missing_credential", False): + return litellm.AuthenticationError( + message=_MISSING_KEY, + llm_provider=request_provider, + model=str(request["model"]).removeprefix(f"{request_provider}/"), + ) + provider_error: Final = _provider_error(error, request, request_provider) + if provider_error is None: + return failures.map_native_failure( + error, str(request["model"]), request_provider, request, optional_str(request.get("api_base")) + ) + try: + return _exception_type( + model=str(request["model"]), + custom_llm_provider=request_provider, + original_exception=provider_error, + extra_kwargs=dict(request), + ) + except Exception as mapped: # noqa: BLE001 # exception_type raises the mapped public error + mapped.__context__ = error + return mapped + + def _drop_params(kwargs: Mapping[str, object]) -> bool: return bool(litellm.drop_params) or normalize_drop_params(kwargs.get("drop_params")) is True diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 8edd2707405..6a0c867979b 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -5,6 +5,7 @@ from __future__ import annotations from typing import Final import litellm +from litellm.anthropic_beta_headers_manager import is_foundry_api_base from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.opencode.common_utils import is_opencode_api_base from litellm.rust_bridge.detectors import api_base, overridden, present, streaming @@ -15,6 +16,12 @@ def provider_affinity(request: NativeCall) -> bool: return is_opencode_api_base(api_base(request)) +def foundry_beta_row(request: NativeCall) -> bool: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + return is_foundry_api_base(AnthropicModelInfo.get_api_base(api_base(request))) + + def anthropic_workload_identity(request: NativeCall) -> bool: """Python mints an Anthropic federation token when no static credential is configured.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo diff --git a/litellm/secret_managers/aws_secret_manager.py b/litellm/secret_managers/aws_secret_manager.py index 88642ae141e..c87edc4b1dd 100644 --- a/litellm/secret_managers/aws_secret_manager.py +++ b/litellm/secret_managers/aws_secret_manager.py @@ -70,7 +70,7 @@ class AWSKeyManagementService_V2: is_litellm_license_in_env = True if is_litellm_license_in_env is False: raise ValueError( - "AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your envionment." + "AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment." ) def load_aws_kms(self, use_aws_kms: bool | None): diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index 61e47e9d7af..6bba56a5afe 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,5 +1,6 @@ from typing import Final +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse @@ -37,16 +38,19 @@ class TextCompletionStreamWrapper: response["object"] = "text_completion" response["created"] = chunk.get("created", None) response["model"] = chunk.get("model", None) - text_choices: Final = TextChoices() if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices raise Exception if chunk["choices"]: - delta: Final = chunk["choices"][0]["delta"] + chat_choice: Final = chunk["choices"][0] + text_choices: Final = TextChoices( + **{field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} + ) + delta: Final = chat_choice["delta"] text_choices["text"] = delta["content"] text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + text_choices["index"] = chat_choice["index"] + text_choices["finish_reason"] = chat_choice["finish_reason"] response["choices"] = [text_choices] else: response["choices"] = [] diff --git a/litellm/types/decisions.py b/litellm/types/decisions.py index 33b2c677d31..3b28225f0ff 100644 --- a/litellm/types/decisions.py +++ b/litellm/types/decisions.py @@ -210,7 +210,7 @@ class NoulAnswer(LiteLLMPydanticObjectBase): class ChoiceAnswer(LiteLLMPydanticObjectBase): type: Literal["choice"] choice: str - confidence: float + confidence: float | None = None probabilities: Mapping[str, float] model_config = ConfigDict(extra="allow", frozen=True) @@ -219,7 +219,7 @@ class ChoiceAnswer(LiteLLMPydanticObjectBase): class ScoreAnswer(LiteLLMPydanticObjectBase): type: Literal["score"] score: float - confidence: float + confidence: float | None = None legend: Mapping[str, DecisionsJSON] probabilities: Mapping[str, float] @@ -431,7 +431,7 @@ class OpenAIChoiceAnswer(LiteLLMPydanticObjectBase): name: str | None choice: str | bool probabilities: tuple[OpenAIChoiceProbability, ...] - confidence: float + confidence: float | None = None model_config = ConfigDict(frozen=True) @@ -449,7 +449,7 @@ class OpenAIScoreAnswer(LiteLLMPydanticObjectBase): name: str | None score: float probabilities: tuple[OpenAIScoreProbability, ...] - confidence: float + confidence: float | None = None model_config = ConfigDict(frozen=True) @@ -595,7 +595,7 @@ class DecisionsIRChoiceProbability: @dataclass(frozen=True, slots=True) class DecisionsIRChoiceAnswer: choice: str | bool - confidence: float + confidence: float | None probabilities: tuple[DecisionsIRChoiceProbability, ...] extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) @@ -610,7 +610,7 @@ class DecisionsIRScoreProbability: @dataclass(frozen=True, slots=True) class DecisionsIRScoreAnswer: score: float - confidence: float + confidence: float | None probabilities: tuple[DecisionsIRScoreProbability, ...] extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 1f90552e869..9f63d27f5f3 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -55,6 +55,7 @@ RoutingStrategyName: TypeAlias = Literal[ TRUSTED_CALLBACK_VARS_FIELD: Final = "litellm_trusted_callback_vars" ADDRESSED_RESPONSE_ID_FIELD: Final = "_litellm_addressed_response_id" +ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: Final = "_litellm_addressed_comparison_response_id" STREAM_USAGE_INJECTED_FIELD: Final = "_litellm_stream_usage_injected" WIRE_NAME: Final = "wire_name" @@ -407,6 +408,9 @@ class ProxyRequestState: secret_fields: "SecretFields | None" = None trusted_callback_vars: Mapping[str, str] | None = field(default=None, metadata=wire(TRUSTED_CALLBACK_VARS_FIELD)) addressed_response_id: str | None = field(default=None, metadata=wire(ADDRESSED_RESPONSE_ID_FIELD)) + addressed_comparison_response_id: str | None = field( + default=None, metadata=wire(ADDRESSED_COMPARISON_RESPONSE_ID_FIELD) + ) strip_stream_usage: bool | None = field(default=None, metadata=wire("_litellm_strip_stream_usage")) stream_usage_injected: bool | None = field(default=None, metadata=wire(STREAM_USAGE_INJECTED_FIELD)) client_side_timeout: bool | None = None diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index dc35af81091..802f6794839 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -660,6 +660,7 @@ class MessageChunk(TypedDict, total=False): stop_sequence: str | None usage: UsageDelta safeguard_results: ReadOnly[Sequence[SafeguardResult]] + input_transformations: ReadOnly[Sequence[Mapping[str, object]]] class MessageStartBlock(TypedDict): @@ -849,6 +850,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): ADVISOR_TOOL_2026_03_01 = "advisor-tool-2026-03-01" PER_TURN_CONTROL_2026_07_01 = "per-turn-control-2026-07-01" DANGEROUS_TOOL_USE_2026_09_03 = "dangerous-tool-use-2026-09-03" + MCP_CLIENT_2025_04_04 = "mcp-client-2025-04-04" # Tool search beta header constant (for Anthropic direct API and Microsoft Foundry) diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 4f419385630..51168071c85 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -1,4 +1,4 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Literal, TypeAlias from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -106,3 +106,4 @@ class AnthropicMessagesResponse(TypedDict, total=False): usage: AnthropicUsage | None context_management: NotRequired[ContextManagementResponse] safeguard_results: NotRequired[ReadOnly[Sequence[SafeguardResult]]] + input_transformations: NotRequired[ReadOnly[Sequence[Mapping[str, object]]]] diff --git a/litellm/types/llms/custom_llm.py b/litellm/types/llms/custom_llm.py index 6b583b0a83d..bb12449617d 100644 --- a/litellm/types/llms/custom_llm.py +++ b/litellm/types/llms/custom_llm.py @@ -1,11 +1,12 @@ from typing import TYPE_CHECKING -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict if TYPE_CHECKING: + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_llm import CustomLLM class CustomLLMItem(TypedDict): provider: str - custom_handler: "CustomLLM" + custom_handler: ReadOnly["CustomLLM | BaseVideoConfig"] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index d7e55b57a07..c9739eee91d 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1,8 +1,9 @@ import builtins import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import ItemsView, Iterable, Mapping, Sequence from enum import Enum from os import PathLike +from types import MappingProxyType from typing import IO, TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, Union import httpx @@ -692,9 +693,19 @@ class PromptCacheBreakpoint(TypedDict): ttl: NotRequired[ReadOnly[Literal["30m", "12h"]]] +class ChatCompletionPromptCacheBreakpoint(TypedDict): + mode: ReadOnly[Literal["explicit"]] + + +class AzureResponsesPromptCacheBreakpoint(TypedDict): + mode: ReadOnly[Literal["explicit"]] + ttl: NotRequired[ReadOnly[Literal["30m"]]] + + class PromptCacheOptions(TypedDict, total=False): mode: ReadOnly[Literal["implicit", "explicit"]] ttl: ReadOnly[Literal["30m", "12h"]] + comparison_response_id: ReadOnly[str] class ChatCompletionThinkingBlock(TypedDict, total=False): @@ -1236,9 +1247,12 @@ AllEmbeddingInputValues = str | list[str] | list[int] | list[list[int]] OpenAIAudioTranscriptionOptionalParams = Literal[ "language", + "languages", + "keywords", "prompt", "temperature", "response_format", + "stream", "timestamp_granularities", "include", ] @@ -1754,17 +1768,40 @@ class ResponsesAPIStreamEvents(str, Enum): ERROR = "error" -class ResponseCreatedEvent(BaseLiteLLMOpenAIResponseObject): +_NO_EXTRAS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class BaseResponsesStreamEvent(BaseLiteLLMOpenAIResponseObject): # frozen-ok: guardrails rewrite events in place + def __getitem__(self, key: str) -> object: + return self._mapping_holding(key)[key] + + def get(self, key: str, default: object = None) -> object: + return self._mapping_holding(key).get(key, default) + + def __contains__(self, key: object) -> bool: + return key in self.__dict__ or key in self._extras() + + def items(self) -> ItemsView[str, object]: + return MappingProxyType({**self.__dict__, **self._extras()}).items() + + def _extras(self) -> Mapping[str, object]: + return self.__pydantic_extra__ or _NO_EXTRAS + + def _mapping_holding(self, key: str) -> Mapping[str, object]: + return self.__dict__ if key in self.__dict__ else self._extras() + + +class ResponseCreatedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_CREATED] response: ResponsesAPIResponse -class ResponseInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class ResponseInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS] response: ResponsesAPIResponse -class ResponseCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class ResponseCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_COMPLETED] response: ResponsesAPIResponse _hidden_params: dict = PrivateAttr(default_factory=dict) @@ -1778,12 +1815,12 @@ class ResponseCompletedEvent(BaseLiteLLMOpenAIResponseObject): self._hidden_params = hidden_params -class ResponseFailedEvent(BaseLiteLLMOpenAIResponseObject): +class ResponseFailedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_FAILED] response: ResponsesAPIResponse -class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): +class ResponseIncompleteEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE] response: ResponsesAPIResponse @@ -1791,7 +1828,7 @@ class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): ResponsesTerminalEvent: TypeAlias = ResponseCompletedEvent | ResponseIncompleteEvent | ResponseFailedEvent -class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): +class ResponsePartAddedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_PART_ADDED] item_id: str output_index: int @@ -1799,7 +1836,7 @@ class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): part: dict -class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class ReasoningSummaryTextDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA] item_id: str output_index: int @@ -1807,7 +1844,7 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): delta: str -class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject): +class ReasoningSummaryTextDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE] item_id: str output_index: int @@ -1816,7 +1853,7 @@ class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject): text: str -class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject): +class ReasoningSummaryPartDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE] item_id: str output_index: int @@ -1825,13 +1862,13 @@ class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject): part: BaseLiteLLMOpenAIResponseObject -class OutputItemAddedEvent(BaseLiteLLMOpenAIResponseObject): +class OutputItemAddedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] output_index: int item: BaseLiteLLMOpenAIResponseObject | None -class OutputItemDoneEvent(BaseLiteLLMOpenAIResponseObject): +class OutputItemDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE] output_index: int sequence_number: int = 1 @@ -1851,7 +1888,7 @@ class OpenAIChatCompletionLogprobsContent(TypedDict, total=False): top_logprobs: list[OpenAIChatCompletionLogprobsContentTopLogprobs] -class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject): +class ContentPartAddedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.CONTENT_PART_ADDED] item_id: str output_index: int @@ -1879,7 +1916,7 @@ class ContentPartDonePartReasoningText(BaseLiteLLMOpenAIResponseObject): PART_UNION_TYPES = ContentPartDonePartOutputText | ContentPartDonePartRefusal | ContentPartDonePartReasoningText -class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject): +class ContentPartDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.CONTENT_PART_DONE] item_id: str output_index: int @@ -1887,7 +1924,7 @@ class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject): part: PART_UNION_TYPES -class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class OutputTextDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA] item_id: str output_index: int @@ -1895,7 +1932,7 @@ class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): delta: str -class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject): +class OutputTextAnnotationAddedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED] item_id: str output_index: int @@ -1904,7 +1941,7 @@ class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject): annotation: dict -class OutputTextDoneEvent(BaseLiteLLMOpenAIResponseObject): +class OutputTextDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE] item_id: str output_index: int @@ -1912,7 +1949,7 @@ class OutputTextDoneEvent(BaseLiteLLMOpenAIResponseObject): text: str -class RefusalDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class RefusalDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REFUSAL_DELTA] item_id: str output_index: int @@ -1920,7 +1957,7 @@ class RefusalDeltaEvent(BaseLiteLLMOpenAIResponseObject): delta: str -class RefusalDoneEvent(BaseLiteLLMOpenAIResponseObject): +class RefusalDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REFUSAL_DONE] item_id: str output_index: int @@ -1928,72 +1965,72 @@ class RefusalDoneEvent(BaseLiteLLMOpenAIResponseObject): refusal: str -class FunctionCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class FunctionCallArgumentsDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] item_id: str output_index: int delta: str -class FunctionCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject): +class FunctionCallArgumentsDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE] item_id: str output_index: int arguments: str -class FileSearchCallInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class FileSearchCallInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_IN_PROGRESS] output_index: int item_id: str -class FileSearchCallSearchingEvent(BaseLiteLLMOpenAIResponseObject): +class FileSearchCallSearchingEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_SEARCHING] output_index: int item_id: str -class FileSearchCallCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class FileSearchCallCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_COMPLETED] output_index: int item_id: str -class WebSearchCallInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class WebSearchCallInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_IN_PROGRESS] output_index: int item_id: str -class WebSearchCallSearchingEvent(BaseLiteLLMOpenAIResponseObject): +class WebSearchCallSearchingEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_SEARCHING] output_index: int item_id: str -class WebSearchCallCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class WebSearchCallCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_COMPLETED] output_index: int item_id: str # MCP List Tools Events -class MCPListToolsInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class MCPListToolsInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS] sequence_number: int output_index: int item_id: str -class MCPListToolsCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class MCPListToolsCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED] sequence_number: int output_index: int item_id: str -class MCPListToolsFailedEvent(BaseLiteLLMOpenAIResponseObject): +class MCPListToolsFailedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_FAILED] sequence_number: int output_index: int @@ -2001,14 +2038,14 @@ class MCPListToolsFailedEvent(BaseLiteLLMOpenAIResponseObject): # MCP Call Events -class MCPCallInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_IN_PROGRESS] sequence_number: int output_index: int item_id: str -class MCPCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallArgumentsDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA] output_index: int item_id: str @@ -2016,7 +2053,7 @@ class MCPCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject): sequence_number: int -class MCPCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallArgumentsDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DONE] output_index: int item_id: str @@ -2024,21 +2061,21 @@ class MCPCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject): sequence_number: int -class MCPCallCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_COMPLETED] sequence_number: int item_id: str output_index: int -class MCPCallFailedEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallFailedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_FAILED] sequence_number: int item_id: str output_index: int -class ImageGenerationPartialImageEvent(BaseLiteLLMOpenAIResponseObject): +class ImageGenerationPartialImageEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE] partial_image_index: int b64_json: str @@ -2053,7 +2090,7 @@ class ErrorEventError(BaseLiteLLMOpenAIResponseObject): param: str | dict[str, object] | None = None -class ErrorEvent(BaseLiteLLMOpenAIResponseObject): +class ErrorEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.ERROR] sequence_number: int error: ErrorEventError | None = None @@ -2062,7 +2099,7 @@ class ErrorEvent(BaseLiteLLMOpenAIResponseObject): param: str | Mapping[str, object] | None = None -class GenericEvent(BaseLiteLLMOpenAIResponseObject): +class GenericEvent(BaseResponsesStreamEvent): type: str model_config = ConfigDict(extra="allow", protected_namespaces=()) @@ -2511,6 +2548,17 @@ class OpenAIRealtimeResponseUsage(TypedDict): output_token_details: NotRequired[ReadOnly[OpenAIRealtimeUsageTokenDetails]] +class OpenAIRealtimeTranslationDurationUsage(TypedDict): + type: ReadOnly[Literal["duration"]] + output_seconds: NotRequired[ReadOnly[float]] + input_seconds: NotRequired[ReadOnly[float]] + + +class OpenAIRealtimeTranslationClosedEvent(TypedDict): + type: ReadOnly[Literal["session.closed"]] + usage: ReadOnly[OpenAIRealtimeTranslationDurationUsage] + + class OpenAIRealtimeEventTypes(Enum): SESSION_CREATED = "session.created" # Beta delta event names @@ -2554,6 +2602,7 @@ OpenAIRealtimeEvents = ( | OpenAIRealtimeInputAudioTranscriptionCompleted | OpenAIRealtimeTranscriptionSessionCreated | OpenAIRealtimeErrorEvent + | OpenAIRealtimeTranslationClosedEvent ) OpenAIRealtimeStreamList = list[OpenAIRealtimeEvents] diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 2f269dc87d9..b6bab77ab33 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -1,9 +1,11 @@ +from collections.abc import Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from enum import Enum -from typing import Final, Protocol +from typing import Final, Literal, Protocol, TypeAlias from pydantic import Field -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -28,6 +30,10 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__ SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response" +WebsocketUpstreamFrameKind: TypeAlias = Literal[ + "starts_turn", "ends_turn", "opens_session", "updates_session", "bills_usage", "refuses_turn", "relayed" +] + class BackgroundResponseSettlement(str, Enum): SETTLES = "settles" @@ -63,7 +69,7 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): """ The body of the request """ - response_body: dict | None # only tracked for non-streaming responses + response_body: dict | Sequence[Mapping[str, object]] | None # only tracked for non-streaming responses """ The body of the response """ @@ -75,6 +81,8 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): Optional field, we use this for cost tracking only if it's set. """ + end_time: ReadOnly[datetime] + @dataclass(frozen=True, slots=True) class WebSocketSessionUsage: diff --git a/litellm/types/realtime.py b/litellm/types/realtime.py index 6da73982291..4a9b4228bc8 100644 --- a/litellm/types/realtime.py +++ b/litellm/types/realtime.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from typing import Any, Literal from pydantic import ConfigDict @@ -66,6 +67,41 @@ class RealtimeExpiresAfter(LiteLLMBaseModel): seconds: int | None = None +class RealtimeAudioTranscriptionConfig(LiteLLMBaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + model: str | None = None + delay: str | None = None + keywords: Sequence[str] | None = None + language: str | None = None + languages: Sequence[str] | None = None + prompt: str | None = None + + +class RealtimeAudioInputConfig(LiteLLMBaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + format: str | Mapping[str, object] | None = None + noise_reduction: Mapping[str, object] | None = None + transcription: RealtimeAudioTranscriptionConfig | None = None + turn_detection: Mapping[str, object] | None = None + + +class RealtimeAudioOutputConfig(LiteLLMBaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + format: str | Mapping[str, object] | None = None + language: str | None = None + voice: str | Mapping[str, object] | None = None + + +class RealtimeSessionAudioConfig(LiteLLMBaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + input: RealtimeAudioInputConfig | None = None + output: RealtimeAudioOutputConfig | None = None + + class RealtimeSessionConfig(LiteLLMBaseModel): """ Session configuration nested inside the client_secrets request body. @@ -75,12 +111,12 @@ class RealtimeSessionConfig(LiteLLMBaseModel): Extra/unknown fields are passed through unchanged. """ - model_config = ConfigDict(extra="allow") + model_config = ConfigDict(extra="allow", frozen=True) type: str | None = None model: str | None = None instructions: str | None = None - audio: dict[str, object] | None = None + audio: RealtimeSessionAudioConfig | None = None include: list[str] | None = None max_output_tokens: int | str | None = None output_modalities: list[str] | None = None @@ -129,17 +165,20 @@ class RealtimeTranscriptionSessionRequest(LiteLLMBaseModel): unchanged to the provider. """ - model_config = ConfigDict(extra="allow") + model_config = ConfigDict(extra="allow", frozen=True) # LiteLLM-only routing hint — stripped before forwarding upstream. model: str | None = None input_audio_transcription: dict[str, Any] | None = None + audio: RealtimeSessionAudioConfig | None = None def resolved_model(self) -> str | None: if self.model: return self.model if self.input_audio_transcription: return self.input_audio_transcription.get("model") + if self.audio and self.audio.input and self.audio.input.transcription: + return self.audio.input.transcription.model return None @@ -151,7 +190,7 @@ class RealtimeTranscriptionSessionResponse(LiteLLMBaseModel): ephemeral key. Unknown fields pass through unchanged. """ - model_config = ConfigDict(extra="allow") + model_config = ConfigDict(extra="allow", frozen=True) client_secret: dict[str, object] | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index 41a986837da..4e3e4d92b3c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1061,6 +1061,7 @@ class ModelGroupInfo(LiteLLMBaseModel): supports_vision: bool = Field(default=False) supports_web_search: bool = Field(default=False) supports_url_context: bool = Field(default=False) + supports_continuation: bool = Field(default=False) supports_reasoning: bool = Field(default=False) supports_function_calling: bool = Field(default=False) supports_fast_mode: bool = Field(default=False) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index da1d99957a7..57f4d971eb3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -227,6 +227,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_tool_search: bool | None supports_mid_conversation_system: bool | None supports_url_context: bool | None + supports_continuation: ReadOnly[bool | None] supports_none_reasoning_effort: bool | None supports_minimal_reasoning_effort: bool | None supports_low_reasoning_effort: bool | None @@ -246,6 +247,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False): bedrock_converse_supports_strict_tools: bool | None supports_regex_lookaround: ReadOnly[bool | None] token_count_framing: ReadOnly[int | None] + token_count_tool_framing: ReadOnly[int | None] + token_count_responses_framing: ReadOnly[int | None] + token_count_responses_tool_framing: ReadOnly[int | None] token_count_model: ReadOnly[TokenCountModelInfo | None] token_count_tokenizer: ReadOnly[TokenCountTokenizer | None] @@ -550,6 +554,11 @@ class CallTypes(str, Enum): aextraction = "aextraction" arealtime = "_arealtime" aresponses_websocket = "_aresponses_websocket" + acreate_realtime_client_secret = "acreate_realtime_client_secret" + arealtime_calls = "arealtime_calls" + acreate_realtime_transcription_session = "acreate_realtime_transcription_session" + acreate_realtime_translation_client_secret = "acreate_realtime_translation_client_secret" + arealtime_translation_calls = "arealtime_translation_calls" create_batch = "create_batch" acreate_batch = "acreate_batch" aretrieve_batch = "aretrieve_batch" @@ -790,6 +799,8 @@ CallTypesLiteral = Literal[ "acreate_realtime_client_secret", "arealtime_calls", "acreate_realtime_transcription_session", + "acreate_realtime_translation_client_secret", + "arealtime_translation_calls", ] MCP_GUARDRAIL_CALL_TYPES: Final[frozenset[str]] = frozenset( @@ -798,6 +809,10 @@ MCP_GUARDRAIL_CALL_TYPES: Final[frozenset[str]] = frozenset( # Mapping of API routes to their corresponding call types API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { + "/v1/realtime/translations/client_secrets": (CallTypes.acreate_realtime_translation_client_secret,), + "/realtime/translations/client_secrets": (CallTypes.acreate_realtime_translation_client_secret,), + "/v1/realtime/translations/calls": (CallTypes.arealtime_translation_calls,), + "/realtime/translations/calls": (CallTypes.arealtime_translation_calls,), # Chat Completions "/chat/completions": [CallTypes.acompletion, CallTypes.completion], "/v1/chat/completions": [CallTypes.acompletion, CallTypes.completion], @@ -1116,9 +1131,11 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { "/v1/responses/{response_id}/input_items": (CallTypes.alist_input_items,), "/openai/v1/responses/{response_id}/input_items": (CallTypes.alist_input_items,), # Realtime API - "/realtime": [CallTypes.arealtime], - "/v1/realtime": [CallTypes.arealtime], - "/openai/v1/realtime": [CallTypes.arealtime], + "/realtime": (CallTypes.arealtime,), + "/v1/realtime": (CallTypes.arealtime,), + "/openai/v1/realtime": (CallTypes.arealtime,), + "/realtime/translations": (CallTypes.arealtime,), + "/v1/realtime/translations": (CallTypes.arealtime,), "/live/sessions": (CallTypes.arealtime,), "/v1/live/sessions": (CallTypes.arealtime,), "/openai/v1/live/sessions": (CallTypes.arealtime,), @@ -2947,9 +2964,19 @@ class TranscriptionUsageTokensObject(LiteLLMBaseModel): input_token_details: TranscriptionUsageInputTokenDetailsObject | None = None +class TranscriptionDetectedLanguage(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + code: str + + +_DETECTED_LANGUAGES: Final = TypeAdapter(tuple[TranscriptionDetectedLanguage, ...]) + + class TranscriptionResponse(OpenAIObject): text: str | None = None usage: TranscriptionUsageDurationObject | TranscriptionUsageTokensObject | None = None + languages: Sequence[TranscriptionDetectedLanguage] | None = None _hidden_params: dict = {} _response_headers: dict | None = None @@ -2962,8 +2989,28 @@ class TranscriptionResponse(OpenAIObject): def hidden_params(self, hidden_params: dict[str, object]) -> None: # mutable-ok: API requires mutation self._hidden_params = hidden_params - def __init__(self, text=None) -> None: - super().__init__(text=text) + def __init__( + self, + text: str | None = None, + usage: TranscriptionUsageDurationObject | TranscriptionUsageTokensObject | Mapping[str, object] | None = None, + languages: Sequence[TranscriptionDetectedLanguage | Mapping[str, object]] | None = None, + **kwargs: object, # kwargs-ok: OpenAI-compatible responses preserve provider extension fields + ) -> None: + super().__init__( + text=text, usage=usage, languages=TranscriptionResponse.detected_languages(languages), **kwargs + ) + if self.languages is None: + del self.languages + + @staticmethod + def detected_languages(value: object) -> tuple[TranscriptionDetectedLanguage, ...] | None: + try: + return _DETECTED_LANGUAGES.validate_python(value, from_attributes=True) + except ValidationError: + return None + + def set_audio_transcription_duration(self, duration: float) -> None: + self._hidden_params["audio_transcription_duration"] = duration def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator @@ -4231,12 +4278,21 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( "supports_assistant_prefill", "supports_bedrock_native_structured_output", "supports_pdf_input", + "supports_continuation", "reasoning_modes", "reasoning_contexts", } ) -DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) +DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset( + { + "token_count_framing", + "token_count_tool_framing", + "token_count_responses_framing", + "token_count_responses_tool_framing", + "token_count_model", + } +) SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) @@ -4343,6 +4399,7 @@ bedrock_batch_litellm_params: Final = BEDROCK_BATCH_KWARG_NAMES TRUSTED_CALLBACK_VARS_FIELD: Final = _litellm_params.TRUSTED_CALLBACK_VARS_FIELD ADDRESSED_RESPONSE_ID_FIELD: Final = _litellm_params.ADDRESSED_RESPONSE_ID_FIELD +ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: Final = _litellm_params.ADDRESSED_COMPARISON_RESPONSE_ID_FIELD anthropic_wif_litellm_params: Final = tuple(sorted(ANTHROPIC_WIF_KWARGS_KEYS)) openai_wif_litellm_params: Final = tuple(sorted(OPENAI_WIF_KWARGS_KEYS)) @@ -4566,6 +4623,7 @@ class LlmProviders(str, Enum): CLF_AI_GATEWAY = "clf_ai_gateway" PERPLEXITY = "perplexity" TYPESAFE = "typesafe" + PARALLEL_AI = "parallel_ai" NACE = "nace" STRANDS_DECIDER = "strands_decider" LAYA = "laya" diff --git a/litellm/utils.py b/litellm/utils.py index b0826f84006..e003f0a3987 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -83,6 +83,10 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + CHAT_ALLOWED_TOOLS_LAYOUT, + parse_chat_allowed_tools_tool_choice, +) from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, normalize_drop_params, @@ -98,7 +102,11 @@ from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens -from litellm.litellm_core_utils.tool_fields import MalformedField, resolve_deployment_tool_fields +from litellm.litellm_core_utils.tool_fields import ( + MalformedField, + resolve_deployment_tool_fields, + tools_without_hosted_mcp, +) from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge.catalog import decision from litellm.rust_bridge.configuration import Decision @@ -459,6 +467,7 @@ if TYPE_CHECKING: ) from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.custom_llm import CustomLLM from litellm.types.llms.anthropic import ( ANTHROPIC_API_ONLY_HEADERS, AnthropicThinkingParam, @@ -1115,6 +1124,38 @@ def _is_streaming_response_for_correlation(result: object) -> bool: return isinstance(result, CustomStreamWrapper) +def _returns_requested_stream(kwargs: dict[str, object], call_type: CallTypes | str, result: object) -> bool: + if call_type in (CallTypes.transcription, CallTypes.atranscription): + return isinstance(result, (openai.Stream, openai.AsyncStream)) + return _is_streaming_request(kwargs=kwargs, call_type=call_type) + + +def _transcription_runs_unstreamed( + call_type: str, requested_kwargs: Mapping[str, object], call_kwargs: Mapping[str, object] +) -> bool: + from litellm.litellm_core_utils.audio_utils.transcription_streaming import transcription_stream_requested + + return ( + call_type == CallTypes.atranscription.value + and transcription_stream_requested(requested_kwargs.get("stream")) + and call_kwargs.get("stream") is False + ) + + +def _guardrail_view_of_requested_stream(call_kwargs: Mapping[str, object]) -> dict[str, object]: + from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming + + return guardrail_request_data_with_streaming(call_kwargs, is_streaming=True) + + +def _replayed_when_run_unstreamed(result: object, runs_unstreamed: bool) -> object: + from litellm.litellm_core_utils.audio_utils.transcription_streaming import GuardedTranscriptionReplay + + if not runs_unstreamed or not isinstance(result, TranscriptionResponse): + return result + return GuardedTranscriptionReplay(result) + + def _is_converted_stream_result(result: object) -> bool: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator @@ -1550,7 +1591,6 @@ def function_setup( applied_guardrails=applied_guardrails, supports_correlation_logging=is_async_call, ) - ## check if metadata is passed in litellm_params: Final[dict[str, object]] = {"api_base": ""} if "metadata" in kwargs: @@ -1798,7 +1838,28 @@ async def async_pre_call_deployment_hook(kwargs: Mapping[str, object], call_type if result is not None: modified_kwargs = result - return modified_kwargs + return _unstreamed_when_post_call_guardrails_apply(modified_kwargs, typed_call_type) + + +def _unstreamed_when_post_call_guardrails_apply( + kwargs: dict[str, object], call_type: CallTypes | None +) -> dict[str, object]: + from litellm.litellm_core_utils.audio_utils.transcription_streaming import transcription_stream_requested + + if ( + call_type is not CallTypes.atranscription + or not transcription_stream_requested(kwargs.get("stream")) + or _deployment_provider(kwargs) not in (LlmProviders.OPENAI.value, LlmProviders.AZURE.value) + ): + return kwargs + streamed_request: Final = _guardrail_view_of_requested_stream(kwargs) + custom_guardrail_class: Final = _get_cached_custom_guardrail() + if not any( + isinstance(callback, custom_guardrail_class) and callback.runs_post_call_deployment_hook(streamed_request) + for callback in litellm.callbacks + ): + return kwargs + return {**kwargs, "stream": False} async def async_post_call_success_deployment_hook( @@ -2146,7 +2207,7 @@ def client(original_function): # MODEL CALL result = original_function(*args, **kwargs) end_time = datetime.datetime.now() - if _is_streaming_request(kwargs=kwargs, call_type=call_type) or _is_converted_stream_result(result): + if _returns_requested_stream(kwargs, call_type, result) or _is_converted_stream_result(result): if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks = [] for idx, chunk in enumerate(result): @@ -2155,6 +2216,15 @@ def client(original_function): else: logging_obj.stream = True logging_obj.model_call_details["stream"] = True + if call_type == CallTypes.transcription.value and isinstance(result, openai.Stream): + from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( + litellm_exception_mapper, + wrap_transcription_stream, + ) + + result = wrap_transcription_stream( + result, logging_obj, start_time, litellm_exception_mapper(logging_obj.model_call_details) + ) # RETURN RESULT update_response_metadata: _ResponseMetadataUpdater = litellm_utils.update_response_metadata update_response_metadata( @@ -2330,6 +2400,7 @@ def client(original_function): kwargs["litellm_logging_obj"] = logging_obj modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type) + transcription_runs_unstreamed: Final = _transcription_runs_unstreamed(call_type, kwargs, modified_kwargs) if modified_kwargs is not None: kwargs = modified_kwargs @@ -2380,7 +2451,9 @@ def client(original_function): if _is_converted_stream_result(_caching_handler_response.cached_result): logging_obj.stream = True logging_obj.model_call_details["stream"] = True - return _caching_handler_response.cached_result + return _replayed_when_run_unstreamed( + _caching_handler_response.cached_result, transcription_runs_unstreamed + ) elif _caching_handler_response.embedding_all_elements_cache_hit is True: return _caching_handler_response.final_embedding_cached_response @@ -2443,7 +2516,7 @@ def client(original_function): raise end_time = datetime.datetime.now() - streaming_requested: Final = _is_streaming_request(kwargs=kwargs, call_type=call_type) + streaming_requested: Final = _returns_requested_stream(kwargs, call_type, result) if streaming_requested or _is_converted_stream_result(result): logging_obj.stream = True logging_obj.model_call_details["stream"] = True @@ -2457,6 +2530,15 @@ def client(original_function): chunks.append(chunk) return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: + if call_type == CallTypes.atranscription.value and isinstance(result, openai.AsyncStream): + from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( + litellm_exception_mapper, + wrap_transcription_stream, + ) + + result = wrap_transcription_stream( + result, logging_obj, start_time, litellm_exception_mapper(logging_obj.model_call_details) + ) _update_response_metadata( result=result, logging_obj=logging_obj, @@ -2483,7 +2565,9 @@ def client(original_function): _call_type_enum: Final = _CALL_TYPE_ENUM_MAP.get(call_type) if _call_type_enum is not None: result = await async_post_call_success_deployment_hook( - request_data=call_kwargs, + request_data=_guardrail_view_of_requested_stream(call_kwargs) + if transcription_runs_unstreamed + else call_kwargs, response=result, call_type=_call_type_enum, ) @@ -2534,7 +2618,7 @@ def client(original_function): is_litellm_internal_call=_is_litellm_internal_call, ) - return result + return _replayed_when_run_unstreamed(result, transcription_runs_unstreamed) except Exception as e: traceback_exception: Final = traceback.format_exc() # Reuse the timestamp taken right when the deployment call itself failed, before @@ -3183,6 +3267,14 @@ def supports_url_context(model: str, custom_llm_provider: str | None = None) -> ) +def supports_continuation(model: str, custom_llm_provider: str | None = None) -> bool: + return supports_factory( + model=model, + custom_llm_provider=custom_llm_provider, + key="supports_continuation", + ) + + def supports_native_streaming(model: str, custom_llm_provider: str | None) -> bool: """ Check if the given model supports native streaming and return a boolean value. @@ -4134,6 +4226,54 @@ def _register_on_current_map( _runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload +_runtime_registration_order: Final[dict[str, int]] = {} # mutable-ok: tells which of a row's registered names came last +_runtime_registered_names: Final[dict[str, tuple[str, ...]]] = {} # mutable-ok: names registered onto each live row +_REGISTRATION_TICKS: Final = itertools.count() + + +@dataclass(frozen=True, slots=True) +class _RegistrationsOverCatalog(Mapping[str, Mapping[str, object]]): + catalog: Mapping[str, Mapping[str, object]] + registrations: Mapping[str, Mapping[str, object]] + registration_order: Mapping[str, int] + registered_names: Mapping[str, tuple[str, ...]] + + def __getitem__(self, key: str) -> Mapping[str, object]: + catalog_entry: Final = self.catalog.get(key) + names: Final = dict.fromkeys((*self.registered_names.get(key, ()), key)) + present: Final = tuple(name for name in names if name in self.registrations) + if not present: + return self.catalog[key] + oldest_first: Final = sorted(present, key=lambda name: self.registration_order.get(name, -1)) + registered: Final = (self.registrations[name] for name in oldest_first) + return MappingProxyType( + dict(itertools.chain.from_iterable(entry.items() for entry in (catalog_entry or {}, *registered))) + ) + + def __iter__(self) -> Iterator[str]: + return iter({**dict.fromkeys(self.catalog), **dict.fromkeys(self.registrations)}) + + def __len__(self) -> int: + return len(self.catalog.keys() | self.registrations.keys()) + + +def _record_registered_name(written_key: str, registered_key: str) -> None: + names: Final = _runtime_registered_names.get(written_key, ()) + if registered_key not in names: + _runtime_registered_names[written_key] = (*names, registered_key) + + +def catalog_with_runtime_registrations() -> Mapping[str, Mapping[str, object]]: + """The loaded catalog plus ``register_model`` calls, without what routers register for their deployments.""" + catalog: Final = GetModelCostMap.loaded_model_cost_map() + if not _runtime_registered_model_cost: + return catalog + return _RegistrationsOverCatalog( + catalog, + MappingProxyType(_runtime_registered_model_cost), + MappingProxyType(_runtime_registration_order), + MappingProxyType(_runtime_registered_names), + ) def runtime_model_cost_registrations( @@ -4199,7 +4339,14 @@ def reapply_runtime_model_cost_registrations() -> None: if _LiveDeploymentReplay.callback is not None: _LiveDeploymentReplay.callback() if _runtime_registered_model_cost: - register_model(model_cost=dict(_runtime_registered_model_cost)) + register_model( + model_cost=dict( + sorted( + _runtime_registered_model_cost.items(), + key=lambda registration: _runtime_registration_order.get(registration[0], -1), + ) + ) + ) def cost_map_omits_token_price(*keys: object) -> bool: @@ -4273,9 +4420,12 @@ def register_model( dict(_runtime_registered_model_cost.get(_registered_key, {})), _RATE_TABLE.validate_python(_registered_value), ) + _runtime_registration_order[_registered_key] = next(_REGISTRATION_TICKS) for key, value in loaded_model_cost.items(): write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name) + if persist_across_reloads: + _record_registered_name(write.key, str(key)) # Invalidate case-insensitive lookup map since model_cost was modified _invalidate_model_cost_lowercase_map() @@ -4316,40 +4466,61 @@ def _get_non_default_params(passed_params: dict, default_params: dict, additiona return non_default_params +_OPENAI_SDK_ONLY_TRANSCRIPTION_PARAMS: Final = frozenset(("languages", "keywords", "stream")) + + def get_optional_params_transcription( model: str, custom_llm_provider: str, language: str | None = None, + languages: Sequence[str] | None = None, + keywords: Sequence[str] | None = None, prompt: str | None = None, response_format: str | None = None, temperature: int | None = None, timestamp_granularities: list[Literal["word", "segment"]] | None = None, + stream: bool | None = None, drop_params: bool | None = None, **kwargs, ): from litellm.constants import OPENAI_TRANSCRIPTION_PARAMS # retrieve all parameters passed to the function - passed_params: Final = locals() + passed_params: Final = locals().copy() passed_params.pop("OPENAI_TRANSCRIPTION_PARAMS") - passed_params.pop("custom_llm_provider") - passed_params.pop("drop_params") - drop_params = normalize_drop_params(drop_params) + passed_params.pop("model") + custom_llm_provider = passed_params.pop("custom_llm_provider") + drop_params = normalize_drop_params(passed_params.pop("drop_params")) special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs") + from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( + transcription_stream_flag, + transcription_stream_unset, + ) + for k, v in special_params.items(): passed_params[k] = v - default_params: Final = { - "language": None, - "prompt": None, - "response_format": None, - "temperature": None, # openai defaults this to 0 - "timestamp_granularities": None, - } + sends_through_openai_sdk: Final = custom_llm_provider in ("openai", "azure") + drops_unsupported_params: Final = drop_params is True or litellm.drop_params is True + transcription_params: Final = ( + OPENAI_TRANSCRIPTION_PARAMS + if sends_through_openai_sdk or drops_unsupported_params + else [param for param in OPENAI_TRANSCRIPTION_PARAMS if param not in _OPENAI_SDK_ONLY_TRANSCRIPTION_PARAMS] + ) + default_params: Final = {param: None for param in (*transcription_params, "prompt", "temperature")} + stream_unset: Final = transcription_stream_unset(stream) + malformed_stream: Final = not stream_unset and transcription_stream_flag(stream) is None + if malformed_stream and not drops_unsupported_params: + raise UnsupportedParamsError( + status_code=400, + message="`stream` must be a boolean. To drop it from the call, set `litellm.drop_params = True`.", + ) non_default_params: Final = { - k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k]) + k: v + for k, v in passed_params.items() + if k in default_params and v != default_params[k] and not (k == "stream" and (stream_unset or malformed_stream)) } optional_params = {} @@ -4358,9 +4529,7 @@ def get_optional_params_transcription( if len(non_default_params.keys()) > 0: keys: Final = list(non_default_params.keys()) for k in keys: - if ( - drop_params is True or litellm.drop_params is True - ) and k not in supported_params: # drop the unsupported non-default values + if drops_unsupported_params and k not in supported_params: non_default_params.pop(k, None) elif k not in supported_params: raise UnsupportedParamsError( @@ -4414,9 +4583,12 @@ def get_optional_params_transcription( optional_params=optional_params, passed_params=passed_params, custom_llm_provider=custom_llm_provider, - openai_params=list(default_params) if custom_llm_provider == "bedrock" else OPENAI_TRANSCRIPTION_PARAMS, + openai_params=list(default_params) if custom_llm_provider == "bedrock" else transcription_params, additional_drop_params=kwargs.get("additional_drop_params", None), ) + extra_body: Final = optional_params.get("extra_body") + if isinstance(extra_body, dict) and not extra_body: + optional_params.pop("extra_body") return optional_params @@ -5186,7 +5358,12 @@ def remove_sensitive_keys_from_dict(d: dict) -> dict: return d -def pre_process_optional_params(passed_params: dict, non_default_params: dict, custom_llm_provider: str) -> dict: +_TOOL_LIST: Final = TypeAdapter(list[object]) + + +def pre_process_optional_params( + passed_params: dict, non_default_params: dict, custom_llm_provider: str, model: str, drop_params: bool | None +) -> dict: """For .completion(), preprocess optional params""" optional_params: dict = {} @@ -5243,14 +5420,20 @@ def pre_process_optional_params(passed_params: dict, non_default_params: dict, c and custom_llm_provider not in litellm.openai_compatible_providers ): if custom_llm_provider == "ollama": + if "tools" in non_default_params: + sent_tools: Final = _TOOL_LIST.validate_python(non_default_params.pop("tools")) + non_default_params.pop("tool_choice", None) # causes ollama requests to hang + prompt_tools: Final = tools_without_hosted_mcp( + sent_tools, model=model, route=custom_llm_provider, drop_params=bool(drop_params) + ) + if sent_tools and not prompt_tools: + return optional_params + optional_params["functions_unsupported_model"] = prompt_tools + elif "functions" in non_default_params: + optional_params["functions_unsupported_model"] = non_default_params.pop("functions") # ollama actually supports json output optional_params["format"] = "json" litellm.add_function_to_prompt = True # so that main.py adds the function call to the prompt - if "tools" in non_default_params: - optional_params["functions_unsupported_model"] = non_default_params.pop("tools") - non_default_params.pop("tool_choice", None) # causes ollama requests to hang - elif "functions" in non_default_params: - optional_params["functions_unsupported_model"] = non_default_params.pop("functions") elif litellm.add_function_to_prompt: # if user opts to add it to prompt instead optional_params["functions_unsupported_model"] = non_default_params.pop( "tools", non_default_params.pop("functions", None) @@ -5385,6 +5568,8 @@ def get_optional_params( passed_params=passed_params, non_default_params=non_default_params, custom_llm_provider=custom_llm_provider, + model=model, + drop_params=drop_params, ) def _check_valid_arg(supported_params: list[str]): @@ -5976,6 +6161,14 @@ EXTRA_BODY_ROUTING_KEYS: Final = frozenset({"model"}) LITELLM_ONLY_PARAMS: Final = frozenset(("cache_control_injection_points",)) +def header_auth_params(custom_llm_provider: str) -> frozenset[str]: + if custom_llm_provider != "azure_ai": + return frozenset() + from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_CREDENTIAL_PARAM_KEYS + + return AZURE_ENTRA_CREDENTIAL_PARAM_KEYS + + def add_provider_specific_params_to_optional_params( optional_params: dict, passed_params: dict, @@ -6010,8 +6203,10 @@ def add_provider_specific_params_to_optional_params( **extra_body, } - dropped_keys: Final = EXTRA_BODY_ROUTING_KEYS | frozenset( - param for param in (additional_drop_params or ()) if isinstance(param, str) + dropped_keys: Final = ( + EXTRA_BODY_ROUTING_KEYS + | header_auth_params(custom_llm_provider) + | frozenset(param for param in (additional_drop_params or ()) if isinstance(param, str)) ) ensure_extra_body_is_safe: Final = litellm_utils.ensure_extra_body_is_safe optional_params["extra_body"] = ensure_extra_body_is_safe( @@ -7152,6 +7347,7 @@ def get_model_info_helper( "cache_read_input_audio_token_cost_batches", None ), explicit_cache_read_input_token_cost=_model_info.get("explicit_cache_read_input_token_cost", None), + cache_read_input_image_token_cost=_model_info.get("cache_read_input_image_token_cost", None), prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( "cache_read_input_token_cost_above_200k_tokens", None @@ -7360,6 +7556,7 @@ def get_model_info_helper( ), supports_web_search=_model_info.get("supports_web_search", None), supports_url_context=_model_info.get("supports_url_context", None), + supports_continuation=_model_info.get("supports_continuation", None), supports_reasoning=_model_info.get("supports_reasoning", None), supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", None), supports_legacy_thinking=_model_info.get("supports_legacy_thinking", None), @@ -7398,6 +7595,9 @@ def get_model_info_helper( supports_image_size=_model_info.get("supports_image_size", None), supported_audio_formats=_model_info.get("supported_audio_formats", None), token_count_framing=_model_info.get("token_count_framing", None), + token_count_tool_framing=_model_info.get("token_count_tool_framing", None), + token_count_responses_framing=_model_info.get("token_count_responses_framing", None), + token_count_responses_tool_framing=_model_info.get("token_count_responses_tool_framing", None), token_count_model=_model_info.get("token_count_model", None), token_count_tokenizer=_model_info.get("token_count_tokenizer", None), vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), @@ -7516,6 +7716,7 @@ def get_model_info( supports_pdf_input: Optional[bool] supports_web_search: Optional[bool] supports_url_context: Optional[bool] + supports_continuation: Optional[bool] supports_reasoning: Optional[bool] Raises: Exception: If the model is not mapped yet. @@ -8868,6 +9069,9 @@ def bundled_provider_models(provider: str) -> tuple[str, ...]: return tuple(model for model in bundled if not is_hidden_as_retired(cost_map.get(model, {}))) +AZURE_MODEL_PLACEHOLDER: Final = "Azure-LLM" + + def get_valid_models( check_provider_endpoint: bool | None = None, custom_llm_provider: str | None = None, @@ -8923,7 +9127,7 @@ def get_valid_models( continue if provider == "azure": - valid_models.append("Azure-LLM") + valid_models.append(AZURE_MODEL_PLACEHOLDER) elif provider_config is not None and check_provider_endpoint and provider is not None: valid_models.extend( _get_valid_models_from_provider_api( @@ -9440,9 +9644,28 @@ def validate_chat_completion_tool_choice( tool_choice_type = tool_choice.get("type") if tool_choice_type in ("auto", "none", "required") and "function" not in tool_choice: return tool_choice_type + if ( + tool_choice_type == "allowed_tools" + and ( + parse_chat_allowed_tools_tool_choice( + cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic + object, tool_choice + ) + ) + is None + ) + ): + raise BadRequestError( + message=( + f"Invalid tool choice, tool_choice={tool_choice}. On chat completions an allowed_tools " + f"tool_choice must be {CHAT_ALLOWED_TOOLS_LAYOUT}" + ), + model=model, + llm_provider="", + ) # Standard OpenAI format: {"type": "function", "function": {...}} - if tool_choice.get("type") is None or tool_choice.get("function") is None: + if tool_choice_type != "allowed_tools" and (tool_choice_type is None or tool_choice.get("function") is None): raise BadRequestError( message=f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec", model=model, @@ -10052,12 +10275,19 @@ class ProviderConfigManager: return litellm.LayaDecisionsConfig() if provider == LlmProviders.BESPOKE: return litellm.BespokeDecisionsConfig() - if provider in (LlmProviders.OPENAI, LlmProviders.AZURE, LlmProviders.AZURE_AI, LlmProviders.EU_JEV): + if provider in ( + LlmProviders.OPENAI, + LlmProviders.AZURE, + LlmProviders.AZURE_AI, + LlmProviders.EU_JEV, + LlmProviders.PARALLEL_AI, + ): configs: Final = { LlmProviders.OPENAI: litellm.OpenAIDecisionsConfig, LlmProviders.AZURE: litellm.AzureDecisionsConfig, LlmProviders.AZURE_AI: litellm.AzureAIDecisionsConfig, LlmProviders.EU_JEV: litellm.EuJevDecisionsConfig, + LlmProviders.PARALLEL_AI: litellm.ParallelAIDecisionsConfig, } return configs[provider]() return None @@ -10178,6 +10408,7 @@ class ProviderConfigManager: def get_provider_audio_transcription_config( model: str, provider: LlmProviders, + base_model: str | None = None, ) -> BaseAudioTranscriptionConfig | None: model_cost_entry: Final = _get_model_cost_entry_for_provider_config( model=model, @@ -10205,7 +10436,13 @@ class ProviderConfigManager: return XAIAudioTranscriptionConfig() elif litellm.LlmProviders.OPENAI == provider: - if "gpt-4o" in model or (model.startswith("gpt-") and "transcribe" in model): + if "gpt-transcribe" in (model, base_model): + from litellm.llms.openai.transcriptions.gpt_transformation import ( + OpenAIGPTTranscribeAudioTranscriptionConfig, + ) + + return OpenAIGPTTranscribeAudioTranscriptionConfig() + elif "gpt-4o" in model or (model.startswith("gpt-") and "transcribe" in model): return litellm.OpenAIGPTAudioTranscriptionConfig() else: return litellm.OpenAIWhisperAudioTranscriptionConfig() @@ -10505,12 +10742,14 @@ class ProviderConfigManager: from litellm.llms.ollama.common_utils import OllamaModelInfo return OllamaModelInfo() - elif LlmProviders.VLLM == provider or LlmProviders.HOSTED_VLLM == provider: - from litellm.llms.vllm.common_utils import ( - VLLMModelInfo, # experimental approach, to reduce bloat on __init__.py - ) + elif LlmProviders.VLLM == provider: + from litellm.llms.vllm.common_utils import VLLMModelInfo return VLLMModelInfo() + elif LlmProviders.HOSTED_VLLM == provider: + from litellm.llms.vllm.common_utils import HostedVLLMModelInfo + + return HostedVLLMModelInfo() elif LlmProviders.LEMONADE == provider: return litellm.LemonadeChatConfig() elif LlmProviders.CLARIFAI == provider: @@ -10896,8 +11135,13 @@ class ProviderConfigManager: @staticmethod def get_provider_video_config( model: str | None, - provider: LlmProviders, + provider: LlmProviders | str, ) -> BaseVideoConfig | None: + if not isinstance(provider, LlmProviders): + custom_video_config: Final = ProviderConfigManager.get_custom_provider_video_config(provider) + if custom_video_config is not None or provider not in LlmProvidersSet: + return custom_video_config + return ProviderConfigManager.get_provider_video_config(model=model, provider=LlmProviders(provider)) if LlmProviders.OPENAI == provider: from litellm.llms.openai.videos.transformation import OpenAIVideoConfig @@ -10948,6 +11192,34 @@ class ProviderConfigManager: return VolcEngineVideoConfig() return None + @staticmethod + def get_custom_llm_handler(provider: str | None) -> CustomLLM | None: + """The `CustomLLM` registered for `provider` in `litellm.custom_provider_map`, if its handler is one.""" + from litellm.llms.custom_llm import CustomLLM + + return next( + ( + item["custom_handler"] + for item in litellm.custom_provider_map + if item["provider"] == provider and isinstance(item["custom_handler"], CustomLLM) + ), + None, + ) + + @staticmethod + def get_custom_provider_video_config(provider: str) -> BaseVideoConfig | None: + """The `BaseVideoConfig` registered for `provider` in `litellm.custom_provider_map`, if its handler is one.""" + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig + + return next( + ( + item["custom_handler"] + for item in litellm.custom_provider_map + if item["provider"] == provider and isinstance(item["custom_handler"], BaseVideoConfig) + ), + None, + ) + @staticmethod def get_provider_container_config( provider: LlmProviders, diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 85e651b687c..66eb10e30db 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -26,6 +26,17 @@ from litellm.types.videos.utils import decode_video_id_with_provider from litellm.utils import ProviderConfigManager, client from litellm.videos.utils import VideoGenerationRequestUtils + +def _unsupported_video_call( + call_name: str, custom_llm_provider: str, model: str | None = None +) -> litellm.BadRequestError: + return litellm.BadRequestError( + message=f"{call_name} is not supported for {custom_llm_provider}", + model=model or "", + llm_provider=custom_llm_provider, + ) + + #################### Initialize provider clients #################### llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() @@ -40,7 +51,7 @@ async def avideo_generation( size: str | None = None, user: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -170,7 +181,7 @@ def video_generation( size: str | None = None, user: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -208,11 +219,11 @@ def video_generation( # get provider config video_generation_provider_config: BaseVideoConfig | None = ProviderConfigManager.get_provider_video_config( model=model, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_generation_provider_config is None: - raise ValueError(f"video generation is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video generation", custom_llm_provider, model) local_vars.update(kwargs) # Get VideoGenerationOptionalRequestParams with only valid parameters @@ -329,11 +340,11 @@ def video_content( # get provider config video_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_provider_config is None: - raise ValueError(f"video support download is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video support download", custom_llm_provider) local_vars.update(kwargs) # For video content download, we don't need complex optional parameter handling @@ -456,7 +467,7 @@ async def avideo_remix( video_id: str, prompt: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -559,7 +570,7 @@ def video_remix( video_id: str, prompt: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -598,11 +609,11 @@ def video_remix( # get provider config video_remix_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_remix_provider_config is None: - raise ValueError(f"video remix is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video remix", custom_llm_provider) local_vars.update(kwargs) # For video remix, we need the video_id and prompt @@ -660,7 +671,7 @@ async def avideo_list( order: str | None = None, api_key: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -774,7 +785,7 @@ def video_list( limit: int | None = None, order: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -810,11 +821,11 @@ def video_list( # get provider config video_list_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_list_provider_config is None: - raise ValueError(f"video list is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video list", custom_llm_provider) local_vars.update(kwargs) # For video list, we need the query parameters @@ -871,7 +882,7 @@ def video_list( async def avideo_status( video_id: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -970,7 +981,7 @@ def video_status( def video_status( video_id: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -1031,11 +1042,11 @@ def video_status( # get provider config video_status_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_status_provider_config is None: - raise ValueError(f"video status is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video status", custom_llm_provider) local_vars.update(kwargs) # For video status, we need the video_id @@ -1088,7 +1099,7 @@ async def avideo_create_character( name: str, video: FileTypes, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1143,7 +1154,7 @@ def video_create_character( name: str, video: FileTypes, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1172,11 +1183,11 @@ def video_create_character( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video create character is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video create character", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"name": name} @@ -1218,7 +1229,7 @@ def video_create_character( async def avideo_get_character( character_id: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1268,7 +1279,7 @@ async def avideo_get_character( def video_get_character( character_id: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1297,11 +1308,11 @@ def video_get_character( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video get character is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video get character", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"character_id": character_id} @@ -1343,7 +1354,7 @@ async def avideo_edit( video_id: str, prompt: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1398,7 +1409,7 @@ def video_edit( video_id: str, prompt: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1430,11 +1441,11 @@ def video_edit( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video edit is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video edit", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"video_id": video_id, "prompt": prompt} @@ -1480,7 +1491,7 @@ async def avideo_extension( prompt: str, seconds: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1534,7 +1545,7 @@ def video_extension( prompt: str, seconds: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1564,11 +1575,11 @@ def video_extension( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video extension is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video extension", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e1fdfc89a9f..974f1bb0c8b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1,4 +1,168 @@ { + "azure/gpt-realtime-2.1-2026-07-07": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 2.4e-05, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "cache_read_input_image_token_cost": 5e-07, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-realtime-2.1-mini-2026-07-07": { + "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, + "input_cost_per_audio_token": 1e-05, + "input_cost_per_image_token": 8e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "cache_read_input_image_token_cost": 8e-08, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-realtime-translate-2026-05-06": { + "input_cost_per_second": 0.0005666666666666667, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/a-new-chapter-for-realtime-ai-reasoning-translation-and-real-time-transcription/4517124", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_native_streaming": true + }, + "azure/gpt-realtime-translate-2026-05-07": { + "input_cost_per_second": 0.0005666666666666667, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "source": "https://ai.azure.com/catalog/models/gpt-realtime-translate", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_native_streaming": true + }, + "azure/gpt-realtime-whisper-2026-05-06": { + "input_cost_per_second": 0.0002833333333333333, + "litellm_provider": "azure", + "mode": "audio_transcription", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/gpt-realtime-whisper", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "max_input_tokens": 32000, + "max_output_tokens": 4096 + }, + "azure/gpt-realtime-whisper-2026-05-07": { + "input_cost_per_second": 0.0002833333333333333, + "litellm_provider": "azure", + "mode": "audio_transcription", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/gpt-realtime-whisper", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "max_input_tokens": 32000, + "max_output_tokens": 4096 + }, "sample_spec": { "code_interpreter_cost_per_session": 0.0, "computer_use_input_cost_per_1k_tokens": 0.0, @@ -6424,6 +6588,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6458,6 +6625,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6491,6 +6661,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6525,6 +6698,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6547,6 +6723,7 @@ "azure/gpt-realtime-2.1": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_image_token_cost": 5e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-06-25", "input_cost_per_audio_token": 3.2e-05, @@ -6556,6 +6733,9 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", @@ -6575,12 +6755,15 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true }, "azure/gpt-realtime-2.1-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_image_token_cost": 8e-08, "cache_read_input_token_cost": 6e-08, "deprecation_date": "2027-06-25", "input_cost_per_audio_token": 1e-05, @@ -6590,6 +6773,9 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", @@ -6609,6 +6795,8 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true }, @@ -6626,6 +6814,9 @@ "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6659,6 +6850,9 @@ "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "supported_endpoints": [ "/v1/realtime" ], @@ -6791,6 +6985,9 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "mode": "audio_transcription", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/gpt-realtime-whisper", "supported_endpoints": [ "/v1/realtime", @@ -13478,6 +13675,10 @@ "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", + "token_count_framing": 149, + "token_count_tool_framing": 302, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 100, "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", @@ -13831,7 +14032,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -13845,7 +14046,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/ap-northeast-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14002,7 +14203,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14016,7 +14217,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/ap-south-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14140,12 +14341,12 @@ "input_cost_per_token": 7.4e-07, "input_cost_per_token_flex": 3.7e-07, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, "output_cost_per_token_flex": 1.11e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.html", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md", "supports_audio_input": false, "supports_function_calling": true, "supports_native_structured_output": true, @@ -14250,7 +14451,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14264,7 +14465,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/eu-north-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14786,7 +14987,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14800,7 +15001,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/sa-east-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15035,7 +15236,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15049,7 +15250,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-east-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15156,7 +15357,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15170,7 +15371,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-east-2/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15552,6 +15753,7 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { + "deprecation_date": "2027-03-30", "input_cost_per_token": 2.65e-06, "litellm_provider": "bedrock", "max_input_tokens": 8000, @@ -15562,6 +15764,7 @@ "supports_pdf_input": true }, "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": { + "deprecation_date": "2027-03-30", "input_cost_per_token": 3e-07, "litellm_provider": "bedrock", "max_input_tokens": 8000, @@ -15708,7 +15911,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15722,7 +15925,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-west-2/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -23519,12 +23722,12 @@ ] }, "deepgram/streaming/nova-3": { - "input_cost_per_second": 8e-05, + "input_cost_per_second": 0.00012833, "litellm_provider": "deepgram", "metadata": { - "calculation": "$0.0048/60 seconds = $0.00008000 per second", + "calculation": "$0.0077/60 seconds = $0.00012833 per second", "note": "Nova-3 monolingual streaming, pay as you go", - "original_pricing_per_minute": 0.0048 + "original_pricing_per_minute": 0.0077 }, "mode": "audio_transcription", "output_cost_per_second": 0.0, @@ -23533,13 +23736,28 @@ "/v1/listen" ] }, - "deepgram/streaming/nova-3-multilingual": { - "input_cost_per_second": 9.667e-05, + "deepgram/streaming/nova-3-medical": { + "input_cost_per_second": 0.00012833, "litellm_provider": "deepgram", "metadata": { - "calculation": "$0.0058/60 seconds = $0.00009667 per second", + "calculation": "$0.0077/60 seconds = $0.00012833 per second", + "note": "Nova-3 Medical streaming, launch price from Deepgram's 2025-03-03 announcement. deepgram.com/pricing does not list Nova-3 Medical, and this matches its current Nova-3 monolingual streaming pay as you go rate", + "original_pricing_per_minute": 0.0077 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://www.businesswire.com/news/home/20250303470439/en", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/nova-3-multilingual": { + "input_cost_per_second": 0.00015333, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0092/60 seconds = $0.00015333 per second", "note": "Nova-3 multilingual (language=multi) streaming, pay as you go", - "original_pricing_per_minute": 0.0058 + "original_pricing_per_minute": 0.0092 }, "mode": "audio_transcription", "output_cost_per_second": 0.0, @@ -24519,7 +24737,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock_converse", - "max_input_tokens": 164000, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -24530,7 +24748,7 @@ "supports_reasoning": true, "supports_native_structured_output": true, "supports_tool_choice": true, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -29217,9 +29435,9 @@ }, "gemini-3-pro-image": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 3.6e-07, "deprecation_date": "2027-05-28", @@ -31203,6 +31421,7 @@ "cache_read_input_token_cost_batches": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, + "deprecation_date": "2026-12-01", "input_cost_per_audio_token": 1e-06, "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 3e-07, @@ -31786,6 +32005,7 @@ "output_cost_per_token_flex": 5e-06, "supports_url_context": true, "cache_read_input_token_cost_above_200k_tokens_batches": 2.5e-07, + "deprecation_date": "2026-12-01", "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 }, @@ -31889,7 +32109,7 @@ }, "gemini/gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_batches": 2e-08, + "cache_read_input_token_cost_batches": 1.5e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, @@ -34788,6 +35008,7 @@ "gpt-4.1": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_priority": 8.75e-07, + "deprecation_date": "2027-01-14", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_priority": 3.5e-06, @@ -34873,6 +35094,7 @@ "gpt-4.1-mini": { "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_priority": 1.75e-07, + "deprecation_date": "2027-01-14", "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, "input_cost_per_token_priority": 7e-07, @@ -35032,6 +35254,7 @@ "gpt-4o": { "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_priority": 2.125e-06, + "deprecation_date": "2027-01-14", "input_cost_per_token": 2.5e-06, "input_cost_per_token_batches": 1.25e-06, "input_cost_per_token_priority": 4.25e-06, @@ -35376,6 +35599,7 @@ "gpt-4o-mini": { "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.25e-07, + "deprecation_date": "2027-01-14", "input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_priority": 2.5e-07, @@ -39044,6 +39268,9 @@ "max_input_tokens": 128000, "max_output_tokens": 32000, "mode": "realtime", + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, "source": "https://developers.openai.com/api/docs/pricing", @@ -39099,6 +39326,8 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true }, @@ -39135,6 +39364,8 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true }, @@ -43081,18 +43312,18 @@ "mode": "embedding" }, "mistral/mistral-large-4": { - "cache_read_input_token_cost": 1.4e-07, - "input_cost_per_token": 1.36e-06, + "cache_read_input_token_cost": 7e-08, + "input_cost_per_token": 7e-07, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 4.18e-06, + "output_cost_per_token": 2.1e-06, "reasoning_effort_levels": [ "none", "high" ], - "source": "https://docs.mistral.ai/models/mistral-large-4", + "source": "https://docs.mistral.ai/models/mistral-large-4-0", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -43102,18 +43333,18 @@ "supports_vision": true }, "mistral/mistral-large-4-0": { - "cache_read_input_token_cost": 1.4e-07, - "input_cost_per_token": 1.36e-06, + "cache_read_input_token_cost": 7e-08, + "input_cost_per_token": 7e-07, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 4.18e-06, + "output_cost_per_token": 2.1e-06, "reasoning_effort_levels": [ "none", "high" ], - "source": "https://docs.mistral.ai/models/mistral-large-4", + "source": "https://docs.mistral.ai/models/mistral-large-4-0", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -46387,6 +46618,19 @@ "supports_response_schema": true, "supports_tool_choice": false }, + "parallel_ai/decision": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "parallel_ai", + "max_input_tokens": 32000, + "mode": "evaluation", + "source": "https://docs.parallel.ai/getting-started/pricing#decisions-api", + "supported_endpoints": [ + "/v1/decisions", + "/v1/systemone" + ] + }, "parallel_ai/search": { "input_cost_per_query": 0.005, "litellm_provider": "parallel_ai", @@ -49614,6 +49858,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -49658,6 +49903,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -52547,9 +52793,9 @@ }, "vertex_ai/gemini-3-pro-image": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 3.6e-07, "deprecation_date": "2027-05-28", @@ -53957,23 +54203,23 @@ "supports_vision": true }, "wandb/meta-llama/Llama-3.1-8B-Instruct": { - "max_input_tokens": 131000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "input_cost_per_token": 2.2e-07, "output_cost_per_token": 2.2e-07, "litellm_provider": "wandb", "mode": "chat", - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/deepseek-ai/DeepSeek-V3.1": { "supports_reasoning": true, - "max_input_tokens": 161000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "input_cost_per_token": 5.5e-07, "output_cost_per_token": 1.65e-06, "litellm_provider": "wandb", "mode": "chat", - "source": "https://wandb.ai/site/pricing/tokens/" + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference" }, "wandb/meta-llama/Llama-3.3-70B-Instruct": { "max_input_tokens": 128000, @@ -54457,6 +54703,8 @@ "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", + "token_count_framing": 185, + "token_count_tool_framing": 214, "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", @@ -59875,6 +60123,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html", "use_openai_responses_path": true, "supported_endpoints": [ @@ -59896,7 +60145,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { "input_cost_per_token": 4.4e-06, @@ -59911,6 +60161,7 @@ "max_input_tokens": 1050000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/responses" @@ -59930,7 +60181,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -60336,6 +60588,7 @@ "mode": "responses", "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, "output_cost_per_token_ultrafast": 0.00033, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60359,7 +60612,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-6-sol": { "input_cost_per_token": 2.2e-06, @@ -60374,6 +60628,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60397,7 +60652,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-6-luna": { "input_cost_per_token": 1.1e-07, @@ -60412,6 +60668,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60435,7 +60692,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", + "supports_web_search": true }, "us.openai.gpt-6-astra": { "supports_bedrock_runtime_chat_completions_response_format": true, @@ -62630,7 +62888,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 23, + "token_count_tool_framing": 12 }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, @@ -62736,7 +62996,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 76, + "token_count_tool_framing": 75 }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, @@ -63349,6 +63611,10 @@ "input_cost_per_second": 7.5e-05, "litellm_provider": "openai", "mode": "audio_transcription", + "provider_specific_entry": { + "transcription_json_or_text_only": 1, + "transcription_usage_duration": 1 + }, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/transcriptions", @@ -63361,7 +63627,8 @@ "supported_output_modalities": [ "text" ], - "supports_audio_input": true + "supports_audio_input": true, + "supports_native_streaming": true }, "gpt-live-transcribe": { "input_cost_per_second": 0.000283333333333, @@ -63379,7 +63646,8 @@ "supported_output_modalities": [ "text" ], - "supports_audio_input": true + "supports_audio_input": true, + "supports_native_streaming": true }, "gpt-live-1": { "input_cost_per_second": 0.000833333333333, @@ -63413,7 +63681,8 @@ "audio" ], "supports_audio_input": true, - "supports_audio_output": true + "supports_audio_output": true, + "supports_native_streaming": true }, "claude-mythos-5": { "supports_anthropic_compaction": true, @@ -64022,6 +64291,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", @@ -64029,6 +64299,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64048,6 +64319,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64090,6 +64362,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64120,6 +64393,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", @@ -64127,6 +64401,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64146,6 +64421,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64220,6 +64496,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64240,6 +64517,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64263,6 +64541,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64338,6 +64617,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64374,6 +64654,7 @@ "output_cost_per_token_priority": 6.375e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_vision": true }, @@ -64410,6 +64691,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64480,6 +64762,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64503,6 +64786,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67296,6 +67580,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67314,6 +67599,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67332,6 +67618,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67350,6 +67637,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67365,6 +67653,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67407,6 +67696,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67425,6 +67715,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67443,6 +67734,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67460,6 +67752,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67478,6 +67771,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67496,6 +67790,7 @@ "output_cost_per_token_priority": 4.05e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -68281,6 +68576,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -68326,6 +68622,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -68554,6 +68851,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -68594,6 +68892,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -68633,7 +68932,9 @@ "cache_read_input_token_cost_above_272k_tokens": 5.4e-07, "output_cost_per_token": 1.62e-05, "output_cost_per_token_above_272k_tokens": 2.43e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { "litellm_provider": "bedrock_mantle", @@ -68668,7 +68969,9 @@ "cache_read_input_token_cost_above_272k_tokens": 5.4e-08, "output_cost_per_token": 1.62e-06, "output_cost_per_token_above_272k_tokens": 2.43e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { "litellm_provider": "bedrock_mantle", @@ -68697,7 +69000,9 @@ "input_cost_per_token": 3.375e-06, "cache_read_input_token_cost": 3.375e-07, "output_cost_per_token": 2.025e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/xai.grok-4.3": { "use_openai_responses_path": true, @@ -71248,6 +71553,9 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/realtime" @@ -71305,7 +71613,11 @@ "supported_output_modalities": [ "text" ], - "supports_audio_input": true + "supports_audio_input": true, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "supports_native_streaming": true }, "azure/gpt-transcribe": { "deprecation_date": "2028-02-01", @@ -71315,7 +71627,8 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/audio/transcriptions", - "/v1/realtime/transcription_sessions" + "/v1/realtime/transcription_sessions", + "/v1/realtime" ], "supported_modalities": [ "text", @@ -71324,7 +71637,13 @@ "supported_output_modalities": [ "text" ], - "supports_audio_input": true + "supports_audio_input": true, + "provider_specific_entry": { + "realtime_ga_only": 1, + "transcription_json_or_text_only": 1, + "transcription_usage_duration": 1 + }, + "supports_native_streaming": true }, "azure/gpt-realtime-translate": { "deprecation_date": "2027-11-09", @@ -71342,7 +71661,11 @@ "audio" ], "supports_audio_input": true, - "supports_audio_output": true + "supports_audio_output": true, + "provider_specific_entry": { + "realtime_ga_only": 1 + }, + "supports_native_streaming": true }, "aihubmix/agnes-2.5-flash": { "input_cost_per_token": 3e-08, @@ -72781,7 +73104,8 @@ "max_input_tokens": 65536, "mode": "evaluation", "output_cost_per_token": 0.0, - "source": "https://docs.typesafe.ai/models" + "source": "https://docs.typesafe.ai/models", + "supports_box_questions": true }, "typesafe/jev-preview": { "input_cost_per_token": 4.2e-08, @@ -73731,6 +74055,7 @@ "output_cost_per_token_priority": 1.875e-05, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -75312,6 +75637,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_ultrafast": 9.9e-05, "output_cost_per_token_ultrafast": 6.6e-05, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/chat/completions", @@ -75334,6 +75660,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, @@ -76629,23 +76956,19 @@ }, "vertex_ai/gemini-robotics-er-2-preview": { "cache_read_input_token_cost": 1e-07, - "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1e-06, - "input_cost_per_token_batches": 5e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 131072, "max_output_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 5e-06, "output_cost_per_token": 5e-06, - "output_cost_per_token_batches": 2.5e-06, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/v1/completions" ], "supported_modalities": [ "text", @@ -76675,23 +76998,19 @@ }, "gemini-robotics-er-2-preview": { "cache_read_input_token_cost": 1e-07, - "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1e-06, - "input_cost_per_token_batches": 5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, "max_output_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 5e-06, "output_cost_per_token": 5e-06, - "output_cost_per_token_batches": 2.5e-06, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/v1/completions" ], "supported_modalities": [ "text", @@ -76758,6 +77077,7 @@ "supports_audio_input": true, "supports_audio_output": false, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76817,6 +77137,7 @@ ], "supports_audio_input": true, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76874,6 +77195,7 @@ "text" ], "supports_audio_input": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -77231,6 +77553,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -77389,6 +77712,7 @@ "output_cost_per_token_priority": 9e-07, "source": "https://platform.minimax.io/docs/guides/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -79022,6 +79346,7 @@ "mode": "responses", "output_cost_per_token": 5.5e-06, "output_cost_per_token_above_272k_tokens": 8.25e-06, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-terra.html", "supported_endpoints": [ "/v1/chat/completions", @@ -79044,6 +79369,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, @@ -81164,6 +81490,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81180,6 +81507,7 @@ "output_cost_per_token": 3.08e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81196,6 +81524,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81212,6 +81541,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 700a35edb06..40ea9d15ff4 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -3158,7 +3158,7 @@ "batches": false, "messages": false, "interactions": false, - "systemone": false, + "systemone": true, "a2a": false, "generateContent": false } diff --git a/pyproject.toml b/pyproject.toml index 5c23a64ebb3..7a653b28f81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "tiktoken>=0.12.0,<1.0; python_version == '3.14'", "tiktoken>=0.14.0,<1.0; python_version >= '3.15'", "tokenizers>=0.21.0,<1.0", - "huggingface-hub>=0.34.0,<2.0", + "huggingface-hub>=0.34.0,<3.0", "click>=8.0.0,<9.0", "jinja2>=3.1.6,<4.0", "aiohttp>=3.14.2,<4.0", diff --git a/scripts/check_openapi_docs.py b/scripts/check_openapi_docs.py index 26bb034bfa6..e29cac29212 100755 --- a/scripts/check_openapi_docs.py +++ b/scripts/check_openapi_docs.py @@ -18,6 +18,7 @@ from collections.abc import Iterator, Mapping from typing import Final, NamedTuple GATED_METHODS: Final = frozenset({"post"}) +SCHEMA_MEDIA_TYPES: Final = ("application/json", "application/sdp") class Violation(NamedTuple): @@ -30,18 +31,21 @@ def _as_mapping(node: object) -> Mapping[str, object]: return node if isinstance(node, Mapping) else {} -def _json_schema(container: object) -> object: +def _content_schema(container: object) -> object: content: Final = _as_mapping(_as_mapping(container).get("content")) - return _as_mapping(content.get("application/json")).get("schema") + return next( + (schema for media_type in SCHEMA_MEDIA_TYPES if (schema := _as_mapping(content.get(media_type)).get("schema"))), + None, + ) def _operation_codes(operation: Mapping[str, object]) -> Iterator[str]: - if not _json_schema(operation.get("requestBody")): + if not _content_schema(operation.get("requestBody")): yield "OA001" description: Final = operation.get("description") if not isinstance(description, str) or not description.strip(): yield "OA002" - if not _json_schema(_as_mapping(operation.get("responses")).get("200")): + if not _content_schema(_as_mapping(operation.get("responses")).get("200")): yield "OA003" diff --git a/tests/code_coverage_tests/check_e2e_no_raw_requests.py b/tests/code_coverage_tests/check_e2e_no_raw_requests.py index bea177d3b7d..94cba6334e0 100644 --- a/tests/code_coverage_tests/check_e2e_no_raw_requests.py +++ b/tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -3,7 +3,7 @@ raw HTTP client imports (requests, urllib.request, httpx, aiohttp, http.client) banned in suite code, and tests/e2e_harness, which tests that harness, is held to the same ban. Importing requests' exception types for catching is fine anywhere; a small allowlist grandfathers the files that legitimately make raw calls (the transport -itself, the root conftest liveness probe, the claude_code version resolver's constant +itself, the root conftest liveness probe, the cli_harness npm version resolver's constant registry URL fetch, and the mcp OAuth client, whose httpx client is the object the official mcp SDK's streamable_http_client requires and so cannot go through the sync requests transport). Referenced by tests/e2e/AGENTS.md.""" @@ -24,7 +24,7 @@ BANNED_MODULES = ("requests", "urllib.request", "http.client", "httpx", "aiohttp ALLOWED_RAW_CLIENT_FILES = { "e2e/e2e_http.py": ("requests",), "e2e/conftest.py": ("requests",), - "e2e/claude_code/pr_gate_version_resolver.py": ("urllib.request",), + "e2e/cli_harness/npm_version_resolver.py": ("urllib.request",), } EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"}) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 1b1ee484afc..57ae9f1905f 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -138,6 +138,7 @@ ignored_function_names = [ "_resolve_pre_call_check_input_tokens", # Tested through the token-count failure pre-call check tests in test_router.py "_params_neither_dropped_nor_allowed", # Tested through the allowed-or-dropped param pre-call check tests in test_router.py "_deployment_rejects_params", # Tested through the unsupported-param pre-call check tests in test_router.py + "_cache_flags_left_to_the_catalog", # Tested through the wildcard prompt cache flag tests in test_anthropic_cache_control_hook.py "_deployment_lacks_params", # Tested through the unsupported-param pre-call check tests in test_router.py "_pre_call_check_requested_chat_params", # Tested through the unsupported-param pre-call check tests in test_router.py "_deployment_supports_param", # Tested through the deployment-capability pre-call check tests in test_router.py diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index cf1e41a59a6..7843c1a4091 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,11 +44,13 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests -- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture. Ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`). The CLI-driving path stays bespoke +- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture. Ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`). The CLI-driving path stays bespoke. `pr_gate_version_resolver.py` only forwards to `cli_harness/npm_version_resolver.py`, because the Buildkite runner images still run that path and install @latest without it +- `cli_harness/` - no tests, only the plumbing every suite that drives a real coding-agent CLI shares: `npm_version_resolver.py` picks the CLI release to install (the newest npm release at least three days old, run by CI as `python tests/e2e/cli_harness/npm_version_resolver.py `, stdlib only), `lockdown.py` gives the CLI subprocess an allowlisted env and a fresh empty HOME so it never sees the runner's credentials or dotfiles, and `rate_limiter.py` is the cross-process token bucket per upstream that keeps parallel CLI runs under one request rate. Covered by `tests/e2e_harness/cli_harness/`. Code only one CLI suite needs stays in that suite +- `codex/` - the Codex compatibility suite: drives the real `codex exec --json` CLI through the gateway, one class per row (the command round trip, a file edit with the patch tool, a long session that follows a chain of files one command at a time, and auto-compaction, where Codex reads a random value, compacts through the gateway and is asked for the value after the file is deleted), each row parametrized over `SETUPS` in `setups.py`, one `pytest.param` per setup (GPT on Bedrock Mantle, Bedrock Runtime, Azure OpenAI and OpenAI, plus Claude on Bedrock). A row with one setup broken by a known bug parametrizes over `with_expected_failure(, reason=...)`, which is `SETUPS` with a strict xfail on that one setup. The `setup` fixture registers that setup's deployment under a fresh name via `/model/new` and the `codex` fixture makes a key scoped to it, and both are deleted on teardown. `harness.py` gives each test a throwaway HOME and CODEX_HOME whose `config.toml` points a custom `responses` provider at the gateway and has Codex read its model list from the gateway's `/v1/models` (`model_catalog_url` on the provider), with the sandbox on (`workspace-write`) and approvals off, and runs the CLI under `cli_harness`'s env lockdown and rate limiter. `run.py` reads the JSONL events into a `CodexRun`: the commands with their output, the patched files, the last agent message as `answer`, the usage and the errors, where an error's `status_code` is what lets the suite's 5xx rerun rule catch a turn Codex gave up on. `codex exec --json` prints nothing when Codex compacts, so `rollout.py` reads the compactions from the thread's rollout file under CODEX_HOME/sessions, and a compaction is remote when its replacement history kept the provider's `compaction` item. `exec` and `resume` take `compact_after_tokens`, which forces Codex's remote compaction (`compaction_trigger` items over `/v1/responses`, the path Codex otherwise only takes for OpenAI and Azure) on the gateway provider and compacts once the context grows past that many tokens. A test passes on a random value the model can only learn through the gateway, never on the model's wording or the exact command it typed. It needs `codex` on PATH and fails without it. Covered by `tests/e2e_harness/codex/` - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream, see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport. The Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json`. Its specs follow the dashboard's rule against testing what shadcn and Base UI already guarantee, in `ui/litellm-dashboard/AGENTS.md` ## MCP suite: real Datadog only @@ -324,5 +326,6 @@ other... - spin up a local proxy by running the litellm proxy locally (`litellm --config .yml --port 4000`, see CONTRIBUTING.md), make sure all tests pass. if a test fails due to an internally found issue, let users know to create a linear ticket for it. - do not use xfail markers, tests should be written in a form that the end user expects it to pass +- the `codex/` suite is the one exception: a single setup that a LiteLLM or provider bug breaks may carry `pytest.mark.xfail(strict=True, reason=...)` with a plain one-sentence reason on that setup's `pytest.param`, so it keeps running and goes red once the bug is fixed. Strict only, and never on a flaky setup (one that passes some runs and fails others) - a cell addresses the stack through its front door (`PROXY_BASE_URL`) the way a customer does, never a gateway pod by address, and proves a cross-replica property with N independent calls through that door, naming the miss odds (2^-N at two pods) in its docstring. `PROXY_REPLICA_URLS` is for read-backs that are per pod by nature (polling until a management write has converged on every replica, `/metrics`, RSS), never for steering the scenario a cell asserts on at a chosen pod diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 6f7a7b39a39..599cfe6aa75 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -16,23 +16,18 @@ from __future__ import annotations import json import os import re -import shutil import subprocess import sys import tempfile import time from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field -from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Dict, Final, List, Mapping, Optional, Sequence, Tuple, Union +from cli_harness.lockdown import isolated_home, locked_down_env +from cli_harness.rate_limiter import RateBucket, RateLimiter, get_default_limiter from e2e_metadata import step -from claude_code.rate_limiter import ( - RateLimiter, - get_default_limiter, - infer_provider, -) - CLAUDE_CLI_DEFAULT = "claude" # 120s is fine for a single isolated CLI call against an unloaded # upstream, but the matrix run launches up to 75 concurrent calls and @@ -77,61 +72,36 @@ Anthropic / Bedrock / Vertex) to reset. Both knobs are env-tunable so a matrix run can trade wall time for resilience without code edits; retries=0 disables the behavior entirely.""" -# Env vars the `claude` Node CLI legitimately needs to function: -# locating its own binary + node, basic locale/terminal plumbing. -# Deliberately excludes every credential-bearing var that the -# surrounding CI job sets for the proxy (ANTHROPIC_API_KEY, AWS_*, -# AZURE_*, VERTEXAI_CREDENTIALS, GITHUB_TOKEN, OPENAI_API_KEY, -# DATABASE_URL, ...). The CLI talks to the proxy via the explicit -# ANTHROPIC_BASE_URL/ANTHROPIC_AUTH_TOKEN we set below — it has no -# business reading the proxy's upstream credentials, and a -# compromised CLI release shouldn't be able to exfiltrate them out -# of the CI environment. -# -# `HOME` is intentionally NOT in this list: see `_make_isolated_home` -# below. We give the CLI a fresh empty per-invocation HOME so a -# compromised claude package or a model-directed `Read` tool call -# can't reach files like `~/.config/gh/hosts.yml`, `~/.ssh/id_rsa`, -# or `~/.bash_history` on the cron VM (and on the CircleCI executor -# the same isolation prevents accidentally exposing checkout-adjacent -# files even though the runner home is ephemeral there). -_CLI_ENV_ALLOWLIST: tuple = ( - "PATH", - "USER", - "LOGNAME", - "SHELL", - "TERM", - "TMPDIR", - "LANG", - "LC_ALL", - "LC_CTYPE", - "NODE_PATH", - "NVM_DIR", - "NVM_BIN", -) +def infer_provider(model: str) -> RateBucket: + """Map a model alias to its compat-matrix rate bucket. -def _make_isolated_home() -> str: - """Create a fresh empty HOME directory for a single `claude` subprocess. - - The CLI needs *a* writable HOME (it caches per-session state under - `$HOME/.claude/projects//`), but it has no legitimate need - for the *user's* HOME. Handing it the real one means a compromised - `@anthropic-ai/claude-code` release, or a model-directed `Read` - tool call during a PDF/vision cell, can read host files like - `~/.config/gh/hosts.yml` (GitHub CLI host token), `~/.ssh/`, - `~/.bash_history`, or any other dotfile under the runtime user's - home. On the cron VM the runtime user is a real interactive - account (`mateo`) with a populated home directory, so this is a - real exfiltration surface. - - The directory is created under `tempfile.gettempdir()` (which is - `/tmp` on Linux; under systemd's `PrivateTmp=true` that's a - per-service tmpfs that the service user can't otherwise reach). - Caller is responsible for `shutil.rmtree`-ing it after the - subprocess exits. + The matrix column layout is fixed; aliases registered in the proxy + encode the provider via a suffix (`-bedrock-converse`, + `-bedrock-invoke`, `-azure`, `-vertex`, `-openai`, `-azure-openai`, + `-bedrock-mantle`) or its absence (Anthropic). Order matters: + `-azure-openai` also ends with `-openai`, and the bedrock suffixes + all contain `bedrock`, so the more-specific suffixes are tested + first. """ - return tempfile.mkdtemp(prefix="claude-cli-home-") + if not model: + raise ValueError("model must be a non-empty string") + lower: Final = model.lower() + if lower.endswith("-azure-openai"): + return RateBucket.AZURE_OPENAI + if lower.endswith("-openai"): + return RateBucket.OPENAI + if lower.endswith("-bedrock-mantle"): + return RateBucket.BEDROCK_MANTLE + if lower.endswith("-bedrock-converse"): + return RateBucket.BEDROCK_CONVERSE + if lower.endswith("-bedrock-invoke"): + return RateBucket.BEDROCK_INVOKE + if lower.endswith("-azure"): + return RateBucket.AZURE + if lower.endswith("-vertex"): + return RateBucket.VERTEX_AI + return RateBucket.ANTHROPIC _FIXED_CLI_USER_ID = "0" * 64 @@ -176,7 +146,7 @@ def _stable_cli_state() -> Tuple[str, str]: Pinning both makes the bodies repeatable across builds. This narrows what survives rather than widening it: HOME stays fresh and - empty per invocation, so the isolation `_make_isolated_home` describes is + empty per invocation, so the isolation `cli_harness.lockdown` describes is unchanged, and the CLI's own state no longer outlives the pod either. The working directory is deliberately not the checkout, so a model-directed `Read` sees an empty directory instead of the repository. @@ -291,38 +261,26 @@ def run_claude( cmd.append("--") cmd.append(prompt) - # Build a minimal env for the CLI subprocess: only the allowlisted - # process-runtime vars from os.environ, plus the explicit proxy - # creds, plus any caller-supplied overrides. See _CLI_ENV_ALLOWLIST - # above for the security rationale. - env: Dict[str, str] = { - key: os.environ[key] for key in _CLI_ENV_ALLOWLIST if key in os.environ - } - env["ANTHROPIC_BASE_URL"] = base_url - env["ANTHROPIC_AUTH_TOKEN"] = api_key - env["DISABLE_AUTOUPDATER"] = "1" - # Hand the CLI a fresh empty HOME so a compromised claude package - # or a model-directed Read tool call can't see the runtime user's - # real dotfiles. Created here, removed in the `finally` below - # regardless of how the subprocess exits. - isolated_home = _make_isolated_home() - env["HOME"] = isolated_home - config_dir, workspace = _stable_cli_state() - env["CLAUDE_CONFIG_DIR"] = config_dir - if extra_env: - env.update(extra_env) - # Throttle by provider *before* launching the CLI. Doing this here # (rather than per-test) means every code path that lands on # `run_claude` is rate-limited automatically — including # `run_claude_models_parallel`, which is the hot path during the # full matrix run. limiter = rate_limiter if rate_limiter is not None else get_default_limiter() - provider = infer_provider(model) - limiter.acquire(provider) - run_fn = runner or subprocess.run - try: + with isolated_home("claude-cli-home-") as home: + config_dir, workspace = _stable_cli_state() + env = locked_down_env( + home=home, + overrides={ + "ANTHROPIC_BASE_URL": base_url, + "ANTHROPIC_AUTH_TOKEN": api_key, + "DISABLE_AUTOUPDATER": "1", + "CLAUDE_CONFIG_DIR": config_dir, + **(extra_env or {}), + }, + ) + limiter.acquire(infer_provider(model)) try: completed = run_fn( cmd, @@ -340,15 +298,6 @@ def run_claude( ) from exc except subprocess.TimeoutExpired as exc: raise ClaudeCLIError(f"claude CLI timed out after {timeout}s") from exc - finally: - # Best-effort cleanup. If the subprocess wrote a `.claude/` - # session dir under the isolated HOME, we remove it here so - # parallel matrix runs don't accumulate per-call tmpdirs. - # `ignore_errors=True` because rmtree races with any - # not-yet-reaped child (SIGTERM'd `claude` on a host-side - # timeout) are benign — the next matrix run starts from a - # fresh tmpdir anyway. - shutil.rmtree(isolated_home, ignore_errors=True) events = _parse_stream_json(completed.stdout or "") text = _extract_assistant_text(events) diff --git a/tests/e2e/claude_code/cron_vm/README.md b/tests/e2e/claude_code/cron_vm/README.md index 4c16b6e237a..7451e86e09b 100644 --- a/tests/e2e/claude_code/cron_vm/README.md +++ b/tests/e2e/claude_code/cron_vm/README.md @@ -52,7 +52,7 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: the stable venv doesn't install) is never loaded. The tag's own `tests/e2e/` is deliberately not used. 3. **Resolves and installs the Claude Code CLI under test**: - `pr_gate_version_resolver.py` (run on the venv, the image has no + `tests/e2e/cli_harness/npm_version_resolver.py` (run on the venv, the image has no Python of its own) picks the newest `@anthropic-ai/claude-code` npm release published at least 3 days ago, the same buffer the PR gate uses, unless `CLAUDE_CODE_VERSION` pins one, and diff --git a/tests/e2e/claude_code/cron_vm/run_daily.sh b/tests/e2e/claude_code/cron_vm/run_daily.sh index 17b628b8c78..749d5166a34 100755 --- a/tests/e2e/claude_code/cron_vm/run_daily.sh +++ b/tests/e2e/claude_code/cron_vm/run_daily.sh @@ -8,7 +8,7 @@ # Releases API. # 2. Update a long-lived worktree at $WORKTREE to that tag and `uv sync` it. # 3. Resolve the Claude Code CLI version under test (the newest npm -# release at least 3 days old, via pr_gate_version_resolver.py, or +# release at least 3 days old, via cli_harness/npm_version_resolver.py, or # $CLAUDE_CODE_VERSION when set) and download that release's # linux-x64 binary into the run's scratch dir, checksum-verified # against the vendor's release manifest (install_claude_code.sh). @@ -329,7 +329,7 @@ PROXY_CONFIG="${WORKTREE}/tests/e2e/claude_code/test_config.yaml" if [[ -z "${CLAUDE_CODE_VERSION}" ]]; then CLAUDE_CODE_VERSION="$( cd "${WORKTREE}" \ - && "${WORKTREE_UV}" run --no-sync python "${POPULATOR_DIR}/../pr_gate_version_resolver.py" + && "${WORKTREE_UV}" run --no-sync python "${POPULATOR_DIR}/../../cli_harness/npm_version_resolver.py" @anthropic-ai/claude-code )" || die "could not resolve the Claude Code version to test" log "resolved claude code: ${CLAUDE_CODE_VERSION}" else diff --git a/tests/e2e/claude_code/http_probe.py b/tests/e2e/claude_code/http_probe.py index a6c587f98cf..390ef4c24d4 100644 --- a/tests/e2e/claude_code/http_probe.py +++ b/tests/e2e/claude_code/http_probe.py @@ -31,8 +31,7 @@ from __future__ import annotations from dataclasses import dataclass from typing import TYPE_CHECKING -from pydantic import BaseModel - +from cli_harness.rate_limiter import RateLimiter, get_default_limiter from e2e_http import ( NetworkError, RateLimitedError, @@ -59,12 +58,9 @@ from models import ( JsonSchemaProperty, ToolInputSchema, ) +from pydantic import BaseModel -from claude_code.rate_limiter import ( - RateLimiter, - get_default_limiter, - infer_provider, -) +from claude_code.cli_driver import infer_provider if TYPE_CHECKING: from proxy_client import ProxyClient diff --git a/tests/e2e/claude_code/pr_gate_version_resolver.py b/tests/e2e/claude_code/pr_gate_version_resolver.py index 756dafb7d73..72392f8a0ae 100644 --- a/tests/e2e/claude_code/pr_gate_version_resolver.py +++ b/tests/e2e/claude_code/pr_gate_version_resolver.py @@ -1,155 +1,16 @@ -"""Claude Code PR-Gate Version Resolver. - -Resolves the `@anthropic-ai/claude-code` npm version that the PR-gate CI -job installs. Selects the newest version (by publish timestamp) whose -publish timestamp is at least 3 days old. The 3-day window is a security -review buffer — see PRD #26476, "Version resolvers". - -Two surfaces: - -- ``resolve_pr_gate_version(...)`` — the importable function. Accepts - pre-fetched npm metadata (for unit tests) or a custom ``fetcher`` - callable. The default fetcher hits the public npm registry. -- ``python -m claude_code.pr_gate_version_resolver`` — prints the - resolved version string to stdout, suitable for piping into a shell - ``$(...)`` substitution inside the CircleCI job. - -The CLI form is what CircleCI runs at job start; engineers reading the -job log can see the selected version on a single line above the -``npm install -g`` step (acceptance criterion: "the selected Claude -Code version is logged in the CI output"). -""" +"""The path the Buildkite runner images run (`python3 tests/e2e/claude_code/pr_gate_version_resolver.py`) to pick the +Claude Code release they install. They fall back to @latest when this file is missing, so keep it until they call +`cli_harness/npm_version_resolver.py` themselves.""" from __future__ import annotations -import json +import subprocess import sys -import urllib.request -from datetime import datetime, timedelta, timezone -from typing import Callable, Mapping, Optional - -PACKAGE_NAME = "@anthropic-ai/claude-code" -NPM_REGISTRY_URL = "https://registry.npmjs.org/{package}" -DEFAULT_MIN_AGE = timedelta(days=3) -DEFAULT_FETCH_TIMEOUT_SECONDS = 30 - -# npm's `time` map mixes per-version timestamps with these meta keys. -_TIME_META_KEYS = frozenset({"created", "modified"}) - - -class NoEligibleVersionError(RuntimeError): - """Raised when no version in the npm metadata satisfies the min-age cutoff.""" - - -def _parse_npm_timestamp(value: str) -> datetime: - """Parse the ISO-8601 timestamps npm emits (always UTC, may use ``Z``).""" - if value.endswith("Z"): - value = value[:-1] + "+00:00" - parsed = datetime.fromisoformat(value) - if parsed.tzinfo is None: - parsed = parsed.replace(tzinfo=timezone.utc) - return parsed - - -def _default_fetcher(package_name: str) -> dict: - """Fetch the npm packument for ``package_name`` over HTTPS. - - Uses urllib (stdlib) so this module has no extra dependencies in the - CI environment. Returns the raw JSON dict. - """ - # urllib.parse.quote would encode the leading '@' / '/' which the - # npm registry expects literally; do a minimal hand-roll instead. - url = NPM_REGISTRY_URL.format(package=package_name.replace("/", "%2F")) - req = urllib.request.Request(url, headers={"Accept": "application/json"}) - with urllib.request.urlopen( # noqa: S310 — registry URL is constant - req, timeout=DEFAULT_FETCH_TIMEOUT_SECONDS - ) as response: - body = response.read().decode("utf-8") - return json.loads(body) - - -def resolve_pr_gate_version( - *, - metadata: Optional[Mapping] = None, - fetcher: Optional[Callable[[str], Mapping]] = None, - as_of: Optional[datetime] = None, - min_age: timedelta = DEFAULT_MIN_AGE, - package_name: str = PACKAGE_NAME, -) -> str: - """Return the newest npm version of ``package_name`` published >= ``min_age`` ago. - - "Newest" means newest by **publish time**, not semver string order — - if a patch lands on an older major after a newer release, the - patched line is the eligible one. A version npm has unpublished keeps - its ``time`` entry but drops out of ``versions``, so only versions - still present in ``versions`` are candidates. - - Args: - metadata: Pre-fetched npm packument (skips the HTTP call). Useful - for unit tests. - fetcher: Callable taking a package name and returning the - packument. Defaults to a stdlib HTTPS fetcher. - as_of: The clock used to decide whether a version is "old - enough". Defaults to ``datetime.now(timezone.utc)``. - min_age: Minimum publish age. Defaults to 3 days. - package_name: Defaults to ``@anthropic-ai/claude-code``. - - Raises: - NoEligibleVersionError: when no version in the registry meets - the age cutoff. - """ - if metadata is None: - fetch = fetcher or _default_fetcher - metadata = fetch(package_name) - - times = metadata.get("time") or {} - versions = metadata.get("versions") or {} - if as_of is None: - as_of = datetime.now(timezone.utc) - cutoff = as_of - min_age - - eligible: list[tuple[datetime, str]] = [] - for version, raw_ts in times.items(): - if version in _TIME_META_KEYS: - continue - if version not in versions: - continue - if not isinstance(raw_ts, str): - continue - if "-" in version: - continue - published = _parse_npm_timestamp(raw_ts) - if published <= cutoff: - eligible.append((published, version)) - - if not eligible: - raise NoEligibleVersionError( - f"no version of {package_name} is at least {min_age} old " - f"as of {as_of.isoformat()}" - ) - - eligible.sort(key=lambda pair: pair[0], reverse=True) - return eligible[0][1] - - -def _main(argv: list[str]) -> int: - """Print the resolved version to stdout. Exit code 0 on success. - - Stderr carries the human-readable announcement so the version can be - captured cleanly with ``$(python -m ...)`` in shell. - """ - try: - version = resolve_pr_gate_version() - except Exception as exc: # noqa: BLE001 — CLI surface, want everything - print(f"pr_gate_version_resolver: {exc}", file=sys.stderr) # noqa: T201 - return 1 - print( # noqa: T201 - f"pr_gate_version_resolver: selected {PACKAGE_NAME}@{version}", - file=sys.stderr, - ) - print(version) # noqa: T201 - return 0 +from pathlib import Path +from typing import Final +_RESOLVER: Final = Path(__file__).resolve().parent.parent / "cli_harness" / "npm_version_resolver.py" +_PACKAGE: Final = "@anthropic-ai/claude-code" if __name__ == "__main__": - raise SystemExit(_main(sys.argv[1:])) + raise SystemExit(subprocess.run((sys.executable, "-I", str(_RESOLVER), _PACKAGE), check=False).returncode) diff --git a/tests/e2e/claude_code/rate_limiter.py b/tests/e2e/claude_code/rate_limiter.py deleted file mode 100644 index 5818338ff7f..00000000000 --- a/tests/e2e/claude_code/rate_limiter.py +++ /dev/null @@ -1,378 +0,0 @@ -"""Cross-process token-bucket rate limiter for the Claude Code compat suite. - -The compat matrix runs 75 live `claude` CLI invocations (25 cells × 3 -Claude tiers per cell). When pytest-xdist fans these out across worker -processes, each worker would maintain its own in-memory rate limiter -and the *aggregate* request rate hitting any one upstream provider -would be `workers × per-worker rate` — exactly the situation that -trips Anthropic / Azure / Bedrock / Vertex 429s in the middle of a -matrix run and silently flips green cells red. - -The fix is a **shared, cross-process** token bucket per provider, -backed by a small JSON state file and an OS-level `flock`. Each -`run_claude` invocation acquires one token (sleeping if the bucket is -empty) before launching the CLI; refills happen lazily based on wall -time, so workers can be killed and restarted without losing or -double-spending budget. - -Configuration is driven entirely by environment variables so a -binary-search workflow can shift per-provider rates without code -edits: - - LITELLM_COMPAT_RATE_ANTHROPIC (req/s, default 5.0) - LITELLM_COMPAT_RATE_AZURE (req/s, default 5.0) - LITELLM_COMPAT_RATE_VERTEX_AI (req/s, default 5.0) - LITELLM_COMPAT_RATE_BEDROCK_CONVERSE (req/s, default 5.0) - LITELLM_COMPAT_RATE_BEDROCK_INVOKE (req/s, default 5.0) - LITELLM_COMPAT_RATE_OPENAI (req/s, default 5.0) - LITELLM_COMPAT_RATE_AZURE_OPENAI (req/s, default 5.0) - LITELLM_COMPAT_RATE_BEDROCK_MANTLE (req/s, default 5.0) - LITELLM_COMPAT_RATE_BURST (per-bucket burst override; - default = rate) - LITELLM_COMPAT_RATE_STATE_DIR (state file directory; - default = $TMPDIR/litellm-claude-compat-ratelimit) - -A rate of 0 (or any non-positive value) disables throttling for that -provider — useful when you trust the upstream to handle the burst or -when running unit-test-shaped workloads that never actually hit the -network. - -The provider id is inferred from the model id by `infer_provider`, -mirroring the matrix's column layout (`anthropic`, `azure`, -`vertex_ai`, `bedrock_converse`, `bedrock_invoke`, `openai`, -`azure_openai`, `bedrock_mantle`). The `vertex_ai_gpt` matrix column -has no bucket: its cells are static not_applicable stubs that never -reach the network. -""" - -from __future__ import annotations - -import contextlib -import json -import math -import os -import tempfile -import time -from dataclasses import dataclass -from pathlib import Path -from typing import Dict, Iterator, Mapping, Optional - -# `fcntl` is POSIX-only; the suite is Linux/macOS only, so we don't -# attempt a Windows fallback. Importing at module load fails fast on -# the (currently non-existent) Windows runner so we don't silently -# degrade to no-locking behavior. -import fcntl - -PROVIDER_ANTHROPIC = "anthropic" -PROVIDER_AZURE = "azure" -PROVIDER_VERTEX_AI = "vertex_ai" -PROVIDER_BEDROCK_CONVERSE = "bedrock_converse" -PROVIDER_BEDROCK_INVOKE = "bedrock_invoke" -PROVIDER_OPENAI = "openai" -PROVIDER_AZURE_OPENAI = "azure_openai" -PROVIDER_BEDROCK_MANTLE = "bedrock_mantle" - -ALL_PROVIDERS = ( - PROVIDER_ANTHROPIC, - PROVIDER_AZURE, - PROVIDER_VERTEX_AI, - PROVIDER_BEDROCK_CONVERSE, - PROVIDER_BEDROCK_INVOKE, - PROVIDER_OPENAI, - PROVIDER_AZURE_OPENAI, - PROVIDER_BEDROCK_MANTLE, -) - -DEFAULT_RATE = 5.0 # req/s per provider, conservative starting point -RATE_ENV_PREFIX = "LITELLM_COMPAT_RATE_" -BURST_ENV = "LITELLM_COMPAT_RATE_BURST" -STATE_DIR_ENV = "LITELLM_COMPAT_RATE_STATE_DIR" -DEFAULT_STATE_DIR_NAME = "litellm-claude-compat-ratelimit" - - -def infer_provider(model: str) -> str: - """Map a model alias to its compat-matrix provider id. - - The matrix column layout is fixed; aliases registered in the proxy - encode the provider via a suffix (`-bedrock-converse`, - `-bedrock-invoke`, `-azure`, `-vertex`, `-openai`, `-azure-openai`, - `-bedrock-mantle`) or its absence (Anthropic). Order matters: - `-azure-openai` also ends with `-openai`, and the bedrock suffixes - all contain `bedrock`, so the more-specific suffixes are tested - first. - """ - if not model: - raise ValueError("model must be a non-empty string") - lower = model.lower() - if lower.endswith("-azure-openai"): - return PROVIDER_AZURE_OPENAI - if lower.endswith("-openai"): - return PROVIDER_OPENAI - if lower.endswith("-bedrock-mantle"): - return PROVIDER_BEDROCK_MANTLE - if lower.endswith("-bedrock-converse"): - return PROVIDER_BEDROCK_CONVERSE - if lower.endswith("-bedrock-invoke"): - return PROVIDER_BEDROCK_INVOKE - if lower.endswith("-azure"): - return PROVIDER_AZURE - if lower.endswith("-vertex"): - return PROVIDER_VERTEX_AI - return PROVIDER_ANTHROPIC - - -@dataclass(frozen=True) -class ProviderConfig: - """Static config snapshot for one provider's bucket. - - Captured up front (rather than re-read per acquire) so the limiter's - behavior in a single process is stable even if env vars are mutated - mid-run. A fresh `RateLimiter` picks up env changes on construction. - """ - - rate_per_sec: float - burst: float - - @property - def enabled(self) -> bool: - return self.rate_per_sec > 0 and self.burst > 0 - - -def load_config( - env: Optional[Mapping[str, str]] = None, -) -> Dict[str, ProviderConfig]: - """Build a {provider: ProviderConfig} from env, applying defaults. - - Parsing failures fall back to the default rate rather than - crashing the test session — a typo in `LITELLM_COMPAT_RATE_AZURE` - should not silently disable throttling, but it also shouldn't - abort 75 live tests with a `ValueError` ten minutes in. - """ - src = env if env is not None else os.environ - - def _as_float(value: Optional[str], default: float) -> float: - if value is None or value == "": - return default - try: - return float(value) - except (TypeError, ValueError): - return default - - burst_override = _as_float(src.get(BURST_ENV), -1.0) - - out: Dict[str, ProviderConfig] = {} - for provider in ALL_PROVIDERS: - env_key = RATE_ENV_PREFIX + provider.upper() - rate = _as_float(src.get(env_key), DEFAULT_RATE) - burst = burst_override if burst_override > 0 else max(rate, 1.0) - out[provider] = ProviderConfig(rate_per_sec=rate, burst=burst) - return out - - -def _state_dir(env: Optional[Mapping[str, str]] = None) -> Path: - """Resolve the directory holding per-provider state files. - - A user-supplied `LITELLM_COMPAT_RATE_STATE_DIR` wins for tests and - container environments where `$TMPDIR` may be ephemeral or shared - in surprising ways. The directory is created lazily with - `parents=True, exist_ok=True` so first-run setup needs no - fixture wiring. - """ - src = env if env is not None else os.environ - explicit = src.get(STATE_DIR_ENV) - if explicit: - return Path(explicit) - return Path(tempfile.gettempdir()) / DEFAULT_STATE_DIR_NAME - - -class RateLimiter: - """Cross-process token bucket per provider. - - Each `acquire(provider)` call: - 1. Opens (creating if needed) `/.json`. - 2. Holds an exclusive `flock` while reading + updating the - {tokens, last_refill} state. - 3. Refills tokens based on `now - last_refill`, capped at burst. - 4. If tokens >= 1, subtracts one and returns immediately. - 5. Otherwise, computes the wall-time delay needed to earn a - single token at the configured rate, releases the lock, and - sleeps. After sleeping it retries — staying under the lock - while sleeping would serialize all workers behind whichever - one held it longest. - - This bounds the *aggregate* req/s seen by the upstream, regardless - of how many xdist workers, threads, or processes are concurrently - running tests against the same provider. - - A `_clock` / `_sleep` injection seam keeps the unit tests fast and - deterministic; production callers should never override either. - """ - - def __init__( - self, - config: Optional[Mapping[str, ProviderConfig]] = None, - state_dir: Optional[Path] = None, - clock: Optional[callable] = None, - sleep: Optional[callable] = None, - ) -> None: - self._config = dict(config) if config is not None else load_config() - self._state_dir = Path(state_dir) if state_dir is not None else _state_dir() - self._clock = clock or time.monotonic - self._sleep = sleep or time.sleep - # `_state_dir.mkdir` once on construction is fine; concurrent - # workers all racing to create the same directory is benign. - self._state_dir.mkdir(parents=True, exist_ok=True) - - def acquire(self, provider: str) -> float: - """Block until one token is available for `provider`. - - Returns the cumulative wall-time spent waiting (0.0 when the - bucket had budget and we returned immediately). Callers can - log this to attribute slow cells to throttling vs. upstream - latency — same role `DriverResult.duration_ms` plays for the - actual CLI invocation. - """ - cfg = self._config.get(provider) - if cfg is None or not cfg.enabled: - return 0.0 - - path = self._state_path(provider) - total_waited = 0.0 - while True: - now = self._clock() - sleep_for = self._try_consume(path, cfg, now) - if sleep_for <= 0: - return total_waited - self._sleep(sleep_for) - total_waited += sleep_for - - def _state_path(self, provider: str) -> Path: - return self._state_dir / f"{provider}.json" - - def _try_consume(self, path: Path, cfg: ProviderConfig, now: float) -> float: - """Atomically refill and try to take one token. - - Returns 0.0 if a token was consumed, or a positive sleep - duration (seconds) if the caller must wait before retrying. - - We hold an exclusive `flock` only across the read-modify-write - of the JSON state — never across a `sleep` — so workers don't - serialize while one of them is parked. - """ - # `os.open` + `os.O_CREAT | os.O_RDWR` gives us a fd we can - # both lock and read/write through. Opening with `"a+"` then - # seeking is equivalent but uglier; this version is closer to - # the canonical flock recipe. - fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o600) - try: - fcntl.flock(fd, fcntl.LOCK_EX) - try: - tokens, last_refill = self._read_state(fd, cfg, now) - tokens, last_refill = self._refill(tokens, last_refill, now, cfg) - if tokens >= 1.0: - tokens -= 1.0 - self._write_state(fd, tokens, last_refill) - return 0.0 - # Not enough budget. Persist the refilled state so a - # subsequent caller doesn't have to redo the math, then - # release the lock and tell the caller how long to - # sleep before retrying. - self._write_state(fd, tokens, last_refill) - deficit = 1.0 - tokens - # `deficit / rate` seconds will earn exactly enough - # for one token. Add a tiny safety margin so we don't - # wake up nanoseconds early and spin. - return (deficit / cfg.rate_per_sec) + 1e-3 - finally: - fcntl.flock(fd, fcntl.LOCK_UN) - finally: - os.close(fd) - - @staticmethod - def _read_state(fd: int, cfg: ProviderConfig, now: float) -> tuple: - """Read {tokens, last_refill} from `fd`, defaulting to a full - bucket on a missing/empty/corrupt file. - - New files start full so the very first request never waits; - corrupt files are treated like new files because the - alternative — refusing to run — is worse than briefly - over-spending one bucket's worth of budget. - """ - os.lseek(fd, 0, os.SEEK_SET) - raw = os.read(fd, 4096).decode("utf-8") - if not raw.strip(): - return cfg.burst, now - try: - obj = json.loads(raw) - tokens = float(obj.get("tokens", cfg.burst)) - last_refill = float(obj.get("last_refill", now)) - return tokens, last_refill - except (ValueError, TypeError): - return cfg.burst, now - - @staticmethod - def _refill( - tokens: float, last_refill: float, now: float, cfg: ProviderConfig - ) -> tuple: - """Apply elapsed time to the bucket, capped at burst. - - Negative elapsed (clock went backward, e.g. across host - sleep/resume or a manually-tweaked monotonic mock) is clamped - to zero so we never *remove* tokens. - """ - elapsed = max(0.0, now - last_refill) - tokens = min(cfg.burst, tokens + elapsed * cfg.rate_per_sec) - return tokens, now - - @staticmethod - def _write_state(fd: int, tokens: float, last_refill: float) -> None: - payload = json.dumps({"tokens": tokens, "last_refill": last_refill}).encode( - "utf-8" - ) - os.lseek(fd, 0, os.SEEK_SET) - os.ftruncate(fd, 0) - os.write(fd, payload) - - -# A single process-wide instance is all we need: provider-keyed state -# is stored in files, so multiple `RateLimiter` instances would just -# duplicate the in-process bookkeeping. We expose a getter rather than -# the instance directly so unit tests can install a custom limiter -# scoped to a tmp directory without monkeypatching globals. -_default: Optional[RateLimiter] = None - - -def get_default_limiter() -> RateLimiter: - global _default - if _default is None: - _default = RateLimiter() - return _default - - -def reset_default_limiter() -> None: - """Drop the cached default limiter; the next `get_default_limiter` - call rebuilds it from the current environment. - - Useful between unit tests that patch env vars: without this they'd - keep reading the stale config snapshot from the first call. - """ - global _default - _default = None - - -@contextlib.contextmanager -def use_limiter(limiter: RateLimiter) -> Iterator[RateLimiter]: - """Temporarily install `limiter` as the process default. - - The driver's `run_claude` calls `get_default_limiter()`; tests that - want a controlled tmp-dir-backed limiter use this contextmanager - to swap one in without touching env vars or the on-disk default - state. - """ - global _default - previous = _default - _default = limiter - try: - yield limiter - finally: - _default = previous diff --git a/tests/e2e/claude_code/run_compat.sh b/tests/e2e/claude_code/run_compat.sh index c42416283e7..e3410391661 100755 --- a/tests/e2e/claude_code/run_compat.sh +++ b/tests/e2e/claude_code/run_compat.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash # Run the Claude Code compat matrix end-to-end against a live LiteLLM # proxy, with per-provider rate limits applied via the cross-process -# token bucket in `tests/e2e/claude_code/rate_limiter.py`. +# token bucket in `tests/e2e/cli_harness/rate_limiter.py`. # # Designed for binary-searching the ideal X / Y / Z req/s per provider: # 1. Pick an initial rate (e.g. 5/s for everyone). diff --git a/tests/e2e/cli_harness/__init__.py b/tests/e2e/cli_harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/cli_harness/lockdown.py b/tests/e2e/cli_harness/lockdown.py new file mode 100644 index 00000000000..1d430529851 --- /dev/null +++ b/tests/e2e/cli_harness/lockdown.py @@ -0,0 +1,49 @@ +"""What a coding-agent CLI under test may see of the host. + +The CI job holds the proxy's upstream credentials (provider keys, AWS, Azure, the database URL, GitHub tokens) in its +environment, and the cron VM runs as a real user with a populated home. A compromised CLI release, or a model told to +read a file, must reach neither, so the CLI gets only the allowlisted runtime variables plus what its suite passes in, +and a fresh empty HOME that is deleted when it exits. +""" + +from __future__ import annotations + +import os +import shutil +import tempfile +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from pathlib import Path +from types import MappingProxyType +from typing import Final + +CLI_ENV_ALLOWLIST: Final = ( + "PATH", + "USER", + "LOGNAME", + "SHELL", + "TERM", + "TMPDIR", + "LANG", + "LC_ALL", + "LC_CTYPE", + "NODE_PATH", + "NVM_DIR", + "NVM_BIN", +) + + +@contextmanager +def isolated_home(prefix: str) -> Generator[Path]: + home: Final = Path(tempfile.mkdtemp(prefix=prefix)) + try: + yield home + finally: + shutil.rmtree(home, ignore_errors=True) + + +def locked_down_env( + *, home: Path, overrides: Mapping[str, str], source: Mapping[str, str] = os.environ +) -> Mapping[str, str]: + allowed: Final = {name: source[name] for name in CLI_ENV_ALLOWLIST if name in source} + return MappingProxyType({**allowed, "HOME": str(home), **overrides}) diff --git a/tests/e2e/cli_harness/npm_version_resolver.py b/tests/e2e/cli_harness/npm_version_resolver.py new file mode 100644 index 00000000000..3c3f9b89ce5 --- /dev/null +++ b/tests/e2e/cli_harness/npm_version_resolver.py @@ -0,0 +1,92 @@ +"""The npm release of a coding-agent CLI that CI installs: the newest one published at least three days ago. + +The three days are a security review buffer (PRD #26476, "Version resolvers"), so a release that gets pulled or +patched within days never drives a run. Newest means by publish time, not semver order. + +`python tests/e2e/cli_harness/npm_version_resolver.py ` prints the version on stdout for a shell +`$(...)`. It is stdlib only because the cron image runs it without the project's dependencies. +""" + +from __future__ import annotations + +import json +import sys +import urllib.parse +import urllib.request +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import BinaryIO, Final, cast + +MIN_AGE: Final = timedelta(days=3) +_REGISTRY_URL: Final = "https://registry.npmjs.org/{package}" +_FETCH_TIMEOUT_SECONDS: Final = 30 + + +@dataclass(frozen=True, slots=True) +class Packument: + published_at: Mapping[str, datetime] + versions: frozenset[str] + + @classmethod + def from_json(cls, raw: object) -> Packument: + document: Final = _as_mapping(raw) + times: Final = _as_mapping(document.get("time")) + return cls( + published_at=MappingProxyType({version: _parse_utc(stamp) for version, stamp in _string_values(times)}), + versions=frozenset(_as_mapping(document.get("versions"))), + ) + + +def newest_version_at_least(packument: Packument, *, min_age: timedelta, as_of: datetime) -> str | None: + cutoff: Final = as_of - min_age + eligible: Final = ( + (published, version) + for version, published in packument.published_at.items() + if version in packument.versions and "-" not in version and published <= cutoff + ) + newest: Final = max(eligible, default=None) + return None if newest is None else newest[1] + + +def fetch_packument(package: str) -> Packument: + url: Final = _REGISTRY_URL.format(package=urllib.parse.quote(package, safe="@")) + request: Final = urllib.request.Request(url, headers={"Accept": "application/json"}) + with cast(BinaryIO, urllib.request.urlopen(request, timeout=_FETCH_TIMEOUT_SECONDS)) as response: + return Packument.from_json(cast(object, json.loads(response.read()))) + + +def _as_mapping(raw: object) -> Mapping[str, object]: + if not isinstance(raw, dict): + return {} + return {str(key): value for key, value in cast(dict[object, object], raw).items()} + + +def _string_values(raw: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + return tuple((key, value) for key, value in raw.items() if isinstance(value, str)) + + +def _parse_utc(stamp: str) -> datetime: + parsed: Final = datetime.fromisoformat(stamp.replace("Z", "+00:00")) + return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc) + + +def _main(argv: list[str]) -> int: + if len(argv) != 1: + print("usage: npm_version_resolver.py ", file=sys.stderr) # noqa: T201 # CLI surface + return 2 + package: Final = argv[0] + version: Final = newest_version_at_least( + fetch_packument(package), min_age=MIN_AGE, as_of=datetime.now(timezone.utc) + ) + if version is None: + print(f"npm_version_resolver: no release of {package} is at least {MIN_AGE} old", file=sys.stderr) # noqa: T201 # CLI surface + return 1 + print(f"npm_version_resolver: selected {package}@{version}", file=sys.stderr) # noqa: T201 # CLI surface + print(version) # noqa: T201 # CLI surface + return 0 + + +if __name__ == "__main__": + raise SystemExit(_main(sys.argv[1:])) diff --git a/tests/e2e/cli_harness/rate_limiter.py b/tests/e2e/cli_harness/rate_limiter.py new file mode 100644 index 00000000000..bd8119ffe87 --- /dev/null +++ b/tests/e2e/cli_harness/rate_limiter.py @@ -0,0 +1,151 @@ +"""A token bucket per upstream, shared by every process on the host, so parallel CLI suites and xdist workers +together stay under one request rate per provider instead of one rate each. + +Each bucket is a small JSON file read and written under an exclusive `flock`, in `LITELLM_COMPAT_RATE_STATE_DIR` +(default: a fixed folder under `$TMPDIR`). `LITELLM_COMPAT_RATE_` sets a bucket's req/s (default 5, 0 or less +turns it off) and `LITELLM_COMPAT_RATE_BURST` overrides every bucket's burst (default: its rate). +""" + +from __future__ import annotations + +import fcntl +import os +import tempfile +import time +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field +from enum import StrEnum +from functools import cache +from pathlib import Path +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict, ValidationError + +DEFAULT_RATE: Final = 5.0 +RATE_ENV_PREFIX: Final = "LITELLM_COMPAT_RATE_" +BURST_ENV: Final = "LITELLM_COMPAT_RATE_BURST" +STATE_DIR_ENV: Final = "LITELLM_COMPAT_RATE_STATE_DIR" +DEFAULT_STATE_DIR_NAME: Final = "litellm-claude-compat-ratelimit" +_WAKE_MARGIN_SECONDS: Final = 1e-3 + + +class RateBucket(StrEnum): + ANTHROPIC = "anthropic" + AZURE = "azure" + VERTEX_AI = "vertex_ai" + BEDROCK_CONVERSE = "bedrock_converse" + BEDROCK_INVOKE = "bedrock_invoke" + OPENAI = "openai" + AZURE_OPENAI = "azure_openai" + BEDROCK_MANTLE = "bedrock_mantle" + + +@dataclass(frozen=True, slots=True) +class ProviderConfig: + rate_per_sec: float + burst: float + + @property + def enabled(self) -> bool: + return self.rate_per_sec > 0 and self.burst > 0 + + +class _BucketState(BaseModel): + model_config = ConfigDict(frozen=True) + + tokens: float + last_refill: float + + +def load_config(env: Mapping[str, str] = os.environ) -> Mapping[RateBucket, ProviderConfig]: + burst_override: Final = _as_float(env.get(BURST_ENV), -1.0) + return MappingProxyType( + { + bucket: _bucket_config(_as_float(env.get(_rate_env(bucket)), DEFAULT_RATE), burst_override) + for bucket in RateBucket + } + ) + + +def default_state_dir(env: Mapping[str, str] = os.environ) -> Path: + explicit: Final = env.get(STATE_DIR_ENV) + return Path(explicit) if explicit else Path(tempfile.gettempdir()) / DEFAULT_STATE_DIR_NAME + + +@dataclass(frozen=True, slots=True) +class RateLimiter: + config: Mapping[RateBucket, ProviderConfig] = field(default_factory=load_config) + state_dir: Path = field(default_factory=default_state_dir) + clock: Callable[[], float] = time.monotonic + sleep: Callable[[float], None] = time.sleep + + def __post_init__(self) -> None: + self.state_dir.mkdir(parents=True, exist_ok=True) + + def acquire(self, bucket: RateBucket) -> float: + cfg: Final = self.config.get(bucket) + if cfg is None or not cfg.enabled: + return 0.0 + return sum(self._waits(self.state_dir / f"{bucket.value}.json", cfg)) + + def _waits(self, path: Path, cfg: ProviderConfig) -> Iterator[float]: + while (wait := _try_consume(path, cfg, self.clock())) > 0: + self.sleep(wait) + yield wait + + +@cache +def get_default_limiter() -> RateLimiter: + return RateLimiter() + + +def _rate_env(bucket: RateBucket) -> str: + return RATE_ENV_PREFIX + bucket.value.upper() + + +def _bucket_config(rate: float, burst_override: float) -> ProviderConfig: + return ProviderConfig(rate_per_sec=rate, burst=burst_override if burst_override > 0 else max(rate, 1.0)) + + +def _as_float(value: str | None, default: float) -> float: + if not value: + return default + try: + return float(value) + except ValueError: + return default + + +def _try_consume(path: Path, cfg: ProviderConfig, now: float) -> float: + fd: Final = os.open(path, os.O_CREAT | os.O_RDWR, 0o600) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + state: Final = _refilled(_read_state(fd, cfg, now), now, cfg) + if state.tokens >= 1.0: + _write_state(fd, _BucketState(tokens=state.tokens - 1.0, last_refill=state.last_refill)) + return 0.0 + _write_state(fd, state) + return (1.0 - state.tokens) / cfg.rate_per_sec + _WAKE_MARGIN_SECONDS + finally: + os.close(fd) + + +def _read_state(fd: int, cfg: ProviderConfig, now: float) -> _BucketState: + os.lseek(fd, 0, os.SEEK_SET) + raw: Final = os.read(fd, 4096) + try: + return _BucketState.model_validate_json(raw) + except ValidationError: + return _BucketState(tokens=cfg.burst, last_refill=now) + + +def _refilled(state: _BucketState, now: float, cfg: ProviderConfig) -> _BucketState: + elapsed: Final = max(0.0, now - state.last_refill) + return _BucketState(tokens=min(cfg.burst, state.tokens + elapsed * cfg.rate_per_sec), last_refill=now) + + +def _write_state(fd: int, state: _BucketState) -> None: + os.lseek(fd, 0, os.SEEK_SET) + os.ftruncate(fd, 0) + os.write(fd, state.model_dump_json().encode("utf-8")) diff --git a/tests/e2e/codex/__init__.py b/tests/e2e/codex/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/codex/conftest.py b/tests/e2e/codex/conftest.py new file mode 100644 index 00000000000..1ce6f2009a5 --- /dev/null +++ b/tests/e2e/codex/conftest.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from cli_harness.rate_limiter import get_default_limiter +from codex.harness import CodexHarness, codex_harness +from codex.setups import Deployment, ModelSetup +from e2e_config import PROXY_BASE_URL, unique_marker +from lifecycle import ResourceManager +from proxy_client import ProxyClient + + +@dataclass(frozen=True, slots=True) +class _Client: + proxy: ProxyClient + + +@pytest.fixture(scope="session") +def client(proxy: ProxyClient) -> _Client: + return _Client(proxy) + + +@pytest.fixture +def setup(deployment: Deployment, proxy: ProxyClient, resources: ResourceManager) -> ModelSetup: + model: Final = f"e2e-codex-{unique_marker()}" + model_id: Final = proxy.create_model(model, deployment.litellm_params, provider_live=True) + resources.defer(lambda: proxy.delete_model(model_id)) + return ModelSetup(model=model, rate_bucket=deployment.rate_bucket) + + +@pytest.fixture +def codex(setup: ModelSetup, resources: ResourceManager) -> Iterator[CodexHarness]: + with codex_harness( + gateway=PROXY_BASE_URL, + api_key=resources.key(models=[setup.model], user_id=None), + rate_bucket=setup.rate_bucket, + limiter=get_default_limiter(), + ) as harness: + yield harness diff --git a/tests/e2e/codex/harness.py b/tests/e2e/codex/harness.py new file mode 100644 index 00000000000..f6dd3700149 --- /dev/null +++ b/tests/e2e/codex/harness.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import subprocess +import tempfile +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from itertools import chain +from pathlib import Path +from typing import Final + +from cli_harness.lockdown import isolated_home, locked_down_env +from cli_harness.rate_limiter import RateBucket, RateLimiter +from codex.rollout import Compaction, rollout_compactions +from codex.run import CodexRun, codex_run +from e2e_metadata import step + +CODEX_CLI: Final = "codex" +EXEC_TIMEOUT_SECONDS: Final = 600.0 +_API_KEY_ENV: Final = "LITELLM_API_KEY" +_PROVIDER: Final = "litellm" +_CONFIG_TOML: Final = """\ +model_provider = "{provider}" +approval_policy = "never" +sandbox_mode = "workspace-write" +check_for_update_on_startup = false + +[analytics] +enabled = false + +[model_providers.{provider}] +name = "LiteLLM" +base_url = "{gateway}/v1" +env_key = "{api_key_env}" +wire_api = "responses" +model_catalog_url = "{gateway}/v1/models" +""" + + +@dataclass(frozen=True, slots=True) +class CodexHarness: + home: Path + api_key: str = field(repr=False) + rate_bucket: RateBucket + limiter: RateLimiter + cli: str = CODEX_CLI + timeout_seconds: float = EXEC_TIMEOUT_SECONDS + + @property + def codex_home(self) -> Path: + return self.home / ".codex" + + @step("Create a workspace folder holding the test's files") + def new_workspace(self, files: Mapping[str, str]) -> Path: + workspace: Final = Path(tempfile.mkdtemp(prefix="workspace-", dir=self.home)).resolve() + for name, content in files.items(): + _ = (workspace / name).write_text(content) + return workspace + + @step('Run Codex on {model} through the proxy with the prompt "{prompt}"') + def exec(self, prompt: str, *, model: str, cwd: Path, compact_after_tokens: int | None = None) -> CodexRun: + return self._run(("exec",), (*_compaction_flags(compact_after_tokens), "--model", model, "--", prompt), cwd) + + @step('Resume Codex thread {thread_id} on {model} through the proxy with the prompt "{prompt}"') + def resume( + self, thread_id: str, prompt: str, *, model: str, cwd: Path, compact_after_tokens: int | None = None + ) -> CodexRun: + return self._run( + ("exec", "resume"), + (*_compaction_flags(compact_after_tokens), "--model", model, "--", thread_id, prompt), + cwd, + ) + + def compactions(self, thread_id: str) -> tuple[Compaction, ...]: + rollouts: Final = sorted(self.codex_home.glob(f"sessions/**/rollout-*-{thread_id}.jsonl")) + return tuple(chain.from_iterable(rollout_compactions(rollout.read_text()) for rollout in rollouts)) + + def _run(self, subcommand: tuple[str, ...], args: tuple[str, ...], cwd: Path) -> CodexRun: + _ = self.limiter.acquire(self.rate_bucket) + completed: Final = subprocess.run( + (self.cli, *subcommand, "--json", "--skip-git-repo-check", *args), + cwd=cwd, + env=locked_down_env( + home=self.home, overrides={"CODEX_HOME": str(self.codex_home), _API_KEY_ENV: self.api_key} + ), + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=self.timeout_seconds, + check=False, + ) + return codex_run(exit_code=completed.returncode, stdout=completed.stdout, stderr=completed.stderr) + + +def _compaction_flags(after_tokens: int | None) -> tuple[str, ...]: + match after_tokens: + case None: + return () + case int(): + return ( + "--config", + f'model_providers.{_PROVIDER}.capabilities.remote_compaction="v2"', + "--config", + 'model_auto_compact_token_limit_scope="body_after_prefix"', + "--config", + f"model_auto_compact_token_limit={after_tokens}", + ) + + +@step("Give Codex a fresh empty home whose config sends every request to the proxy") +@contextmanager +def codex_harness( + *, gateway: str, api_key: str, rate_bucket: RateBucket, limiter: RateLimiter +) -> Generator[CodexHarness]: + with isolated_home("codex-cli-home-") as home: + harness: Final = CodexHarness(home=home, api_key=api_key, rate_bucket=rate_bucket, limiter=limiter) + harness.codex_home.mkdir() + _ = (harness.codex_home / "config.toml").write_text( + _CONFIG_TOML.format(provider=_PROVIDER, gateway=gateway, api_key_env=_API_KEY_ENV) + ) + yield harness diff --git a/tests/e2e/codex/rollout.py b/tests/e2e/codex/rollout.py new file mode 100644 index 00000000000..fcaeaa7ca42 --- /dev/null +++ b/tests/e2e/codex/rollout.py @@ -0,0 +1,49 @@ +"""The compactions Codex recorded in a thread's rollout file, `CODEX_HOME/sessions/**/rollout-*-.jsonl`. + +`codex exec --json` prints no event for a compaction, so the rollout is the only place it shows. Each line is +`{"timestamp", "type", "payload"}` (`codex-rs/history/src/rollout_payload.rs`) and only `compacted` lines are read. +A remote compaction is one whose replacement history carries the provider's `compaction` item, while a local one +replaces the history with a summary Codex wrote itself. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +from pydantic import BaseModel, ConfigDict + +_COMPACTED: Final = "compacted" +_PROVIDER_COMPACTION_ITEM: Final = "compaction" + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _Line(_Frozen): + type: str + payload: object = None + + +class _HistoryItem(_Frozen): + type: str + + +class _CompactedPayload(_Frozen): + replacement_history: tuple[_HistoryItem, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class Compaction: + remote: bool + + +def rollout_compactions(rollout: str) -> tuple[Compaction, ...]: + lines: Final = tuple(_Line.model_validate_json(line) for line in rollout.splitlines() if line.strip()) + return tuple(_compaction(line.payload) for line in lines if line.type == _COMPACTED) + + +def _compaction(payload: object) -> Compaction: + history: Final = _CompactedPayload.model_validate(payload).replacement_history or () + return Compaction(remote=any(item.type == _PROVIDER_COMPACTION_ITEM for item in history)) diff --git a/tests/e2e/codex/run.py b/tests/e2e/codex/run.py new file mode 100644 index 00000000000..5195da67494 --- /dev/null +++ b/tests/e2e/codex/run.py @@ -0,0 +1,187 @@ +"""What one `codex exec --json` run did, read from the JSONL events it printed. + +The event shapes are Codex's own (`codex-rs/exec/src/exec_events.rs`). We take the latest codex release > 3 days old, so an event or item +type this module does not model is skipped, while one it does model is validated strictly and fails loudly when +Codex changes it. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from typing import Annotated, Final, Literal, cast + +from pydantic import BaseModel, ConfigDict, Discriminator, Tag, TypeAdapter + +_HTTP_STATUS: Final = re.compile(r"\bstatus:? (\d{3})\b") + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class TokenUsage(_Frozen): + input_tokens: int + cached_input_tokens: int + cache_write_input_tokens: int = 0 + output_tokens: int + reasoning_output_tokens: int + + +class _AgentMessage(_Frozen): + type: Literal["agent_message"] + text: str + + +class _CommandExecution(_Frozen): + type: Literal["command_execution"] + command: str + aggregated_output: str + exit_code: int | None + + +class _FileUpdate(_Frozen): + path: str + kind: Literal["add", "delete", "update"] + + +class _FileChange(_Frozen): + type: Literal["file_change"] + changes: tuple[_FileUpdate, ...] + status: Literal["in_progress", "completed", "failed"] + + +class _ErrorItem(_Frozen): + type: Literal["error"] + message: str + + +class _Unmodelled(_Frozen): + type: str + + +class _Message(_Frozen): + message: str + + +def _known_or_unmodelled(known: frozenset[str]) -> Callable[[object], str]: + def tag(raw: object) -> str: + kind: Final = cast("Mapping[str, object]", raw).get("type") if isinstance(raw, Mapping) else None + return kind if isinstance(kind, str) and kind in known else "unmodelled" + + return tag + + +_Item = Annotated[ + Annotated[_AgentMessage, Tag("agent_message")] + | Annotated[_CommandExecution, Tag("command_execution")] + | Annotated[_FileChange, Tag("file_change")] + | Annotated[_ErrorItem, Tag("error")] + | Annotated[_Unmodelled, Tag("unmodelled")], + Discriminator(_known_or_unmodelled(frozenset({"agent_message", "command_execution", "file_change", "error"}))), +] + + +class _ThreadStarted(_Frozen): + type: Literal["thread.started"] + thread_id: str + + +class _ItemCompleted(_Frozen): + type: Literal["item.completed"] + item: _Item + + +class _TurnCompleted(_Frozen): + type: Literal["turn.completed"] + usage: TokenUsage + + +class _TurnFailed(_Frozen): + type: Literal["turn.failed"] + error: _Message + + +class _StreamError(_Frozen): + type: Literal["error"] + message: str + + +_Event = Annotated[ + Annotated[_ThreadStarted, Tag("thread.started")] + | Annotated[_ItemCompleted, Tag("item.completed")] + | Annotated[_TurnCompleted, Tag("turn.completed")] + | Annotated[_TurnFailed, Tag("turn.failed")] + | Annotated[_StreamError, Tag("error")] + | Annotated[_Unmodelled, Tag("unmodelled")], + Discriminator( + _known_or_unmodelled(frozenset({"thread.started", "item.completed", "turn.completed", "turn.failed", "error"})) + ), +] +_EVENT: Final = TypeAdapter[_Event](_Event) + + +@dataclass(frozen=True, slots=True) +class Command: + command: str + output: str + exit_code: int | None + + +@dataclass(frozen=True, slots=True) +class CodexError: + message: str + status_code: int | None + + +@dataclass(frozen=True, slots=True) +class CodexRun: + exit_code: int + thread_id: str | None + commands: tuple[Command, ...] + patched_files: tuple[Path, ...] + answer: str + usage: TokenUsage | None + errors: tuple[CodexError, ...] + stderr: str + + +def codex_run(*, exit_code: int, stdout: str, stderr: str) -> CodexRun: + events: Final = tuple(_EVENT.validate_json(line) for line in stdout.splitlines() if line.strip()) + items: Final = tuple(event.item for event in events if isinstance(event, _ItemCompleted)) + patches: Final = tuple(item for item in items if isinstance(item, _FileChange) and item.status == "completed") + answers: Final = tuple(item.text for item in items if isinstance(item, _AgentMessage)) + usages: Final = tuple(event.usage for event in events if isinstance(event, _TurnCompleted)) + return CodexRun( + exit_code=exit_code, + thread_id=next((event.thread_id for event in events if isinstance(event, _ThreadStarted)), None), + commands=tuple( + Command(command=item.command, output=item.aggregated_output, exit_code=item.exit_code) + for item in items + if isinstance(item, _CommandExecution) + ), + patched_files=tuple(Path(change.path) for change in chain.from_iterable(patch.changes for patch in patches)), + answer=answers[-1] if answers else "", + usage=usages[-1] if usages else None, + errors=tuple(_errors(events)), + stderr=stderr, + ) + + +def _errors(events: tuple[_Event, ...]) -> Iterator[CodexError]: + for event in events: + match event: + case _TurnFailed(error=_Message(message=message)) | _StreamError(message=message): + yield _codex_error(message) + case _ItemCompleted(item=_ErrorItem(message=message)): + yield _codex_error(message) + case _: + pass + + +def _codex_error(message: str) -> CodexError: + status: Final = _HTTP_STATUS.search(message) + return CodexError(message=message, status_code=int(status.group(1)) if status else None) diff --git a/tests/e2e/codex/setups.py b/tests/e2e/codex/setups.py new file mode 100644 index 00000000000..4716706c51c --- /dev/null +++ b/tests/e2e/codex/setups.py @@ -0,0 +1,96 @@ +"""The credentials are `os.environ/` references the gateway resolves, the same names the Claude Code matrix's +`test_config.yaml` uses, and Bedrock rides the gateway's default AWS identity.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final, Literal + +import pytest +from _pytest.mark.structures import ParameterSet +from cli_harness.rate_limiter import RateBucket +from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta +from models import LiteLLMParamsBody + +GPT_ON_BEDROCK_MANTLE: Final = "bedrock_mantle/openai.gpt-6.1-sol" +GPT_ON_BEDROCK_RUNTIME: Final = "bedrock/us.openai.gpt-6.1-sol" +GPT_ON_AZURE_OPENAI: Final = "azure/gpt-6.1-sol" +GPT_ON_OPENAI: Final = "openai/gpt-6.1-sol" +CLAUDE_ON_BEDROCK: Final = "bedrock/us.anthropic.claude-sonnet-5-5" +_BEDROCK_REGION: Final = "us-east-1" + +type SetupId = Literal["bedrock_mantle", "bedrock_runtime", "azure_openai", "openai", "claude_bedrock"] + + +@dataclass(frozen=True, slots=True) +class Deployment: + litellm_params: LiteLLMParamsBody + rate_bucket: RateBucket + + +@dataclass(frozen=True, slots=True) +class ModelSetup: + model: str + rate_bucket: RateBucket + + +def _drives(provider: Provider, model: str) -> pytest.MarkDecorator: + return meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(provider,), + models=(model,), + mode=Mode.STREAM, + ) + ) + + +SETUPS: Final = ( + pytest.param( + Deployment( + LiteLLMParamsBody(model=GPT_ON_BEDROCK_MANTLE, aws_region_name=_BEDROCK_REGION), RateBucket.BEDROCK_MANTLE + ), + id="bedrock_mantle", + marks=_drives(Provider.BEDROCK_MANTLE, GPT_ON_BEDROCK_MANTLE), + ), + pytest.param( + Deployment( + LiteLLMParamsBody(model=GPT_ON_BEDROCK_RUNTIME, aws_region_name=_BEDROCK_REGION), + RateBucket.BEDROCK_CONVERSE, + ), + id="bedrock_runtime", + marks=_drives(Provider.BEDROCK, GPT_ON_BEDROCK_RUNTIME), + ), + pytest.param( + Deployment( + LiteLLMParamsBody( + model=GPT_ON_AZURE_OPENAI, api_base="os.environ/AZURE_API_BASE", api_key="os.environ/AZURE_API_KEY" + ), + RateBucket.AZURE_OPENAI, + ), + id="azure_openai", + marks=_drives(Provider.AZURE, GPT_ON_AZURE_OPENAI), + ), + pytest.param( + Deployment(LiteLLMParamsBody(model=GPT_ON_OPENAI, api_key="os.environ/OPENAI_API_KEY"), RateBucket.OPENAI), + id="openai", + marks=_drives(Provider.OPENAI, GPT_ON_OPENAI), + ), + pytest.param( + Deployment( + LiteLLMParamsBody(model=CLAUDE_ON_BEDROCK, aws_region_name=_BEDROCK_REGION), RateBucket.BEDROCK_CONVERSE + ), + id="claude_bedrock", + marks=_drives(Provider.BEDROCK, CLAUDE_ON_BEDROCK), + ), +) + + +def with_expected_failure(setup_id: SetupId, *, reason: str) -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(*setup.values, id=setup.id, marks=(*setup.marks, pytest.mark.xfail(strict=True, reason=reason))) + if setup.id == setup_id + else setup + for setup in SETUPS + ) diff --git a/tests/e2e/codex/test_auto_compaction_e2e.py b/tests/e2e/codex/test_auto_compaction_e2e.py new file mode 100644 index 00000000000..5e0b2491002 --- /dev/null +++ b/tests/e2e/codex/test_auto_compaction_e2e.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from codex.harness import CodexHarness +from codex.setups import ModelSetup, with_expected_failure +from e2e_config import unique_marker + +pytestmark = pytest.mark.e2e + +COMPACT_AFTER_TOKENS: Final = 500 +PADDING: Final = "".join(f"line {number}: the quick brown fox jumps over the lazy dog\n" for number in range(400)) + + +@pytest.mark.parametrize( + "deployment", + with_expected_failure( + "claude_bedrock", + reason="LiteLLM's Responses-to-chat bridge ignores Codex's compaction_trigger item, so Codex gets no" + " compaction back and fails the turn", + ), +) +class TestCodexAutoCompaction: + def test_remembers_a_value_after_the_gateway_compacts_the_session( + self, codex: CodexHarness, setup: ModelSetup + ) -> None: + marker: Final = unique_marker() + workspace: Final = codex.new_workspace(files={"secret.txt": marker, "padding.txt": PADDING}) + first: Final = codex.exec( + "Run `cat secret.txt` and remember the value it prints. Then run `cat padding.txt`. Then reply with OK.", + model=setup.model, + cwd=workspace, + compact_after_tokens=COMPACT_AFTER_TOKENS, + ) + assert first.thread_id is not None, first + (workspace / "secret.txt").unlink() + second: Final = codex.resume( + first.thread_id, + "What value did secret.txt print earlier? Reply with just that value.", + model=setup.model, + cwd=workspace, + compact_after_tokens=COMPACT_AFTER_TOKENS, + ) + assert any(compaction.remote for compaction in codex.compactions(first.thread_id)), (first, second) + assert marker in second.answer, second diff --git a/tests/e2e/codex/test_command_round_trip_e2e.py b/tests/e2e/codex/test_command_round_trip_e2e.py new file mode 100644 index 00000000000..20beba53cb2 --- /dev/null +++ b/tests/e2e/codex/test_command_round_trip_e2e.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from codex.harness import CodexHarness +from codex.setups import SETUPS, ModelSetup +from e2e_config import unique_marker + +pytestmark = pytest.mark.e2e + + +@pytest.mark.parametrize("deployment", SETUPS) +class TestCodexCommandRoundTrip: + def test_answers_with_what_a_command_printed(self, codex: CodexHarness, setup: ModelSetup) -> None: + marker: Final = unique_marker() + workspace: Final = codex.new_workspace(files={"secret.txt": marker}) + run: Final = codex.exec( + "Run `cat secret.txt`, then reply with exactly what it printed.", model=setup.model, cwd=workspace + ) + assert run.exit_code == 0, run + assert any(marker in command.output for command in run.commands), run + assert marker in run.answer, run diff --git a/tests/e2e/codex/test_file_edit_e2e.py b/tests/e2e/codex/test_file_edit_e2e.py new file mode 100644 index 00000000000..5bf544d4eb7 --- /dev/null +++ b/tests/e2e/codex/test_file_edit_e2e.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from codex.harness import CodexHarness +from codex.setups import SETUPS, ModelSetup +from e2e_config import unique_marker + +pytestmark = pytest.mark.e2e + + +@pytest.mark.parametrize("deployment", SETUPS) +class TestCodexFileEdit: + def test_changes_a_file_with_the_patch_tool(self, codex: CodexHarness, setup: ModelSetup) -> None: + marker: Final = unique_marker() + workspace: Final = codex.new_workspace(files={"config.txt": "name = placeholder\n"}) + run: Final = codex.exec( + f"Use the apply_patch tool to change `placeholder` to `{marker}` in config.txt. Change nothing else.", + model=setup.model, + cwd=workspace, + ) + assert run.patched_files == (workspace / "config.txt",), run + assert (workspace / "config.txt").read_text() == f"name = {marker}\n", run diff --git a/tests/e2e/codex/test_long_session_e2e.py b/tests/e2e/codex/test_long_session_e2e.py new file mode 100644 index 00000000000..52769b720e3 --- /dev/null +++ b/tests/e2e/codex/test_long_session_e2e.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from codex.harness import CodexHarness +from codex.setups import SETUPS, ModelSetup +from e2e_config import unique_marker + +pytestmark = pytest.mark.e2e + + +@pytest.mark.parametrize("deployment", SETUPS) +class TestCodexLongSession: + def test_remembers_every_command_output_in_a_chain_of_files(self, codex: CodexHarness, setup: ModelSetup) -> None: + codes: Final = (unique_marker(), unique_marker(), unique_marker()) + second: Final = f"{unique_marker()}.txt" + third: Final = f"{unique_marker()}.txt" + workspace: Final = codex.new_workspace( + files={ + "start.txt": f"code {codes[0]}\nnext {second}\n", + second: f"code {codes[1]}\nnext {third}\n", + third: f"code {codes[2]}\nlast file\n", + } + ) + run: Final = codex.exec( + "Run `cat start.txt`. Each file prints a code and then the name of the next file, until one says last" + " file. Run `cat` on each next file, one command per step. Reply with the three codes in the order you" + " read them.", + model=setup.model, + cwd=workspace, + ) + assert len(run.commands) >= 3, run + assert all(code in run.answer for code in codes), run diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 367135eae26..d99bcff46e0 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -77,6 +77,10 @@ - {id: llm.bedrock_native.bedrock_invoke.basic.stream.works, module: llm, tier: P1, subject_endpoint: bedrock_native, route: bedrock_invoke, capability: basic, streaming: stream, assertions: [works], source: "vendor strategy §9.12 / LIT-4778", rationale: "Bedrock native invoke stream"} - {id: llm.bedrock_native.bedrock_invoke.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: bedrock_native, route: bedrock_invoke, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.12 / LIT-4778", rationale: "Bedrock invoke missing fields and invalid temperature"} - {id: llm.ocr.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: ocr, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.13 / LIT-4778", rationale: "OCR missing document rejected"} +- {id: llm.realtime.openai.translation.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: openai, capability: translation, streaming: stream, assertions: [works], source: "realtime_endpoints/endpoints.py", rationale: "Dedicated translation client-secret, raw SDP, and WebSocket paths emit translated audio and transcript deltas"} +- {id: llm.realtime.azure_openai.translation.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: azure_openai, capability: translation, streaming: stream, assertions: [works], source: "azure/realtime/handler.py", rationale: "Azure GA translation session emits translated audio and transcript deltas"} +- {id: llm.realtime.openai.transcription.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: openai, capability: transcription, streaming: stream, assertions: [works], source: "openai/realtime/handler.py", rationale: "gpt-live-transcribe and gpt-realtime-whisper emit live transcript deltas"} +- {id: llm.realtime.azure_openai.transcription.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: azure_openai, capability: transcription, streaming: stream, assertions: [works], source: "azure/realtime/handler.py", rationale: "Azure GA live transcription emits transcript deltas"} - {id: llm.rerank.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "llms/bedrock/rerank/handler.py", rationale: "Bedrock rerank"} - {id: llm.rerank.together_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: together_ai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/together_ai/rerank/handler.py", rationale: "Together rerank"} - {id: llm.images_generations.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_image_generation_e2e.py:22", rationale: "OpenAI image gen, b64/url"} @@ -94,6 +98,8 @@ - {id: llm.audio_speech.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_ai/text_to_speech/text_to_speech_handler.py", rationale: "Vertex TTS"} - {id: llm.audio_transcriptions.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_transcriptions, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "openai/transcriptions/handler.py", rationale: "OpenAI Whisper"} - {id: llm.audio_transcriptions.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_transcriptions, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.7 / LIT-4778", rationale: "Transcription empty file and missing model are rejected"} +- {id: llm.audio_transcriptions.openai.transcription.stream.works, module: llm, tier: P0, subject_endpoint: audio_transcriptions, route: openai, capability: transcription, streaming: stream, assertions: [works], source: "openai/transcriptions/handler.py", rationale: "gpt-transcribe streams typed transcript delta and done events"} +- {id: llm.audio_transcriptions.azure_openai.transcription.stream.works, module: llm, tier: P0, subject_endpoint: audio_transcriptions, route: azure_openai, capability: transcription, streaming: stream, assertions: [works], source: "azure/audio_transcriptions.py", rationale: "Azure gpt-transcribe streams typed transcript delta and done events over the v1 API"} - {id: llm.audio_transcriptions.azure_openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_transcriptions, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "azure/audio_transcriptions.py", rationale: "Azure STT"} - {id: llm.audio_transcriptions.soniox.basic.nonstream.works, module: llm, tier: P2, subject_endpoint: audio_transcriptions, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "soniox/audio_transcription/handler.py", rationale: "Soniox via OpenAI-compat (smoke)"} - {id: llm.audio_transcriptions.nvidia_riva.basic.nonstream.works, module: llm, tier: P2, subject_endpoint: audio_transcriptions, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "nvidia_riva/audio_transcription/handler.py", rationale: "NVIDIA Riva (smoke)"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index d2ba2c325df..7d9bc380bd0 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -93,6 +93,8 @@ LlmCapability = Literal[ "tool_search_history", "tool_use", "upstream_stream_failure", + "transcription", + "translation", "vision", "web_search", "web_search_server_tool", @@ -162,14 +164,7 @@ class OtherCell(_Base): Cell = Annotated[ - LlmCell - | MgmtCell - | McpCell - | ReliabilityCell - | QuotaCell - | LoggingCell - | GuardrailCell - | OtherCell, + LlmCell | MgmtCell | McpCell | ReliabilityCell | QuotaCell | LoggingCell | GuardrailCell | OtherCell, Field(discriminator="module"), ] diff --git a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md index c85471da90d..49b543bf666 100644 --- a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md @@ -19,6 +19,15 @@ a matching `function_call` output item, the test sends a `function_call_output` back, and the follow-up response incorporates the result (the temperature 72 appears). That raw-websocket tool path is the source of truth for tool calling. +The audio surfaces provision their own OpenAI and Azure deployments. +`TestRealtimeTranslation` streams the spoken weather question into +`/v1/realtime/translations` with French as the target and asserts translated +audio, a French transcript, and a billed spend row. `TestRealtimeLiveTranscription` +streams it into `/v1/realtime?intent=transcription` on gpt-live-transcribe and +asserts the deltas add up to the completed transcript and the session is billed. +`test_azure_ga_text_turn_is_answered_and_billed` guards a plain text turn over +the Azure GA deployment, spend row included + `test_pipecat_tool_smoke` is a realism layer through pipecat for openai, azure, and gemini only (not vertex_ai: native-audio live is flaky under pipecat tool calling while raw-ws tools pass; see pipecat-ai/pipecat#2544). Assertions are diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index 6bd5d61d721..7e7f9f424a2 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -29,9 +29,13 @@ from models import LiteLLMParamsBody _M = TypeVar("_M", bound=BaseModel) +REALTIME_ROUTE: Final = "/v1/realtime" +TRANSLATIONS_ROUTE: Final = "/v1/realtime/translations" -def realtime_ws_url(model: str) -> str: - return f"{ws_base_url()}/v1/realtime?{urlencode({'model': model})}" + +def realtime_ws_url(model: str, *, route: str = REALTIME_ROUTE, intent: str | None = None) -> str: + query: Final = {"model": model} if intent is None else {"model": model, "intent": intent} + return f"{ws_base_url()}{route}?{urlencode(query)}" @dataclass(frozen=True, slots=True) @@ -183,6 +187,75 @@ class InputAudioBufferAppend(BaseModel): audio: str +class InputAudioBufferCommit(BaseModel): + type: str = "input_audio_buffer.commit" + + +class SessionInputAudioBufferAppend(BaseModel): + type: str = "session.input_audio_buffer.append" + audio: str + + +AudioAppend = InputAudioBufferAppend | SessionInputAudioBufferAppend + + +class TranslationOutput(BaseModel): + language: str + + +class TranslationAudio(BaseModel): + output: TranslationOutput + + +class TranslationSessionConfig(BaseModel): + audio: TranslationAudio + + +class TranslationSessionUpdate(BaseModel): + type: str = "session.update" + session: TranslationSessionConfig + + +class PcmFormat(BaseModel): + type: str = "audio/pcm" + rate: int = 24000 + + +class InputTranscription(BaseModel): + model: str + + +class TranscriptionAudioInput(BaseModel): + format: PcmFormat = PcmFormat() + transcription: InputTranscription + + +class TranscriptionAudio(BaseModel): + input: TranscriptionAudioInput + + +class TranscriptionSessionConfig(BaseModel): + type: str = "transcription" + audio: TranscriptionAudio + + +class TranscriptionSessionUpdate(BaseModel): + type: str = "session.update" + session: TranscriptionSessionConfig + + +RealtimeClientEvent = ( + SessionUpdate + | TranslationSessionUpdate + | TranscriptionSessionUpdate + | ConversationItemCreate + | ResponseCreate + | InputAudioBufferAppend + | InputAudioBufferCommit + | SessionInputAudioBufferAppend +) + + def user_message(text: str) -> ConversationItemCreate: return ConversationItemCreate( item=MessageItem(content=[InputTextContent(text=text)]) @@ -320,7 +393,7 @@ class RealtimeSession: connection: Connection @step("Send the realtime event {event.type} over the websocket") - def send(self, event: SessionUpdate | ConversationItemCreate | ResponseCreate | InputAudioBufferAppend) -> None: + def send(self, event: RealtimeClientEvent) -> None: self.connection.send(event.model_dump_json(by_alias=True, exclude_none=True)) @step("Wait for a {stop_type} event on the realtime websocket") @@ -350,9 +423,9 @@ class RealtimeSession: @step("Stream the spoken question as mic chunks, then silence until the turn goes idle") def stream_and_collect( self, - chunks: Sequence[InputAudioBufferAppend], + chunks: Sequence[AudioAppend], *, - tail: InputAudioBufferAppend, + tail: AudioAppend, interval: float, idle: float, timeout: float, @@ -363,9 +436,9 @@ class RealtimeSession: def _live_mic_events( self, - chunks: Sequence[InputAudioBufferAppend], + chunks: Sequence[AudioAppend], *, - tail: InputAudioBufferAppend, + tail: AudioAppend, interval: float, idle: float, timeout: float, @@ -408,13 +481,19 @@ class RealtimeClient: ) return model_name, model_id - @step("Open a /v1/realtime websocket session to {model}") + @step("Open a {route} websocket session to {model}") @contextmanager def connect( - self, *, key: str, model: str, timeout: float = 15.0 + self, + *, + key: str, + model: str, + route: str = REALTIME_ROUTE, + intent: str | None = None, + timeout: float = 15.0, ) -> Generator[RealtimeSession, None, None]: with connect( - realtime_ws_url(model), + realtime_ws_url(model, route=route, intent=intent), additional_headers={"Authorization": f"Bearer {key}"}, open_timeout=timeout, ) as connection: diff --git a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py index 3b7ba3a2430..9645e31f50f 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py @@ -4,7 +4,10 @@ Each test opens a websocket through the proxy, speaks the OpenAI GA realtime event schema, and asserts the proxy normalizes the provider's stream into that schema: the session lifecycle, the canonical response event sequence with a reconstructed transcript and usage, and a full tool-call round-trip (call -> -tool result -> a follow-up response that uses the result). +tool result -> a follow-up response that uses the result). The audio tests +stream the spoken weather question into /v1/realtime/translations and into a +/v1/realtime?intent=transcription session, and assert the translated or +transcribed output and the session's spend row. One GA-speaking client validates every provider; only the model alias changes. Every provider is provisioned at session start, so a missing realtime alias is a @@ -12,19 +15,28 @@ hard failure, not a skip; once configured, a protocol failure is likewise a hard failure. See REALTIME_COVERAGE_MATRIX.md. """ +import base64 +import wave +from pathlib import Path +from types import MappingProxyType from typing import Final import pytest from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta from lifecycle import ResourceManager -from models import LiteLLMParamsBody +from models import LiteLLMParamsBody, SpendLogRow from pydantic import BaseModel from realtime_client import ( PROVIDERS, + TRANSLATIONS_ROUTE, ConversationItemCreate, + DeltaEvent, FunctionCallArgumentsDone, FunctionCallOutputItem, FunctionTool, + InputAudioBufferAppend, + InputAudioBufferCommit, + InputTranscription, JsonSchema, JsonSchemaProperty, RealtimeClient, @@ -35,8 +47,18 @@ from realtime_client import ( ResponseDone, ServerEnvelope, SessionConfig, + SessionInputAudioBufferAppend, SessionUpdate, + TranscriptionAudio, + TranscriptionAudioInput, + TranscriptionSessionConfig, + TranscriptionSessionUpdate, + TranslationAudio, + TranslationOutput, + TranslationSessionConfig, + TranslationSessionUpdate, as_text, + events_of_type, parse_last, realtime_model, transcript, @@ -264,3 +286,204 @@ def test_upstream_handshake_refusal_is_an_error_event_and_policy_close( assert closed.value.rcvd is not None assert closed.value.rcvd.code == 1008, closed.value + + +WEATHER_WAV: Final = Path(__file__).parent / "fixtures" / "weather_question_24k.wav" +PCM_RATE: Final = 24000 +CHUNK_SECONDS: Final = 0.1 +CHUNK_BYTES: Final = int(PCM_RATE * CHUNK_SECONDS) * 2 +SILENCE_B64: Final = base64.b64encode(bytes(CHUNK_BYTES)).decode() +REGISTRY_ROUTES: Final = MappingProxyType({"openai": "openai", "azure": "azure_openai"}) + + +def weather_pcm_chunks() -> tuple[str, ...]: + with wave.open(str(WEATHER_WAV), "rb") as wav: + assert (wav.getnchannels(), wav.getsampwidth(), wav.getframerate()) == (1, 2, PCM_RATE), ( + "fixture must be mono PCM16 at 24 kHz" + ) + pcm: Final = wav.readframes(wav.getnframes()) + return tuple(base64.b64encode(pcm[i : i + CHUNK_BYTES]).decode() for i in range(0, len(pcm), CHUNK_BYTES)) + + +def has_billed_row(rows: list[SpendLogRow]) -> bool: + return any((row.spend or 0) > 0 for row in rows) + + +def audio_marks(provider: RealtimeProvider, capability: str) -> tuple[pytest.MarkDecorator, ...]: + return ( + pytest.mark.covers(f"llm.realtime.{REGISTRY_ROUTES[provider.id]}.{capability}.stream.works"), + meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.REALTIME, + providers=(Provider(provider.id),), + models=(provider.litellm_params.model,), + capabilities=(Capability.AUDIO_INPUT,), + mode=Mode.WEBSOCKET, + ) + ), + ) + + +TRANSLATION_PROVIDERS: Final = ( + RealtimeProvider( + "openai", + "openai-translate", + LiteLLMParamsBody(model="openai/gpt-realtime-translate", api_key="os.environ/OPENAI_API_KEY"), + ), + RealtimeProvider( + "azure", + "azure-translate", + LiteLLMParamsBody( + model="azure/gpt-realtime-translate", + api_base="os.environ/AZURE_API_BASE", + api_key="os.environ/AZURE_API_KEY", + ), + ), +) + +LIVE_TRANSCRIPTION_PROVIDERS: Final = ( + RealtimeProvider( + "openai", + "openai-live-transcribe", + LiteLLMParamsBody(model="openai/gpt-live-transcribe", api_key="os.environ/OPENAI_API_KEY"), + ), + RealtimeProvider( + "azure", + "azure-live-transcribe", + LiteLLMParamsBody( + model="azure/gpt-live-transcribe", + api_base="os.environ/AZURE_API_BASE", + api_key="os.environ/AZURE_API_KEY", + ), + ), +) + +TRANSLATION_PARAMS: Final = tuple( + pytest.param(p, id=p.id, marks=audio_marks(p, "translation")) for p in TRANSLATION_PROVIDERS +) +LIVE_TRANSCRIPTION_PARAMS: Final = tuple( + pytest.param(p, id=p.id, marks=audio_marks(p, "transcription")) for p in LIVE_TRANSCRIPTION_PROVIDERS +) +AZURE_GA_PROVIDER: Final = next(p for p in PROVIDERS if p.id == "azure") + + +class TranscriptionCompleted(BaseModel): + type: str + transcript: str + + +class TestRealtimeTranslation: + @pytest.mark.parametrize("provider", TRANSLATION_PARAMS) + def test_spoken_english_comes_back_as_french_audio_and_transcript_and_is_billed( + self, + client: RealtimeClient, + resources: ResourceManager, + scoped_key: str, + provider: RealtimeProvider, + ) -> None: + model, model_id = client.provision(provider) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + with client.connect(key=scoped_key, model=model, route=TRANSLATIONS_ROUTE) as session: + session.send( + TranslationSessionUpdate( + session=TranslationSessionConfig(audio=TranslationAudio(output=TranslationOutput(language="fr"))) + ) + ) + session.collect_until("session.updated", timeout=30) + events: Final = session.stream_and_collect( + tuple(SessionInputAudioBufferAppend(audio=chunk) for chunk in weather_pcm_chunks()), + tail=SessionInputAudioBufferAppend(audio=SILENCE_B64), + interval=CHUNK_SECONDS, + idle=8, + timeout=60, + ) + + types: Final = tuple(e.type for e in events) + audio_bytes: Final = sum( + len(base64.b64decode(DeltaEvent.model_validate_json(e.payload).delta)) + for e in events_of_type(events, "session.output_audio.delta") + ) + assert audio_bytes > 0, f"translation returned no translated audio; types={types}" + translated: Final = "".join( + DeltaEvent.model_validate_json(e.payload).delta + for e in events_of_type(events, "session.output_transcript.delta") + ).lower() + assert translated.strip(), f"translation returned no translated text; types={types}" + assert "weather" not in translated, f"the spoken weather question came back untranslated: {translated!r}" + + rows: Final = client.proxy.poll_logs_for_key(scoped_key, predicate=has_billed_row) + assert has_billed_row(rows), f"the translation session was not billed: {rows}" + + +class TestRealtimeLiveTranscription: + @pytest.mark.parametrize("provider", LIVE_TRANSCRIPTION_PARAMS) + def test_spoken_question_streams_transcript_deltas_that_match_the_completed_transcript( + self, + client: RealtimeClient, + resources: ResourceManager, + scoped_key: str, + provider: RealtimeProvider, + ) -> None: + model, model_id = client.provision(provider) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + with client.connect(key=scoped_key, model=model, intent="transcription") as session: + session.send( + TranscriptionSessionUpdate( + session=TranscriptionSessionConfig( + audio=TranscriptionAudio( + input=TranscriptionAudioInput(transcription=InputTranscription(model=model)) + ) + ) + ) + ) + session.collect_until("session.updated", timeout=30) + for chunk in weather_pcm_chunks(): + session.send(InputAudioBufferAppend(audio=chunk)) + session.send(InputAudioBufferCommit()) + events: Final = session.collect_until("conversation.item.input_audio_transcription.completed", timeout=60) + + completed: Final = TranscriptionCompleted.model_validate_json(events[-1].payload) + assert "weather" in completed.transcript.lower(), ( + f"transcript of a spoken weather question does not mention weather: {completed.transcript!r}" + ) + streamed: Final = "".join( + DeltaEvent.model_validate_json(e.payload).delta + for e in events_of_type(events, "conversation.item.input_audio_transcription.delta") + ) + assert streamed.strip() == completed.transcript.strip(), ( + f"streamed deltas {streamed!r} do not add up to the completed transcript {completed.transcript!r}" + ) + + rows: Final = client.proxy.poll_logs_for_key(scoped_key, predicate=has_billed_row) + assert has_billed_row(rows), f"the live transcription session was not billed: {rows}" + + +@meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.REALTIME, + providers=(Provider.AZURE,), + models=(AZURE_GA_PROVIDER.litellm_params.model,), + mode=Mode.WEBSOCKET, + ) +) +def test_azure_ga_text_turn_is_answered_and_billed( + client: RealtimeClient, resources: ResourceManager, scoped_key: str +) -> None: + model, model_id = client.provision(AZURE_GA_PROVIDER) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + with client.connect(key=scoped_key, model=model) as session: + session.collect_until("session.created", timeout=20) + session.send(SessionUpdate(session=SessionConfig(instructions="Reply with the single word hello."))) + session.collect_until("session.updated", timeout=20) + session.send(user_message("Say the single word hello.")) + session.send(ResponseCreate()) + events: Final = session.collect_until("response.done", timeout=60) + + assert "hello" in transcript(events).lower(), f"Azure GA realtime did not answer the turn: {transcript(events)!r}" + rows: Final = client.proxy.poll_logs_for_key(scoped_key, predicate=has_billed_row) + assert has_billed_row(rows), f"the Azure GA realtime turn was not billed: {rows}" diff --git a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py index cd1ee8fb03c..43c276dfdf0 100644 --- a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py +++ b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py @@ -8,10 +8,13 @@ the shared multipart transport, since the SDK refuses to send them. A model-less request comes back as one of two 400s depending on whether any wildcard deployment happens to be registered on the shared proxy, so the assertion accepts either phrasing and holds both to naming the model as the problem. +Streamed gpt-transcribe (stream=true) on OpenAI and Azure must send transcript +deltas, a done event, and data: [DONE], and bill the call. """ from __future__ import annotations +from dataclasses import dataclass from pathlib import Path from typing import Final @@ -20,10 +23,10 @@ from e2e_config import unique_marker from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta from e2e_http import UnknownApiError, unwrap from lifecycle import ResourceManager -from models import LiteLLMParamsBody +from models import LiteLLMParamsBody, SpendLogRow from proxy_client import ProxyClient from pydantic import BaseModel, Field -from sdk_clients import SdkClients +from sdk_clients import SdkClients, response_header pytestmark = pytest.mark.e2e @@ -221,3 +224,102 @@ class TestWhisperTranscriptionFormats: assert all(word.start <= word.end for word in transcript.words), ( f"word timings out of order: {transcript.words}" ) + + +STREAM_DONE: Final = "[DONE]" + + +@dataclass(frozen=True, slots=True) +class _StreamingTranscriber: + id: str + registry_route: str + provider: Provider + litellm_params: LiteLLMParamsBody + + +STREAMING_TRANSCRIBERS: Final = ( + _StreamingTranscriber( + "openai", + "openai", + Provider.OPENAI, + LiteLLMParamsBody(model="openai/gpt-transcribe", api_key="os.environ/OPENAI_API_KEY"), + ), + _StreamingTranscriber( + "azure", + "azure_openai", + Provider.AZURE, + LiteLLMParamsBody( + model="azure/gpt-transcribe", + api_base="os.environ/AZURE_API_BASE", + api_key="os.environ/AZURE_API_KEY", + api_version="v1", + ), + ), +) + +STREAMING_PARAMS: Final = tuple( + pytest.param( + t, + id=t.id, + marks=( + pytest.mark.covers(f"llm.audio_transcriptions.{t.registry_route}.transcription.stream.works"), + meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.AUDIO, + providers=(t.provider,), + models=(t.litellm_params.model,), + mode=Mode.STREAM, + ) + ), + ), + ) + for t in STREAMING_TRANSCRIBERS +) + + +class _TranscriptStreamEvent(BaseModel): + type: str + delta: str = "" + text: str = "" + + +def _has_billed_row(rows: list[SpendLogRow]) -> bool: + return any((row.spend or 0) > 0 for row in rows) + + +class TestStreamedAudioTranscriptions: + @pytest.mark.parametrize("transcriber", STREAMING_PARAMS) + def test_stream_sends_transcript_deltas_then_done_and_bills_the_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, transcriber: _StreamingTranscriber + ) -> None: + model: Final = f"e2e-transcribe-stream-{transcriber.id}-{unique_marker()}" + model_id: Final = proxy.create_model(model, transcriber.litellm_params) + resources.defer(lambda: proxy.delete_model(model_id)) + client: Final = sdk.openai(resources.key()) + + with client.audio.transcriptions.with_streaming_response.create( + model=model, file=(WEATHER_WAV.name, WEATHER_WAV.read_bytes(), "audio/wav"), stream=True + ) as response: + content_type: Final = response_header(response.headers, "content-type") + call_id: Final = response_header(response.headers, "x-litellm-call-id") + payloads: Final = tuple( + line.removeprefix("data: ") for line in response.iter_lines() if line.startswith("data: ") + ) + + assert "text/event-stream" in (content_type or ""), f"stream=true did not stream: {content_type!r}" + assert payloads and payloads[-1] == STREAM_DONE, f"the stream did not end with data: [DONE]: {payloads[-3:]}" + assert payloads.count(STREAM_DONE) == 1, f"data: [DONE] was sent more than once: {payloads}" + events: Final = tuple(_TranscriptStreamEvent.model_validate_json(payload) for payload in payloads[:-1]) + deltas: Final = tuple(e.delta for e in events if e.type == "transcript.text.delta") + assert len(deltas) > 1, f"the transcript arrived in {len(deltas)} delta events, not as a stream: {events}" + done: Final = events[-1] + assert done.type == "transcript.text.done", f"the last event before [DONE] is not the done event: {done}" + assert "weather" in done.text.lower(), f"transcript of a spoken weather question lost the word: {done.text!r}" + assert "".join(deltas).strip() == done.text.strip(), ( + f"the deltas {deltas} do not add up to the final transcript {done.text!r}" + ) + + assert call_id, "the streamed response carries no x-litellm-call-id to find its spend log" + rows: Final = proxy.poll_logs_for_request_id(call_id, predicate=_has_billed_row) + assert _has_billed_row(rows), f"the streamed transcription was not billed: {rows}" diff --git a/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py b/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py index 91a31201b0f..59035f8f6b6 100644 --- a/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py @@ -15,48 +15,85 @@ pytestmark = pytest.mark.e2e _BETA_CONFIG: Final = files("litellm").joinpath("anthropic_beta_headers_config.json") _ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-5-5" +_VERTEX_MODEL: Final = "vertex_ai/claude-sonnet-5-5" +_VERTEX_LOCATION: Final = "global" +_AZURE_AI_MODEL: Final = "azure_ai/claude-fable-5-1" _BEDROCK_MODEL_ID: Final = "us.anthropic.claude-fable-5-1" _BETA_CONFIG_FILE: Final = TypeAdapter(dict[str, JsonValue]) _BETA_MAPPING: Final = TypeAdapter(dict[str, str | None]) _CODE_EXECUTION: Final[ToolUnionParam] = {"type": "code_execution_20250825", "name": "code_execution"} +# A live Vertex probe on 2026-10-31: each Claude 4.5 and 4.6 model takes one of these and Claude 5.5 takes neither +_VERTEX_PER_MODEL_BETAS: Final = frozenset({"computer-use-2025-01-24", "computer-use-2025-11-24"}) +_NO_PER_MODEL_BETAS: Final[frozenset[str]] = frozenset() -def _mapped_beta_headers(provider: str) -> str: - mapping: Final = _BETA_MAPPING.validate_python(_BETA_CONFIG_FILE.validate_json(_BETA_CONFIG.read_text())[provider]) - return ",".join(name for name, value in mapping.items() if value is not None) +def _mapped_beta_headers(provider: str, per_model: frozenset[str]) -> str: + config: Final = _BETA_CONFIG_FILE.validate_json(_BETA_CONFIG.read_text()) + mapping: Final = _BETA_MAPPING.validate_python(config[provider]) + header_only: Final = _BETA_MAPPING.validate_python(config.get(f"{provider}_header_only", {})) + forwarded: Final = {name for name, value in {**mapping, **header_only}.items() if value is not None} + assert per_model <= forwarded, per_model - forwarded + return ",".join(sorted(forwarded - per_model)) def _deployment(backend: str) -> LiteLLMParamsBody: - if backend.startswith("anthropic/"): - return LiteLLMParamsBody(model=backend, api_key="os.environ/ANTHROPIC_API_KEY") - return LiteLLMParamsBody( - model=backend, - aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", - aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", - aws_region_name="us-east-1", - ) + match backend.split("/", 1)[0]: + case "anthropic": + return LiteLLMParamsBody(model=backend, api_key="os.environ/ANTHROPIC_API_KEY") + case "vertex_ai": + return LiteLLMParamsBody( + model=backend, + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location=_VERTEX_LOCATION, + vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", + ) + case "azure_ai": + return LiteLLMParamsBody( + model=backend, + api_base="os.environ/AZURE_AI_API_BASE", + api_key="os.environ/AZURE_AI_API_KEY", + ) + case _: + return LiteLLMParamsBody( + model=backend, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="us-east-1", + ) class TestAnthropicBetaHeaders: @pytest.mark.parametrize( - ("backend", "beta_provider", "tools"), + ("backend", "beta_provider", "tools", "per_model"), [ - pytest.param(_ANTHROPIC_MODEL, "anthropic", [_CODE_EXECUTION], id="anthropic"), - pytest.param(f"bedrock/invoke/{_BEDROCK_MODEL_ID}", "bedrock", [], id="bedrock-invoke"), - pytest.param(f"bedrock/converse/{_BEDROCK_MODEL_ID}", "bedrock_converse", [], id="bedrock-converse"), + pytest.param(_ANTHROPIC_MODEL, "anthropic", [_CODE_EXECUTION], _NO_PER_MODEL_BETAS, id="anthropic"), + pytest.param(_VERTEX_MODEL, "vertex_ai", [], _VERTEX_PER_MODEL_BETAS, id="vertex-ai"), + pytest.param(_AZURE_AI_MODEL, "azure_ai", [], _NO_PER_MODEL_BETAS, id="azure-ai"), + pytest.param( + f"bedrock/invoke/{_BEDROCK_MODEL_ID}", "bedrock", [], _NO_PER_MODEL_BETAS, id="bedrock-invoke" + ), + pytest.param( + f"bedrock/converse/{_BEDROCK_MODEL_ID}", + "bedrock_converse", + [], + _NO_PER_MODEL_BETAS, + id="bedrock-converse", + ), ], ) @pytest.mark.covers( "llm.messages.anthropic.basic.nonstream.works", "llm.messages.bedrock_invoke.basic.nonstream.works", "llm.messages.bedrock_converse.basic.nonstream.works", + "llm.messages.vertex.basic.nonstream.works", + "llm.messages.azure_foundry.basic.nonstream.works", ) @meta( Subject( domain=Domain.LLM_TRANSLATION, route=Route.MESSAGES, - providers=(Provider.ANTHROPIC, Provider.BEDROCK), - models=(_ANTHROPIC_MODEL, _BEDROCK_MODEL_ID), + providers=(Provider.ANTHROPIC, Provider.BEDROCK, Provider.VERTEX_AI, Provider.AZURE_AI), + models=(_ANTHROPIC_MODEL, _BEDROCK_MODEL_ID, _VERTEX_MODEL, _AZURE_AI_MODEL), mode=Mode.NONSTREAM, ) ) @@ -68,6 +105,7 @@ class TestAnthropicBetaHeaders: backend: str, beta_provider: str, tools: list[ToolUnionParam], + per_model: frozenset[str], ) -> None: model: Final = f"e2e-messages-beta-{unique_marker()}" model_id: Final = proxy.create_model(model, _deployment(backend)) @@ -78,7 +116,7 @@ class TestAnthropicBetaHeaders: max_tokens=64, messages=[{"role": "user", "content": "Say 'hello' and nothing else"}], tools=tools, - extra_headers={"anthropic-beta": _mapped_beta_headers(beta_provider)}, + extra_headers={"anthropic-beta": _mapped_beta_headers(beta_provider, per_model)}, extra_body=NO_PROXY_CACHE, ) diff --git a/tests/e2e/llm_translation/test_realtime_http_e2e.py b/tests/e2e/llm_translation/test_realtime_http_e2e.py index ee3ce3893c2..3abb74308e8 100644 --- a/tests/e2e/llm_translation/test_realtime_http_e2e.py +++ b/tests/e2e/llm_translation/test_realtime_http_e2e.py @@ -1,14 +1,18 @@ """Vendor §9.19: realtime client_secrets + calls HTTP surface (LIT-4778). Websocket coverage already lives under realtime/; this file pins the HTTP -client-secret mint and the missing-auth contract. +client-secret mint (realtime, translation and transcription sessions), the +translation token scope, and the missing-auth contract. """ from __future__ import annotations +from dataclasses import dataclass +from typing import Final + import pytest from e2e_config import unique_marker -from e2e_http import NoBody, assert_auth_denied, unwrap +from e2e_http import NoBody, PartialBody, UnknownApiError, assert_auth_denied, unwrap from e2e_metadata import Domain, Provider, Route, Subject, meta from lifecycle import ResourceManager from models import LiteLLMParamsBody @@ -18,6 +22,7 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e REALTIME_BACKEND = "openai/gpt-realtime" +TRANSLATIONS_CLIENT_SECRETS: Final = "/v1/realtime/translations/client_secrets" class RealtimeSession(BaseModel): @@ -38,8 +43,65 @@ class RealtimeClientSecretRequest(BaseModel): session: RealtimeSession | None = None +class TurnDetection(BaseModel): + type: str + + +class SessionTranscription(BaseModel): + model: str | None = None + + +class SessionAudioInput(BaseModel): + transcription: SessionTranscription | None = None + turn_detection: TurnDetection | None = None + + +class SessionAudioOutput(BaseModel): + language: str | None = None + + +class SessionAudio(BaseModel): + input: SessionAudioInput | None = None + output: SessionAudioOutput | None = None + + +class TranslationSession(BaseModel): + model: str + audio: SessionAudio + + +class TranscriptionSession(BaseModel): + type: str = "transcription" + audio: SessionAudio + + +class SessionClientSecretRequest(BaseModel): + session: TranslationSession | TranscriptionSession + + +class NullableAudioInput(PartialBody): + turn_detection: TurnDetection | None = None + + +class NullableAudio(PartialBody): + input: NullableAudioInput + + +class NullableRealtimeSession(PartialBody): + type: str + model: str + instructions: str | None = None + audio: NullableAudio + + +class NullableClientSecretRequest(PartialBody): + model: str + session: NullableRealtimeSession + + class RealtimeClientSecretSession(BaseModel): type: str | None = None + audio: SessionAudio | None = None class RealtimeClientSecretResponse(BaseModel): @@ -49,15 +111,85 @@ class RealtimeClientSecretResponse(BaseModel): def _register(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: - model = f"e2e-realtime-http-{unique_marker()}" - model_id = proxy.create_model( - model, - LiteLLMParamsBody(model=REALTIME_BACKEND, api_key="os.environ/OPENAI_API_KEY"), + return _register_backend( + proxy, resources, LiteLLMParamsBody(model=REALTIME_BACKEND, api_key="os.environ/OPENAI_API_KEY") ) + + +def _register_backend( + proxy: ProxyClient, resources: ResourceManager, litellm_params: LiteLLMParamsBody +) -> tuple[str, str]: + model: Final = f"e2e-realtime-http-{unique_marker()}" + model_id: Final = proxy.create_model(model, litellm_params) resources.defer(lambda: proxy.delete_model(model_id)) return model, resources.key() +@dataclass(frozen=True, slots=True) +class _AudioBackend: + id: str + registry_route: str + provider: Provider + litellm_params: LiteLLMParamsBody + + @property + def upstream_model(self) -> str: + return self.litellm_params.model.split("/", 1)[1] + + +def _openai_backend(model: str) -> _AudioBackend: + return _AudioBackend( + "openai", "openai", Provider.OPENAI, LiteLLMParamsBody(model=model, api_key="os.environ/OPENAI_API_KEY") + ) + + +def _azure_backend(model: str) -> _AudioBackend: + return _AudioBackend( + "azure", + "azure_openai", + Provider.AZURE, + LiteLLMParamsBody(model=model, api_base="os.environ/AZURE_API_BASE", api_key="os.environ/AZURE_API_KEY"), + ) + + +TRANSLATION_BACKENDS: Final = ( + _openai_backend("openai/gpt-realtime-translate"), + _azure_backend("azure/gpt-realtime-translate"), +) +TRANSCRIPTION_BACKENDS: Final = ( + _openai_backend("openai/gpt-live-transcribe"), + _azure_backend("azure/gpt-live-transcribe"), +) + + +def _subject(backend: _AudioBackend) -> pytest.MarkDecorator: + return meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.REALTIME, + providers=(backend.provider,), + models=(backend.litellm_params.model,), + ) + ) + + +TRANSLATION_PARAMS: Final = tuple( + pytest.param( + b, + id=b.id, + marks=(pytest.mark.covers(f"llm.realtime.{b.registry_route}.translation.stream.works"), _subject(b)), + ) + for b in TRANSLATION_BACKENDS +) +TRANSCRIPTION_PARAMS: Final = tuple(pytest.param(b, id=b.id, marks=_subject(b)) for b in TRANSCRIPTION_BACKENDS) + + +def _french_translation_request(model: str) -> SessionClientSecretRequest: + return SessionClientSecretRequest( + session=TranslationSession(model=model, audio=SessionAudio(output=SessionAudioOutput(language="fr"))) + ) + + class TestRealtimeHttp: @pytest.mark.covers("llm.realtime.openai.basic.nonstream.works") @meta( @@ -110,3 +242,150 @@ class TestRealtimeHttp: json=NoBody(), ) assert_auth_denied(result, "realtime calls missing auth") + + +class TestRealtimeTranslationClientSecrets: + @pytest.mark.parametrize("backend", TRANSLATION_PARAMS) + def test_translations_route_mints_a_french_translation_session( + self, proxy: ProxyClient, resources: ResourceManager, backend: _AudioBackend + ) -> None: + model, key = _register_backend(proxy, resources, backend.litellm_params) + secret: Final = unwrap( + proxy.transport.post( + TRANSLATIONS_CLIENT_SECRETS, + headers=proxy.transport.bearer(key), + json=_french_translation_request(model), + response_type=RealtimeClientSecretResponse, + ) + ) + assert secret.value, f"client secret value missing: {secret}" + assert secret.session is not None and secret.session.type == "translation", ( + f"the translations route did not mint a translation session: {secret.session}" + ) + output: Final = secret.session.audio.output if secret.session.audio else None + assert output is not None and output.language == "fr", ( + f"the requested target language did not reach the session: {secret.session}" + ) + + @meta( + Subject( + domain=Domain.PROXY_AUTH, + route=Route.REALTIME, + providers=(Provider.OPENAI,), + models=(TRANSLATION_BACKENDS[0].litellm_params.model,), + ) + ) + def test_translation_token_is_refused_by_the_realtime_calls_route( + self, proxy: ProxyClient, resources: ResourceManager + ) -> None: + model, key = _register_backend(proxy, resources, TRANSLATION_BACKENDS[0].litellm_params) + secret: Final = unwrap( + proxy.transport.post( + TRANSLATIONS_CLIENT_SECRETS, + headers=proxy.transport.bearer(key), + json=_french_translation_request(model), + response_type=RealtimeClientSecretResponse, + ) + ) + assert secret.value, f"client secret value missing: {secret}" + result: Final = proxy.transport.send( + "/v1/realtime/calls", headers=proxy.transport.bearer(secret.value), json=NoBody() + ) + assert result.status_code == 401, ( + f"a translation token opened a plain realtime call: {result.status_code} {result.body[:300]}" + ) + assert "not valid for this Realtime endpoint" in result.body, ( + f"the refusal does not say the token belongs to another endpoint: {result.body[:300]}" + ) + + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.REALTIME, + providers=(Provider.OPENAI,), + models=(TRANSLATION_BACKENDS[0].litellm_params.model,), + ) + ) + def test_translation_session_on_the_realtime_route_is_rejected( + self, proxy: ProxyClient, resources: ResourceManager + ) -> None: + model, key = _register_backend(proxy, resources, TRANSLATION_BACKENDS[0].litellm_params) + result: Final = proxy.transport.post( + "/v1/realtime/client_secrets", + headers=proxy.transport.bearer(key), + json=RealtimeClientSecretRequest(model=model, session=RealtimeSession(type="translation", model=model)), + response_type=RealtimeClientSecretResponse, + ) + match result: + case UnknownApiError(status_code=400, body=body): + assert "translations endpoint" in body, f"the 400 does not point at the translations route: {body}" + case other: + pytest.fail(f"a translation session on /v1/realtime/client_secrets expected a 400, got {other!r}") + + +class TestRealtimeTranscriptionClientSecrets: + @pytest.mark.parametrize("backend", TRANSCRIPTION_PARAMS) + def test_transcription_session_runs_on_the_deployment_model( + self, proxy: ProxyClient, resources: ResourceManager, backend: _AudioBackend + ) -> None: + model, key = _register_backend(proxy, resources, backend.litellm_params) + secret: Final = unwrap( + proxy.transport.post( + "/v1/realtime/client_secrets", + headers=proxy.transport.bearer(key), + json=SessionClientSecretRequest( + session=TranscriptionSession( + audio=SessionAudio(input=SessionAudioInput(transcription=SessionTranscription(model=model))) + ) + ), + response_type=RealtimeClientSecretResponse, + ) + ) + assert secret.value, f"client secret value missing: {secret}" + assert secret.session is not None and secret.session.type == "transcription", ( + f"the client secret did not mint a transcription session: {secret.session}" + ) + audio_input: Final = secret.session.audio.input if secret.session.audio else None + transcription: Final = audio_input.transcription if audio_input else None + assert transcription is not None and transcription.model == backend.upstream_model, ( + f"the session does not transcribe with the deployment's model {backend.upstream_model}: {secret.session}" + ) + + +class TestRealtimeClientSecretNulls: + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.REALTIME, + providers=(Provider.OPENAI,), + models=(REALTIME_BACKEND,), + ) + ) + def test_top_level_null_is_dropped_and_nested_null_is_kept( + self, proxy: ProxyClient, resources: ResourceManager + ) -> None: + model, key = _register(proxy, resources) + secret: Final = unwrap( + proxy.transport.post( + "/v1/realtime/client_secrets", + headers=proxy.transport.bearer(key), + json=NullableClientSecretRequest( + model=model, + session=NullableRealtimeSession( + type="realtime", + model=model, + instructions=None, + audio=NullableAudio(input=NullableAudioInput(turn_detection=None)), + ), + ), + response_type=RealtimeClientSecretResponse, + ) + ) + assert secret.value, f"a session with instructions: null was not minted: {secret}" + audio_input: Final = secret.session.audio.input if secret.session and secret.session.audio else None + assert audio_input is not None and "turn_detection" in audio_input.model_fields_set, ( + f"the session does not report turn detection: {secret.session}" + ) + assert audio_input.turn_detection is None, ( + f"turn_detection: null was dropped, so the session fell back to {audio_input.turn_detection}" + ) diff --git a/tests/e2e/llm_translation/test_responses_retrieve_e2e.py b/tests/e2e/llm_translation/test_responses_retrieve_e2e.py index 4a08259c028..0e954af0952 100644 --- a/tests/e2e/llm_translation/test_responses_retrieve_e2e.py +++ b/tests/e2e/llm_translation/test_responses_retrieve_e2e.py @@ -69,9 +69,6 @@ def _retrieve_response(proxy: ProxyClient, key: str, response_id: str) -> Respon class TestResponsesRetrieve: - @pytest.mark.skip( - reason="stage red: product gap (LIT-5446), retrieve returns a different id than the stored response (non-idempotent response-id re-encryption)" - ) @pytest.mark.covers("llm.responses.openai.basic.nonstream.works") @meta( Subject( diff --git a/tests/e2e_harness/AGENTS.md b/tests/e2e_harness/AGENTS.md index 65a3b32df96..83e573ae249 100644 --- a/tests/e2e_harness/AGENTS.md +++ b/tests/e2e_harness/AGENTS.md @@ -2,7 +2,7 @@ Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network -The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, and `claude_code/` covers the driver, builder, probe and version resolver. Put a new harness test under the folder that mirrors the suite folder whose module it covers +The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, `claude_code/` covers the driver, builder and probe, `cli_harness/` covers the version resolver, the CLI env lockdown and the rate limiter, and `codex/` covers how the Codex driver reads the CLI's JSONL events and the compactions in a thread's rollout file. Put a new harness test under the folder that mirrors the suite folder whose module it covers Run them from the repo root. `e2e_config` reads `LITELLM_MASTER_KEY` at import and any value will do, the CI lane sets a dummy: diff --git a/tests/e2e_harness/claude_code/test_pr_gate_version_resolver.py b/tests/e2e_harness/claude_code/test_pr_gate_version_resolver.py deleted file mode 100644 index 0ed7e2bf083..00000000000 --- a/tests/e2e_harness/claude_code/test_pr_gate_version_resolver.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Unit tests for the Claude Code PR-gate version resolver. - -Markerless harness tests: they feed the resolver a hand-built packument and a -fixed clock, so they run without a proxy, never reach the npm registry, and -carry no `e2e` marker. -""" - -from __future__ import annotations - -from datetime import datetime, timezone -from typing import Final, Mapping - -import pytest - -from claude_code.pr_gate_version_resolver import NoEligibleVersionError, resolve_pr_gate_version - -NOW: Final = datetime(2026, 4, 25, 12, 0, tzinfo=timezone.utc) -INSIDE_THE_2_1_88_WINDOW: Final = datetime(2026, 4, 3, 12, 0, tzinfo=timezone.utc) - - -def _packument(times: Mapping[str, str], unpublished: frozenset[str] = frozenset()) -> dict[str, object]: - return { - "name": "@anthropic-ai/claude-code", - "time": {"created": "2024-01-01T00:00:00.000Z", "modified": "2026-04-25T00:00:00.000Z", **times}, - "versions": {version: {"version": version} for version in times if version not in unpublished}, - } - - -def test_skips_a_version_npm_has_unpublished() -> None: - metadata: Final = _packument( - { - "2.1.87": "2026-03-28T20:00:00.000Z", - "2.1.88": "2026-03-30T22:36:48.424Z", - "2.1.89": "2026-03-31T23:32:40.000Z", - }, - unpublished=frozenset({"2.1.88"}), - ) - assert resolve_pr_gate_version(metadata=metadata, as_of=INSIDE_THE_2_1_88_WINDOW) == "2.1.87" - - -def test_raises_when_the_only_old_enough_version_is_unpublished() -> None: - metadata: Final = _packument( - {"2.1.88": "2026-03-30T22:36:48.424Z", "2.1.89": "2026-03-31T23:32:40.000Z"}, - unpublished=frozenset({"2.1.88"}), - ) - with pytest.raises(NoEligibleVersionError): - resolve_pr_gate_version(metadata=metadata, as_of=INSIDE_THE_2_1_88_WINDOW) - - -def test_picks_the_newest_published_version_at_least_min_age_old() -> None: - metadata: Final = _packument( - { - "2.1.118": "2026-04-15T10:00:00.000Z", - "2.1.119": "2026-04-21T10:00:00.000Z", - "2.2.0-rc.1": "2026-04-22T10:00:00.000Z", - "2.1.120": "2026-04-23T10:00:00.000Z", - "2.1.121": "2026-04-25T11:00:00.000Z", - } - ) - assert resolve_pr_gate_version(metadata=metadata, as_of=NOW) == "2.1.119" diff --git a/tests/e2e_harness/claude_code/test_request_determinism.py b/tests/e2e_harness/claude_code/test_request_determinism.py index b7d330b7da6..c6edb02d1c5 100644 --- a/tests/e2e_harness/claude_code/test_request_determinism.py +++ b/tests/e2e_harness/claude_code/test_request_determinism.py @@ -34,7 +34,7 @@ from typing import List, Tuple import pytest from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude -from claude_code.rate_limiter import RateLimiter +from cli_harness.rate_limiter import RateLimiter pytestmark = pytest.mark.cli_determinism diff --git a/tests/e2e_harness/cli_harness/test_lockdown.py b/tests/e2e_harness/cli_harness/test_lockdown.py new file mode 100644 index 00000000000..fcb24f4a830 --- /dev/null +++ b/tests/e2e_harness/cli_harness/test_lockdown.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Final + +from cli_harness.lockdown import isolated_home, locked_down_env + +RUNNER_ENV: Final = { + "PATH": "/usr/bin", + "LANG": "C.UTF-8", + "HOME": "/home/runner", + "OPENAI_API_KEY": "sk-runner-secret", + "AWS_SECRET_ACCESS_KEY": "runner-aws-secret", + "DATABASE_URL": "postgresql://runner", +} + + +def test_cli_sees_runtime_vars_and_its_overrides_but_no_runner_credentials(tmp_path: Path) -> None: + env: Final = locked_down_env(home=tmp_path, overrides={"OPENAI_BASE_URL": "http://gateway"}, source=RUNNER_ENV) + assert dict(env) == { + "PATH": "/usr/bin", + "LANG": "C.UTF-8", + "HOME": str(tmp_path), + "OPENAI_BASE_URL": "http://gateway", + } + + +def test_isolated_home_starts_empty_and_is_gone_after_the_cli_exits() -> None: + with isolated_home("cli-harness-test-") as home: + assert home.is_dir() and not any(home.iterdir()) + (home / ".config").mkdir() + assert not home.exists() diff --git a/tests/e2e_harness/cli_harness/test_npm_version_resolver.py b/tests/e2e_harness/cli_harness/test_npm_version_resolver.py new file mode 100644 index 00000000000..3a84ccd361a --- /dev/null +++ b/tests/e2e_harness/cli_harness/test_npm_version_resolver.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final + +from cli_harness.npm_version_resolver import MIN_AGE, Packument, newest_version_at_least + +NOW: Final = datetime(2026, 4, 25, 12, 0, tzinfo=timezone.utc) +INSIDE_THE_2_1_88_WINDOW: Final = datetime(2026, 4, 3, 12, 0, tzinfo=timezone.utc) + + +def _packument(times: Mapping[str, str], unpublished: frozenset[str] = frozenset()) -> Packument: + return Packument.from_json( + { + "name": "@scope/some-cli", + "time": {"created": "2024-01-01T00:00:00.000Z", "modified": "2026-04-25T00:00:00.000Z", **times}, + "versions": {version: {"version": version} for version in times if version not in unpublished}, + } + ) + + +def test_skips_a_version_npm_has_unpublished() -> None: + packument: Final = _packument( + { + "2.1.87": "2026-03-28T20:00:00.000Z", + "2.1.88": "2026-03-30T22:36:48.424Z", + "2.1.89": "2026-03-31T23:32:40.000Z", + }, + unpublished=frozenset({"2.1.88"}), + ) + assert newest_version_at_least(packument, min_age=MIN_AGE, as_of=INSIDE_THE_2_1_88_WINDOW) == "2.1.87" + + +def test_finds_nothing_when_the_only_old_enough_version_is_unpublished() -> None: + packument: Final = _packument( + {"2.1.88": "2026-03-30T22:36:48.424Z", "2.1.89": "2026-03-31T23:32:40.000Z"}, + unpublished=frozenset({"2.1.88"}), + ) + assert newest_version_at_least(packument, min_age=MIN_AGE, as_of=INSIDE_THE_2_1_88_WINDOW) is None + + +def test_picks_the_newest_stable_version_at_least_min_age_old() -> None: + packument: Final = _packument( + { + "2.1.118": "2026-04-15T10:00:00.000Z", + "2.1.119": "2026-04-21T10:00:00.000Z", + "2.2.0-rc.1": "2026-04-22T10:00:00.000Z", + "2.1.120": "2026-04-23T10:00:00.000Z", + "2.1.121": "2026-04-25T11:00:00.000Z", + } + ) + assert newest_version_at_least(packument, min_age=MIN_AGE, as_of=NOW) == "2.1.119" + + +def test_newest_means_latest_published_not_highest_semver() -> None: + packument: Final = _packument( + { + "3.0.0": "2026-04-01T10:00:00.000Z", + "2.9.9": "2026-04-10T10:00:00.000Z", + } + ) + assert newest_version_at_least(packument, min_age=timedelta(days=3), as_of=NOW) == "2.9.9" diff --git a/tests/e2e_harness/cli_harness/test_rate_limiter.py b/tests/e2e_harness/cli_harness/test_rate_limiter.py new file mode 100644 index 00000000000..78ad3e92450 --- /dev/null +++ b/tests/e2e_harness/cli_harness/test_rate_limiter.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import multiprocessing +from collections.abc import Mapping +from dataclasses import dataclass +from multiprocessing.synchronize import Barrier +from pathlib import Path +from types import MappingProxyType +from typing import Final, NoReturn + +from cli_harness.rate_limiter import ProviderConfig, RateBucket, RateLimiter + +_RACERS: Final = 8 +_ATTEMPTS_PER_RACER: Final = 100 +_BURST: Final = 400 +_NO_REFILL: Final[Mapping[RateBucket, ProviderConfig]] = MappingProxyType( + {RateBucket.OPENAI: ProviderConfig(rate_per_sec=1e-9, burst=_BURST)} +) + + +@dataclass +class _FakeTime: + now: float = 0.0 + + def clock(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.now += seconds + + +def test_two_processes_on_one_host_share_a_single_budget_per_bucket(tmp_path: Path) -> None: + fake_time: Final = _FakeTime() + config: Final = {RateBucket.OPENAI: ProviderConfig(rate_per_sec=2.0, burst=1.0)} + first_process: Final = RateLimiter(config, tmp_path, fake_time.clock, fake_time.sleep) + second_process: Final = RateLimiter(config, tmp_path, fake_time.clock, fake_time.sleep) + + assert first_process.acquire(RateBucket.OPENAI) == 0.0 + assert 0.5 <= second_process.acquire(RateBucket.OPENAI) < 0.51 + assert second_process.acquire(RateBucket.BEDROCK_MANTLE) == 0.0 + + +def test_processes_racing_for_one_bucket_never_get_more_than_its_burst(tmp_path: Path) -> None: + grants: Final = tmp_path / "grants" + + _race(_RACERS, state_dir=tmp_path / "state", grants=grants) + + assert len(grants.read_bytes()) == _BURST + + +def _race(racers: int, *, state_dir: Path, grants: Path) -> None: + context: Final = multiprocessing.get_context("spawn") + start_line: Final = context.Barrier(racers) + processes: Final = tuple( + context.Process(target=_take_tokens_until_told_to_wait, args=(state_dir, grants, start_line)) + for _ in range(racers) + ) + for process in processes: + process.start() + for process in processes: + process.join() + + +def _take_tokens_until_told_to_wait(state_dir: Path, grants: Path, start_line: Barrier) -> None: + limiter: Final = RateLimiter(_NO_REFILL, state_dir, _frozen_clock, _exit_instead_of_sleeping) + with grants.open("ab", buffering=0) as log: + _ = start_line.wait() + for _ in range(_ATTEMPTS_PER_RACER): + _ = limiter.acquire(RateBucket.OPENAI) + _ = log.write(b".") + + +def _frozen_clock() -> float: + return 0.0 + + +def _exit_instead_of_sleeping(_seconds: float) -> NoReturn: + raise SystemExit(0) diff --git a/tests/e2e_harness/codex/test_codex_rollout.py b/tests/e2e_harness/codex/test_codex_rollout.py new file mode 100644 index 00000000000..2ef27a8a9c5 --- /dev/null +++ b/tests/e2e_harness/codex/test_codex_rollout.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import json +import time +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +from cli_harness.rate_limiter import RateBucket, RateLimiter +from codex.harness import CodexHarness +from codex.rollout import Compaction, rollout_compactions + +THREAD_ID: Final = "01a18fea-b2e8-7ed2-a29b-1ce4d2b138c3" +USER_PROMPT: Final = { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Run `cat secret.txt` and remember the value it prints."}], +} + + +def _line(record_type: str, payload: Mapping[str, object]) -> str: + return json.dumps({"timestamp": "2026-10-31T03:03:32.000Z", "type": record_type, "payload": payload}) + + +def _rollout(*compactions: Mapping[str, object]) -> str: + return "\n".join( + ( + _line("session_meta", {"id": THREAD_ID, "cli_version": "0.165.0"}), + _line("response_item", USER_PROMPT), + *(_line("compacted", compaction) for compaction in compactions), + _line("event_msg", {"type": "task_complete"}), + ) + ) + + +REMOTE: Final = { + "message": "", + "replacement_history": [USER_PROMPT, {"type": "compaction", "encrypted_content": "gAAAAABq5VrT"}], +} +LOCAL: Final = { + "message": "The user asked to read secret.txt, which printed 4a6e834a61f8d792.", + "replacement_history": [ + USER_PROMPT, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Summary: 4a6e834a61f8d792"}]}, + ], +} + + +class TestRolloutCompactions: + def test_a_compaction_that_kept_the_providers_compaction_item_is_remote(self) -> None: + assert rollout_compactions(_rollout(REMOTE)) == (Compaction(remote=True),) + + def test_a_compaction_that_codex_summarized_itself_is_local(self) -> None: + assert rollout_compactions(_rollout(LOCAL, {"message": "an older rollout with no replacement history"})) == ( + Compaction(remote=False), + Compaction(remote=False), + ) + + def test_a_rollout_without_compactions_has_none(self) -> None: + assert rollout_compactions(_rollout()) == () + + +def test_the_harness_reads_only_the_rollout_of_the_thread_it_is_asked_about(tmp_path: Path) -> None: + codex: Final = CodexHarness( + home=tmp_path, + api_key="unused", + rate_bucket=RateBucket.OPENAI, + limiter=RateLimiter({}, tmp_path, time.monotonic, time.sleep), + ) + sessions: Final = codex.codex_home / "sessions" / "2026" / "10" / "30" + sessions.mkdir(parents=True) + _ = (sessions / f"rollout-2026-10-30T20-03-32-{THREAD_ID}.jsonl").write_text(_rollout(REMOTE)) + _ = (sessions / "rollout-2026-10-30T20-04-00-01a18feb-9d3e-7ee1-b19f-e35d8894e5da.jsonl").write_text( + _rollout(LOCAL) + ) + + assert codex.compactions(THREAD_ID) == (Compaction(remote=True),) diff --git a/tests/e2e_harness/codex/test_codex_run.py b/tests/e2e_harness/codex/test_codex_run.py new file mode 100644 index 00000000000..415967da114 --- /dev/null +++ b/tests/e2e_harness/codex/test_codex_run.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import configparser +import json +import re +import shlex +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +from codex.run import Command, codex_run + +E2E_PYTEST_INI: Final = Path(__file__).resolve().parents[2] / "e2e" / "pytest.ini" + +ROUND_TRIP: Final = "\n".join( + ( + '{"type":"thread.started","thread_id":"01a18f9f-9d08-72e0-9be3-aa34cf7489d2"}', + '{"type":"turn.started"}', + '{"type":"item.completed","item":{"id":"item_0","type":"agent_message","text":"I will run it."}}', + '{"type":"item.started","item":{"id":"item_1","type":"command_execution","command":"/bin/zsh -lc \'cat' + ' secret.txt\'","aggregated_output":"","exit_code":null,"status":"in_progress"}}', + '{"type":"item.completed","item":{"id":"item_1","type":"command_execution","command":"/bin/zsh -lc \'cat' + ' secret.txt\'","aggregated_output":"marker-7f3a91\\n","exit_code":0,"status":"completed"}}', + '{"type":"item.completed","item":{"id":"item_2","type":"agent_message","text":"marker-7f3a91"}}', + '{"type":"turn.completed","usage":{"input_tokens":22547,"cached_input_tokens":11227,' + '"cache_write_input_tokens":11314,"output_tokens":67,"reasoning_output_tokens":0}}', + ) +) + + +def _line(event: Mapping[str, object]) -> str: + return json.dumps(event) + + +class TestCodexRun: + def test_the_answer_is_the_last_agent_message_and_each_finished_command_keeps_its_output(self) -> None: + run: Final = codex_run(exit_code=0, stdout=ROUND_TRIP, stderr="") + + assert run.thread_id == "01a18f9f-9d08-72e0-9be3-aa34cf7489d2" + assert run.commands == ( + Command(command="/bin/zsh -lc 'cat secret.txt'", output="marker-7f3a91\n", exit_code=0), + ) + assert run.answer == "marker-7f3a91" + assert run.usage is not None and run.usage.input_tokens == 22547 + assert run.errors == () + + def test_only_a_patch_that_applied_counts_as_patching_its_files(self) -> None: + def patch(status: str, path: str) -> str: + item: Final = { + "id": path, + "type": "file_change", + "status": status, + "changes": [{"path": path, "kind": "update"}], + } + return _line({"type": "item.completed", "item": item}) + + run: Final = codex_run( + exit_code=0, stdout="\n".join((patch("completed", "/w/app.py"), patch("failed", "/w/other.py"))), stderr="" + ) + + assert run.patched_files == (Path("/w/app.py"),) + + def test_a_turn_that_gave_up_on_a_5xx_is_rerun_by_the_suites_network_rule(self) -> None: + failed: Final = _line( + {"type": "turn.failed", "error": {"message": "exceeded retry limit, last status: 502 Bad Gateway"}} + ) + + run: Final = codex_run(exit_code=1, stdout=failed, stderr="") + + assert run.errors[0].status_code == 502 + rerun_patterns: Final = _only_rerun_patterns(E2E_PYTEST_INI) + assert any(re.search(pattern, repr(run)) for pattern in rerun_patterns), rerun_patterns + + def test_an_event_or_item_type_it_does_not_model_is_skipped(self) -> None: + newer_codex: Final = "\n".join( + ( + _line({"type": "thread.goal_set", "goal": "ship it"}), + _line({"type": "item.completed", "item": {"id": "item_9", "type": "todo_list", "items": []}}), + ROUND_TRIP, + ) + ) + + assert codex_run(exit_code=0, stdout=newer_codex, stderr="") == codex_run( + exit_code=0, stdout=ROUND_TRIP, stderr="" + ) + + +def _only_rerun_patterns(ini: Path) -> tuple[str, ...]: + parser: Final = configparser.ConfigParser() + _ = parser.read(ini) + args: Final = shlex.split(parser["pytest"]["addopts"]) + return tuple(value for flag, value in zip(args, args[1:]) if flag == "--only-rerun") diff --git a/tests/integration/_support/asgi.py b/tests/integration/_support/asgi.py index e84dc9439d3..30eb9f21760 100644 --- a/tests/integration/_support/asgi.py +++ b/tests/integration/_support/asgi.py @@ -15,7 +15,7 @@ from starlette.types import ASGIApp @contextmanager def asgi_server( - app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 1 + app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 30 ) -> Iterator[str]: with socket.socket() as listener: listener.bind(("127.0.0.1", port)) diff --git a/tests/integration/_support/custom_video_provider.py b/tests/integration/_support/custom_video_provider.py new file mode 100644 index 00000000000..3513d1e3186 --- /dev/null +++ b/tests/integration/_support/custom_video_provider.py @@ -0,0 +1,181 @@ +"""A customer-owned video provider for the proxy's `custom_provider_map`: a job queue with its own wire shape.""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +from httpx._types import RequestFiles + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_llm import CustomLLM +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject +from litellm.types.videos.utils import encode_video_id_with_provider, extract_original_video_id + +_STATUS_BY_STATE: Final = MappingProxyType({"queued": "queued", "running": "in_progress", "done": "completed"}) + + +def _video_object(job: Mapping[str, object], model: str | None, custom_llm_provider: str | None) -> VideoObject: + job_id: Final = str(job["job_id"]) + seconds: Final = job.get("seconds") + video: Final = VideoObject( + id=encode_video_id_with_provider(job_id, custom_llm_provider, model) if custom_llm_provider else job_id, + object="video", + status=_STATUS_BY_STATE[str(job["state"])], + created_at=0, + model=model, + seconds=str(seconds) if seconds is not None else None, + ) + return video + + +class QueueVideoConfig(BaseVideoConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseVideoConfig contract + return ["seconds"] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: BaseVideoConfig contract + return {key: value for key, value in video_create_optional_params.items() if key == "seconds"} + + def validate_environment( + self, + headers: dict, # mutable-ok: BaseVideoConfig contract + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict: # mutable-ok: BaseVideoConfig contract + key: Final = (litellm_params.api_key if litellm_params is not None else None) or api_key or "" + return {**headers, "X-Api-Key": key} + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: BaseVideoConfig contract + ) -> str: + if api_base is None: + raise ValueError("api_base is required") + return f"{api_base.rstrip('/')}/jobs" + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict, # mutable-ok: BaseVideoConfig contract + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[dict, RequestFiles, str]: # mutable-ok: BaseVideoConfig contract + seconds: Final = video_create_optional_request_params.get("seconds") + return ({"model": model, "prompt": prompt, "seconds": seconds}, [], api_base) + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + request_data: dict | None = None, + ) -> VideoObject: + return _video_object(raw_response.json(), model, custom_llm_provider) + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (f"{api_base}/{extract_original_video_id(video_id)}", {}) + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + client: object = None, + ) -> VideoObject: + return _video_object(raw_response.json(), None, custom_llm_provider) + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + variant: str | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (f"{api_base}/{extract_original_video_id(video_id)}/file", {}) + + def transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + client: object = None, + ) -> bytes: + return raw_response.content + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: dict[str, object] | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (api_base, {}) + + def transform_video_list_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig contract + return raw_response.json() + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + extra_body: dict[str, object] | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + raise NotImplementedError("the queue backend has no remix") + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("the queue backend has no remix") + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + raise NotImplementedError("the queue backend has no delete") + + def transform_video_delete_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> VideoObject: + raise NotImplementedError("the queue backend has no delete") + + +VIDEO_CONFIG: Final = QueueVideoConfig() +IMAGE_ONLY_HANDLER: Final = CustomLLM() diff --git a/tests/integration/_support/interception_vendor.py b/tests/integration/_support/interception_vendor.py index 3ca40087441..b46eb998f7a 100644 --- a/tests/integration/_support/interception_vendor.py +++ b/tests/integration/_support/interception_vendor.py @@ -27,7 +27,7 @@ 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"] +Mode: TypeAlias = Literal["search", "multi-search", "plain", "search-down", "followup-fails", "missing"] def deployment(mode: Mode) -> str: @@ -42,8 +42,8 @@ def plain(name: str) -> str: return f"plain-{name}" -def query(name: str) -> str: - return f"q-{name}" +def query(name: str, turn: int = 1) -> str: + return f"q-{name}" if turn == 1 else f"q-{name}-{turn}" def _minted(prefix: str) -> str: @@ -98,9 +98,20 @@ def _chat(name: str, body: Mapping[str, JsonValue]) -> 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): + tool_results: Final = tuple(message for message in messages if message.get("role") == "tool") + if tool_results: if name.endswith("-followup-fails"): return _azure_error(500, FOLLOWUP_FAILURE) + if name.endswith("-multi-search") and len(tool_results) == 1: + call: Final = { + "id": f"call_{name[4:36]}-2", + "type": "function", + "function": {"name": SEARCH_TOOL, "arguments": json.dumps({"query": query(name, turn=2)})}, + } + message: Final = JSON_OBJECT.validate_python({"role": "assistant", "content": None, "tool_calls": [call]}) + return _chat_completion( + _minted(f"chatcmpl-{name}-2-"), message, "tool_calls", FOLLOWUP_USAGE + ) return _chat_completion( _minted(f"chatcmpl-{name}-2-"), {"role": "assistant", "content": answer(name)}, "stop", FOLLOWUP_USAGE ) @@ -153,9 +164,20 @@ def _responses(body: Mapping[str, JsonValue]) -> 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): + function_call_outputs: Final = tuple(item for item in items if item.get("type") == "function_call_output") + if function_call_outputs: if name.endswith("-followup-fails"): return _azure_error(500, FOLLOWUP_FAILURE) + if name.endswith("-multi-search") and len(function_call_outputs) == 1: + call: Final = { + "type": "function_call", + "id": f"fc_{name[4:36]}-2", + "call_id": f"call_{name[4:36]}-2", + "name": SEARCH_TOOL, + "arguments": json.dumps({"query": query(name, turn=2)}), + "status": "completed", + } + return _response(_minted(f"resp_{name}_2_"), name, (call,), RESPONSES_FOLLOWUP_USAGE) return _response( _minted(f"resp_{name}_2_"), name, (_text_item(f"{name}_2", answer(name)),), RESPONSES_FOLLOWUP_USAGE ) diff --git a/tests/integration/authorization/_hidden_alias_budget.py b/tests/integration/authorization/_hidden_alias_budget.py index bee2a585cff..5e9c2f1a95a 100644 --- a/tests/integration/authorization/_hidden_alias_budget.py +++ b/tests/integration/authorization/_hidden_alias_budget.py @@ -150,6 +150,10 @@ def error_type(response: httpx.Response) -> str: return str(object_value(JSON_OBJECT.validate_json(response.content)["error"])["type"]) +def budget_refusal_error_type(route: str) -> str: + return "invalid_request_error" if route == "messages" else "budget_exceeded" + + def settle( send: Callable[[str], httpx.Response], status: int, *, burst: int = GATEWAY_BURST, seconds: float = 60 ) -> None: diff --git a/tests/integration/authorization/test_alias_chain_budget_waiver.py b/tests/integration/authorization/test_alias_chain_budget_waiver.py index 0908fa0691c..e9d922d14f9 100644 --- a/tests/integration/authorization/test_alias_chain_budget_waiver.py +++ b/tests/integration/authorization/test_alias_chain_budget_waiver.py @@ -44,6 +44,7 @@ from integration.authorization._hidden_alias_budget import ( async_anthropic_client, async_openai_client, base_url, + budget_refusal_error_type, error_type, fresh_chat, fresh_message, @@ -191,17 +192,18 @@ def _assert_refused( key: str, model: str, prefix: str, + route: str = "chat", send: Callable[[Gateway, str, str, str], httpx.Response] = fresh_chat, ) -> None: marker: Final = f"{prefix}-" + uuid.uuid4().hex refused: Final = send(candidate, model, key, marker) - assert _refusal(refused) == (BUDGET_EXCEEDED, "budget_exceeded"), refused.text + assert _refusal(refused) == (BUDGET_EXCEEDED, budget_refusal_error_type(route)), refused.text assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 def _assert_refused_on_every_route(rig: ChainRig, candidate: Gateway, key: str, model: str) -> None: for route, send in ROUTES: - _assert_refused(rig, candidate, key, model, f"refused-{route}", send) + _assert_refused(rig, candidate, key, model, f"refused-{route}", route, send) def _assert_provider_failure(response: httpx.Response) -> None: diff --git a/tests/integration/authorization/test_claude_code_gateway_oauth.py b/tests/integration/authorization/test_claude_code_gateway_oauth.py index 23476c439de..11e4485031e 100644 --- a/tests/integration/authorization/test_claude_code_gateway_oauth.py +++ b/tests/integration/authorization/test_claude_code_gateway_oauth.py @@ -474,15 +474,15 @@ def test_device_authorization_polls_complete_sso_and_mints_a_model_restricted_be ) assert model_b_response.status_code == 403, model_b_response.text assert _JSON_OBJECT.validate_json(model_b_response.content) == { + "type": "error", "error": { + "type": "permission_error", "message": ( f"The requested model '{model_b}' is not available for this API key, or the model name is " "invalid. Check the models available to you and try again." ), - "type": "team_model_access_denied", - "param": "model", - "code": "403", - } + "litellm_error_type": "team_model_access_denied", + }, }, model_b_response.text assert upstream.drain() == () diff --git a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py index a53fbfacf4c..fd4c7246801 100644 --- a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py +++ b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py @@ -695,9 +695,11 @@ def test_cli_session_token_is_denied_once_its_team_budget_is_exhausted( key=key, ) assert refused.status_code == 422, f"{refused.status_code} {refused.text}" - error: Final = object_value(JSON_OBJECT.validate_json(refused.content)["error"]) - assert error["type"] == "budget_exceeded", refused.text - assert error["code"] == "422", refused.text + body: Final = JSON_OBJECT.validate_json(refused.content) + assert body["type"] == "error", refused.text + error: Final = object_value(body["error"]) + assert error["type"] == "invalid_request_error", refused.text + assert "param" not in error and "code" not in error, refused.text message: Final = string_value(error["message"]) assert "Budget has been exceeded!" in message, refused.text assert f"Team={team}" in message, refused.text diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass.py b/tests/integration/authorization/test_hidden_alias_budget_bypass.py index 51b36c145d3..e9e55a6aaf7 100644 --- a/tests/integration/authorization/test_hidden_alias_budget_bypass.py +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass.py @@ -38,6 +38,7 @@ from integration.authorization._hidden_alias_budget import ( async_anthropic_client, async_openai_client, base_url, + budget_refusal_error_type, chat_statuses, clear_provider_script, error_type, @@ -599,7 +600,7 @@ def _assert_refused_like_target(rig: AliasRig, candidate: Gateway, key: str, sha target_marker: Final = f"target-{route}-" + uuid.uuid4().hex via_shadow: Final = send(candidate, shadow, key, shadow_marker) by_name: Final = send(candidate, target, key, target_marker) - assert _refusal(by_name) == (BUDGET_EXCEEDED, "budget_exceeded"), by_name.text + assert _refusal(by_name) == (BUDGET_EXCEEDED, budget_refusal_error_type(route)), by_name.text assert via_shadow.status_code == by_name.status_code, via_shadow.text assert error_type(via_shadow) == error_type(by_name), via_shadow.text observed: Final = upstream_requests(rig.gateway.upstream_url) @@ -621,7 +622,7 @@ def _assert_refused_on_every_route(rig: AliasRig, candidate: Gateway, key: str, marker: Final = f"chain-{route}-" + uuid.uuid4().hex refused: Final = send(candidate, model, key, marker) assert refused.status_code == BUDGET_EXCEEDED, refused.text - assert error_type(refused) == "budget_exceeded", refused.text + assert error_type(refused) == budget_refusal_error_type(route), refused.text assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 @@ -828,9 +829,9 @@ def _deployment_id(rig: AliasRig, name: str) -> str: return identities[0] -def _assert_budget_refusal(status_code: int, body: str) -> None: +def _assert_budget_refusal(status_code: int, body: str, route: str = "chat") -> None: assert status_code == BUDGET_EXCEEDED, body - assert "budget_exceeded" in body, body + assert budget_refusal_error_type(route) in body, body def _assert_no_served_row(key: str, marker: str) -> None: @@ -881,7 +882,7 @@ async def test_exhausted_key_is_refused_on_shadowing_alias_through_streamed_anth extra_headers=spend_marker(marker), ) as stream: await stream.get_final_message() - _assert_budget_refusal(refused.value.status_code, str(refused.value.body)) + _assert_budget_refusal(refused.value.status_code, str(refused.value.body), route="messages") assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 diff --git a/tests/integration/authorization/test_responses_id_ownership.py b/tests/integration/authorization/test_responses_id_ownership.py index 2ab26e75f50..c384ebb9434 100644 --- a/tests/integration/authorization/test_responses_id_ownership.py +++ b/tests/integration/authorization/test_responses_id_ownership.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import os import uuid -from collections.abc import Callable +from collections.abc import Callable, Iterator, Mapping from pathlib import Path from typing import Final, Literal from urllib.parse import urlsplit @@ -16,7 +16,7 @@ from integration._support.process import owned_proxy from integration._support.responses_vendor import same_response from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.worker_readiness import wait_for_worker_readiness -from openai.types.responses import ResponseInputParam +from openai.types.responses import ResponseCompletedEvent, ResponseInputParam from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter pytestmark: Final = pytest.mark.shares_worker_state @@ -51,12 +51,6 @@ class _DeleteResponse(BaseModel): deleted: bool -def _deployments_in_group(gateway: Gateway, group: str) -> int: - entries: Final = gateway.get("/model/info")["data"] - assert isinstance(entries, list), entries - return sum(object_value(entry).get("model_name") == group for entry in entries) - - def _wait_for_workers(gateway: Gateway, group: str, deployments: int) -> None: wait_for_worker_readiness(gateway, group, deployments, _PROXY_WORKERS) @@ -138,6 +132,378 @@ def _every_id_route(client: openai.OpenAI, model: str, response_id: str) -> tupl ) +def _comparison_body(response_id: str, nested: bool) -> Mapping[str, object]: + options: Final = {"comparison_response_id": response_id} + if nested: + return { + "prompt_cache_options": {"mode": "explicit", "ttl": "12h"}, + "extra_body": {"prompt_cache_options": options}, + } + return {"prompt_cache_options": options} + + +@pytest.mark.parametrize( + ("prefix", "nested", "stream"), + (("/v1", False, False), ("", True, False), ("/openai/v1", False, True), ("/v1", True, True)), +) +def test_regression_comparison_response_id_owner_repeated_requests_reach_raw_provider_id( + gateway: Gateway, prefix: str, nested: bool, stream: bool +) -> None: + raw_id: Final = f"resp_comparison_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.target == "/v1/models": + return _model_discovery_reply() + assert (request.method, request.target) == ("POST", "/v1/responses") + body: Final = _JSON_OBJECT.validate_json(request.body) + if body["input"] == "create comparison response": + assert body == {"model": _MODEL, "input": "create comparison response"}, request.body + return Reply(body=_response(raw_id)) + assert body == { + "model": _MODEL, + "input": "compare owned response", + "stream": stream, + "prompt_cache_options": {"comparison_response_id": raw_id}, + }, request.body + if stream: + return Reply( + content_type="text/event-stream", + chunks=( + b"data: " + + json.dumps({"type": "response.completed", "response": json.loads(_response(raw_id))}).encode() + + b"\n\n", + b"data: [DONE]\n\n", + ), + ) + return Reply(body=_response(raw_id)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + team: Final = scenario.team(models=[model]) + user: Final = scenario.member(team) + with _client(gateway, _key(scenario, team, user, model), prefix) as client: + created: Final = client.responses.create(model=model, input="create comparison response") + assert created.id != raw_id, created.model_dump_json() + assert len(_contract_requests(wire)) == 1 + + def compare_again() -> None: + result: Final = client.responses.create( + model=model, + input="compare owned response", + stream=stream, + extra_body=_comparison_body(created.id, nested), + ) + if isinstance(result, openai.Stream): + events: Final = tuple(result) + assert isinstance(events[-1], ResponseCompletedEvent), events + assert events[-1].response.status == "completed", events + else: + assert result.status == "completed", result.model_dump_json() + + for _ in range(2): + compare_again() + assert len(_contract_requests(wire)) == 2 + + +@pytest.mark.parametrize("nested", (False, True)) +@pytest.mark.parametrize("stranger", ("user", "team")) +@pytest.mark.parametrize("retained", (False, True), ids=("ordinary", "client-internal-field")) +def test_regression_comparison_response_id_other_identity_is_denied_before_dispatch( + gateway: Gateway, nested: bool, stranger: str, retained: bool +) -> None: + raw_id: Final = f"resp_comparison_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.target == "/v1/models": + return _model_discovery_reply() + return Reply(body=_response(raw_id)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + user_a: Final = scenario.member(team_a) + user_b: Final = scenario.member(team_a) + gateway.post("/team/member_add", {"team_id": team_b, "member": {"role": "user", "user_id": user_a}}) + with ( + _client(gateway, _key(scenario, team_a, user_a, model), "/v1") as owner, + _client( + gateway, + _key( + scenario, team_a if stranger == "user" else team_b, user_b if stranger == "user" else user_a, model + ), + "/v1", + ) as other, + ): + created: Final = owner.responses.create(model=model, input="create private comparison") + assert created.id != raw_id + assert len(_contract_requests(wire)) == 1 + _refused( + lambda: other.responses.create( + model=model, + input="compare another identity", + extra_body={ + **_comparison_body(created.id, nested), + **( + {"_litellm_addressed_comparison_response_id": "resp_unmanaged_comparison"} + if retained + else {} + ), + }, + ), + _OTHER_USER_DETAIL if stranger == "user" else _OTHER_TEAM_DETAIL, + ) + assert _contract_requests(wire) == (), "A cross-identity comparison reached the provider" + + +@pytest.mark.parametrize("nested", (False, True)) +@pytest.mark.parametrize("comparison_id", ("resp_unmanaged_comparison", "resp_!!!not-base64", "cmF3X3ZlbmRvcl9pZA==")) +def test_regression_unmanaged_comparison_response_id_is_denied_by_default( + gateway: Gateway, nested: bool, comparison_id: str +) -> None: + def respond(request: Request) -> Reply: + if request.target == "/v1/models": + return _model_discovery_reply() + return Reply(body=_response(f"resp_result_{uuid.uuid4().hex}")) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + team: Final = scenario.team(models=[model]) + user: Final = scenario.member(team) + with _client(gateway, _key(scenario, team, user, model), "/v1") as client: + _refused( + lambda: client.responses.create( + model=model, input="unmanaged comparison", extra_body=_comparison_body(comparison_id, nested) + ), + _UNMANAGED_DETAIL, + ) + assert _contract_requests(wire) == () + + +def test_comparison_extra_body_owned_id_overrides_foreign_top_level(gateway: Gateway) -> None: + foreign_raw_id: Final = f"resp_ignored_{uuid.uuid4().hex}" + owner_raw_id: Final = f"resp_effective_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.target == "/v1/models": + return _model_discovery_reply() + assert (request.method, request.target) == ("POST", "/v1/responses") + body: Final = _JSON_OBJECT.validate_json(request.body) + if body["input"] in ("seed ignored comparison", "seed effective comparison"): + assert body == {"model": _MODEL, "input": body["input"]}, request.body + return Reply(body=_response(foreign_raw_id if body["input"] == "seed ignored comparison" else owner_raw_id)) + assert body == { + "model": _MODEL, + "input": "compare effective options", + "prompt_cache_options": {"comparison_response_id": owner_raw_id}, + }, request.body + return Reply(body=_response(owner_raw_id)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + team: Final = scenario.team(models=[model]) + owner_user: Final = scenario.member(team) + foreign_user: Final = scenario.member(team) + with ( + _client(gateway, _key(scenario, team, owner_user, model), "/v1") as owner, + _client(gateway, _key(scenario, team, foreign_user, model), "/v1") as foreign, + ): + foreign_response: Final = foreign.responses.create(model=model, input="seed ignored comparison") + assert foreign_response.id != foreign_raw_id + owner_response: Final = owner.responses.create(model=model, input="seed effective comparison") + assert owner_response.id != owner_raw_id + assert len(_contract_requests(wire)) == 2 + compared: Final = owner.responses.create( + model=model, + input="compare effective options", + extra_body={ + "prompt_cache_options": { + "comparison_response_id": foreign_response.id, + "mode": "explicit", + "ttl": "12h", + }, + "extra_body": {"prompt_cache_options": {"comparison_response_id": owner_response.id}}, + }, + ) + assert compared.status == "completed", compared.model_dump_json() + assert len(_contract_requests(wire)) == 1 + + +def test_regression_comparison_extra_body_raw_id_overrides_foreign_top_level(gateway: Gateway) -> None: + raw_id: Final = "resp_raw_effective_comparison" + + def respond(request: Request) -> Reply: + if request.target == "/v1/models": + return _model_discovery_reply() + return Reply(body=_response("resp_foreign_seed")) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + team: Final = scenario.team(models=[model]) + owner_user: Final = scenario.member(team) + foreign_user: Final = scenario.member(team) + with ( + _client(gateway, _key(scenario, team, owner_user, model), "/v1") as owner, + _client(gateway, _key(scenario, team, foreign_user, model), "/v1") as foreign, + ): + foreign_response: Final = foreign.responses.create(model=model, input="seed ignored comparison") + assert foreign_response.id != "resp_foreign_seed" + assert len(_contract_requests(wire)) == 1 + _refused( + lambda: owner.responses.create( + model=model, + input="compare effective options", + extra_body={ + "prompt_cache_options": { + "comparison_response_id": foreign_response.id, + "mode": "explicit", + "ttl": "12h", + }, + "extra_body": {"prompt_cache_options": {"comparison_response_id": raw_id}}, + }, + ), + _UNMANAGED_DETAIL, + ) + assert _contract_requests(wire) == () + + +@pytest.fixture(scope="module") +def allowed_comparison_gateway( + session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[tuple[Gateway, Wire]]: + directory: Final = tmp_path_factory.mktemp("allowed_comparison") + config: Final = directory / "config.yaml" + config.write_text("general_settings:\n allow_unmanaged_response_ids: true\n") + with wire_server(lambda request: Reply(body=_response("resp_decoy"))) as decoy: + with owned_proxy( + session_gateway, + directory, + { + "INTEGRATION_PROXY_WORKERS": "1", + "OPENAI_BASE_URL": f"{decoy.url}/v1", + "OPENAI_API_KEY": "synthetic-env-fallback-key", + }, + config=config, + ) as allowed: + yield allowed, decoy + + +@pytest.mark.parametrize("nested", (False, True)) +@pytest.mark.parametrize("comparison_id", ("resp_unmanaged_comparison", "resp_!!!not-base64", "cmF3X3ZlbmRvcl9pZA==")) +def test_guard_opt_in_forwards_unmanaged_comparison_id( + allowed_comparison_gateway: tuple[Gateway, Wire], nested: bool, comparison_id: str +) -> None: + allowed, decoy = allowed_comparison_gateway + result_id: Final = f"resp_result_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.target == "/v1/models": + return _model_discovery_reply() + return Reply(body=_response(result_id)) + + with wire_server(respond) as wire, allowed.scenario() as scenario: + model: Final = _deployment(scenario, wire) + team: Final = scenario.team(models=[model]) + user: Final = scenario.member(team) + with _client(allowed, _key(scenario, team, user, model), "/v1") as client: + compared: Final = client.responses.create( + model=model, input="unmanaged comparison", extra_body=_comparison_body(comparison_id, nested) + ) + assert compared.model_dump(mode="json", exclude_none=True) == { + "id": compared.id, + "object": "response", + "created_at": 1.0, + "status": "completed", + "model": model, + "output": [], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + }, compared.model_dump_json() + assert compared.id != result_id + requests: Final = _contract_requests(wire) + assert len(requests) == 1 + assert (requests[0].method, requests[0].target) == ("POST", "/v1/responses") + assert _JSON_OBJECT.validate_json(requests[0].body) == { + "model": _MODEL, + "input": "unmanaged comparison", + "prompt_cache_options": {"comparison_response_id": comparison_id}, + }, requests[0].body + assert decoy.drain() == () + + +@pytest.mark.parametrize("stranger", ("user", "team")) +def test_regression_echoed_raw_comparison_id_is_denied_to_other_identity(gateway: Gateway, stranger: str) -> None: + raw_id: Final = f"resp_echo_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.target == "/v1/models": + return _model_discovery_reply() + body: Final = _JSON_OBJECT.validate_json(request.body) + if body["input"] == "create response": + assert body == {"model": _MODEL, "input": "create response"}, request.body + return Reply(body=_response(raw_id)) + assert body == { + "model": _MODEL, + "input": "compare owned response", + "prompt_cache_options": {"comparison_response_id": raw_id}, + }, request.body + return Reply( + body=json.dumps( + {**json.loads(_response(raw_id)), "prompt_cache_options": {"comparison_response_id": raw_id}} + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + user_a: Final = scenario.member(team_a) + user_b: Final = scenario.member(team_a) + gateway.post("/team/member_add", {"team_id": team_b, "member": {"role": "user", "user_id": user_a}}) + with ( + _client(gateway, _key(scenario, team_a, user_a, model), "/v1") as owner, + _client( + gateway, + _key( + scenario, team_a if stranger == "user" else team_b, user_b if stranger == "user" else user_a, model + ), + "/v1", + ) as other, + ): + created: Final = owner.responses.create(model=model, input="create response") + assert created.id != raw_id + assert len(_contract_requests(wire)) == 1 + echoed: Final = owner.responses.create( + model=model, input="compare owned response", extra_body=_comparison_body(created.id, False) + ) + echo_body: Final = _JSON_OBJECT.validate_json(echoed.model_dump_json(exclude_none=True)) + assert echo_body == { + "id": echoed.id, + "object": "response", + "created_at": 1.0, + "status": "completed", + "model": model, + "output": [], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "prompt_cache_options": {"comparison_response_id": raw_id}, + }, echo_body + options: Final = object_value(echo_body["prompt_cache_options"]) + leaked_id: Final = options["comparison_response_id"] + assert leaked_id == raw_id + assert isinstance(leaked_id, str) + assert len(_contract_requests(wire)) == 1 + _refused( + lambda: other.responses.create( + model=model, input="compare owned response", extra_body=_comparison_body(leaked_id, False) + ), + _UNMANAGED_DETAIL, + ) + assert _contract_requests(wire) == () + + @pytest.mark.parametrize("prefix", _ALIASES) def test_response_ids_are_scoped_to_the_issuing_user_and_team(gateway: Gateway, prefix: str) -> None: raw_id: Final = f"resp_owned_{uuid.uuid4().hex}" diff --git a/tests/integration/authorization/test_unknown_key_generic_401.py b/tests/integration/authorization/test_unknown_key_generic_401.py index 025f11d3973..a9ed4b66f27 100644 --- a/tests/integration/authorization/test_unknown_key_generic_401.py +++ b/tests/integration/authorization/test_unknown_key_generic_401.py @@ -50,7 +50,27 @@ class _ErrorResponse(BaseModel): error: _ProxyError +class _AnthropicErrorDetail(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: str + message: str + litellm_error_type: str + + +class _AnthropicErrorResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str + error: _AnthropicErrorDetail + + GENERIC_ERROR: Final = _ProxyError(message=GENERIC_MESSAGE, type="token_not_found_in_db", param="key", code="401") +MESSAGES_PATH: Final = "/v1/messages" +ANTHROPIC_ENVELOPE_PATHS: Final = frozenset({MESSAGES_PATH, "/anthropic/v1/messages"}) +ANTHROPIC_GENERIC_ERROR: Final = _AnthropicErrorDetail( + type="authentication_error", message=GENERIC_MESSAGE, litellm_error_type=GENERIC_ERROR.type +) @dataclass(frozen=True, slots=True) @@ -119,6 +139,22 @@ def _assert_generic_refusal(response: httpx.Response, *secrets: str) -> None: assert secret not in response.text, response.text +def _assert_messages_refusal(response: httpx.Response, *secrets: str) -> None: + assert response.status_code == 401, response.text + body: Final = _AnthropicErrorResponse.model_validate_json(response.content) + assert body.type == "error", response.text + assert body.error == ANTHROPIC_GENERIC_ERROR, response.text + for secret in secrets: + assert secret not in response.text, response.text + + +def _assert_refusal(path: str, response: httpx.Response, *secrets: str) -> None: + if path in ANTHROPIC_ENVELOPE_PATHS: + _assert_messages_refusal(response, *secrets) + else: + _assert_generic_refusal(response, *secrets) + + def _side(gateway: Gateway, peer: Gateway, side: Side) -> Gateway: return gateway if side == "gateway" else peer @@ -129,7 +165,10 @@ def test_an_unknown_key_gets_the_generic_401_without_reaching_the_provider( gateway: Gateway, peer: Gateway, provider: SharedProvider, route: str, side: Side ) -> None: key: Final = _unknown_key() - _assert_generic_refusal(_send(_side(gateway, peer, side), ROUTES[route], key), key, key[-4:], _digest(key)) + request_route: Final = ROUTES[route] + _assert_refusal( + request_route.path, _send(_side(gateway, peer, side), request_route, key), key, key[-4:], _digest(key) + ) assert provider.received() == () @@ -183,8 +222,9 @@ def test_two_different_unknown_keys_in_both_headers_get_the_generic_401_naming_n bearer: Final = _unknown_key() api_key: Final = _unknown_key() headers: Final = {**_headers("authorization", bearer), **_headers("x-api-key", api_key)} - for path, body in (("/v1/chat/completions", CHAT), ("/v1/messages", MESSAGES)): - _assert_generic_refusal( + for path, body in (("/v1/chat/completions", CHAT), (MESSAGES_PATH, MESSAGES)): + _assert_refusal( + path, gateway.client.post(path, json=dict(body), headers=headers), bearer, api_key, @@ -359,16 +399,13 @@ async def _send_async(client: httpx.AsyncClient, route: _Route, key: str) -> htt ) -async def _burst(owned: OwnedProxy, keys: tuple[str, ...]) -> tuple[httpx.Response, ...]: +async def _burst(owned: OwnedProxy, keys: tuple[str, ...]) -> tuple[tuple[_Route, httpx.Response], ...]: + routes: Final = tuple(ROUTES[BURST_ROUTES[index % len(BURST_ROUTES)]] for index in range(len(keys))) async with httpx.AsyncClient(base_url=str(owned.gateway.client.base_url), timeout=60, trust_env=False) as client: - return tuple( - await asyncio.gather( - *( - _send_async(client, ROUTES[BURST_ROUTES[index % len(BURST_ROUTES)]], key) - for index, key in enumerate(keys) - ) - ) + responses: Final = await asyncio.gather( + *(_send_async(client, route, key) for route, key in zip(routes, keys, strict=True)) ) + return tuple(zip(routes, responses, strict=True)) def _log_after(owned: OwnedProxy, offset: int) -> str: @@ -397,7 +434,7 @@ async def test_a_burst_across_a_worker_kill_answers_every_unknown_key_with_the_g lambda text: all(_console_line(key) in text for key in keys), 30, ) - for response, key in zip((*during, *after), keys, strict=True): - _assert_generic_refusal(response, key, _digest(key)) + for (route, response), key in zip((*during, *after), keys, strict=True): + _assert_refusal(route.path, response, key, _digest(key)) assert [window.count(_console_line(key)) for key in keys] == [1] * len(keys) assert [key for key in keys if key in window] == [] diff --git a/tests/integration/caching/test_response_cache_empty_output.py b/tests/integration/caching/test_response_cache_empty_output.py index ce5a4d64ca7..7a20e712d4e 100644 --- a/tests/integration/caching/test_response_cache_empty_output.py +++ b/tests/integration/caching/test_response_cache_empty_output.py @@ -23,7 +23,7 @@ from typing import Final, TypeAlias, TypeVar import httpx import pytest from anthropic import Anthropic, AsyncAnthropic -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, InternalServerError, OpenAI from pydantic import JsonValue from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value @@ -420,14 +420,25 @@ def _streamed_chat(gateway: Gateway, model: str, text: str) -> _Answer: return ids, streamed, tuple(chunk.object for chunk in chunks) -def test_streamed_chat_without_choices_is_never_assembled_or_cached(gateway: Gateway) -> None: +def test_streamed_chat_without_choices_fails_and_is_never_cached(gateway: Gateway) -> None: upstream: Final = Upstream(gateway.upstream_url) + store: Final = redis_store() text: Final = prompt() with gateway.scenario() as scenario: handle: Final = scripted(scenario, CHAT_STREAM_WITHOUT_CHOICES) model: Final = scenario.model(api_base=handle.api_base()) - empty: Final = _never_replayed(partial(_streamed_chat, gateway, model, text), upstream, handle.scenario_id) - assert all(answer.text == "" for answer in empty), empty + stream: Final = partial(_streamed_chat, gateway, model, text) + for attempt in range(1, EMPTY_STREAMS + 1): + with pytest.raises(InternalServerError, match="Stream ended before a finish_reason was received"): + stream() + _assert_prompt_reached_upstream(upstream, handle.scenario_id, text, attempt) + rescript(handle, CHAT_STREAM) + real_ids, real_text, _ = stream() + assert real_text == "streamed response" and len(real_ids) == 1, (real_ids, real_text) + _assert_prompt_reached_upstream(upstream, handle.scenario_id, text, EMPTY_STREAMS + 1) + await_entry(store, next(iter(real_ids))) + replay: Final = _replayed(stream, upstream, handle.scenario_id, real_ids) + assert replay.text == real_text, replay def _streamed_text(gateway: Gateway, model: str, text: str) -> _Answer: diff --git a/tests/integration/compatibility/test_a2a_wire_versions.py b/tests/integration/compatibility/test_a2a_wire_versions.py index 4471322d1e9..b3ca0c99392 100644 --- a/tests/integration/compatibility/test_a2a_wire_versions.py +++ b/tests/integration/compatibility/test_a2a_wire_versions.py @@ -1023,9 +1023,18 @@ def test_a2a_task_methods_pass_through_verbatim_and_push_callbacks_are_validated {"taskId": "t1", "url": callback}, ): rejected = call("tasks/pushNotificationConfig/set", params) - assert (rejected.status_code, rejected.json()) == (400, {"detail": detail}), ( - f"{params}: {rejected.text}" - ) + assert (rejected.status_code, rejected.json()) == ( + 400, + { + "jsonrpc": "2.0", + "id": marker + "tasks/pushNotificationConfig/set", + "error": { + "code": -32000, + "message": detail, + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, + ), f"{params}: {rejected.text}" assert tuple(item for item in wire.drain() if item.method == "POST") == () public: Final = {"taskId": "t1", "pushNotificationConfig": {"url": "https://1.1.1.1/callback"}} @@ -1118,7 +1127,15 @@ def test_a2a_v1_task_methods_and_extended_card_lower_to_0_3_upstream_and_answer_ rejected = call("CreateTaskPushNotificationConfig", params) assert (rejected.status_code, rejected.json()) == ( 400, - {"detail": "Push notification URL must use HTTPS"}, + { + "jsonrpc": "2.0", + "id": marker + "CreateTaskPushNotificationConfig", + "error": { + "code": -32000, + "message": "Push notification URL must use HTTPS", + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, ), f"{params}: {rejected.text}" assert forwarded() == () diff --git a/tests/integration/compatibility/test_audio_transcriptions_wire.py b/tests/integration/compatibility/test_audio_transcriptions_wire.py index 3deeaebaf7d..faebfda715b 100644 --- a/tests/integration/compatibility/test_audio_transcriptions_wire.py +++ b/tests/integration/compatibility/test_audio_transcriptions_wire.py @@ -345,10 +345,6 @@ def test_audio_transcription_forwards_nested_and_repeated_bracket_fields_verbati def test_audio_transcription_stream_returns_server_sent_events(gateway: Gateway) -> None: - pytest.skip( - "BUG: stream=true transcriptions reach the upstream but its SSE body comes back as application/json " - '{"text": "data: ..."}, so OpenAI SDK stream=True callers get no transcript events' - ) wav: Final = _WAV_BYTES + uuid.uuid4().bytes def respond(request: Request) -> Reply: diff --git a/tests/integration/compatibility/test_codex_model_catalog.py b/tests/integration/compatibility/test_codex_model_catalog.py index a3ba032e612..fbf05b0f526 100644 --- a/tests/integration/compatibility/test_codex_model_catalog.py +++ b/tests/integration/compatibility/test_codex_model_catalog.py @@ -351,6 +351,41 @@ def test_stock_model_keeps_codex_entry(gateway: Gateway) -> None: assert catalog[name] == expected, catalog[name] +def test_stock_model_alias_takes_configured_max_input_tokens(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + stock: Final = _bundled_stock()["gpt-6-sol"] + stock_window: Final = stock["context_window"] + ceiling: Final = stock["max_context_window"] + assert isinstance(stock_window, int) and isinstance(ceiling, int) and stock_window < ceiling, stock + raised: Final = _model_name("codex-raised-window") + lowered: Final = _model_name("codex-lowered-window") + half_set: Final = _model_name("codex-half-set-window") + half_lowered: Final = _model_name("codex-half-lowered-window") + unset: Final = _model_name("codex-unset-window") + names: Final = (raised, lowered, half_set, half_lowered, unset) + _new_model(scenario, raised, model="openai/gpt-6-sol", model_info={"max_input_tokens": ceiling + 50000}) + _new_model(scenario, lowered, model="openai/gpt-6-sol", model_info={"max_input_tokens": stock_window // 2}) + _new_model(scenario, half_set, model="openai/gpt-6-sol", model_info={"max_input_tokens": ceiling}) + _new_model(scenario, half_set, model="openai/gpt-6-sol") + _new_model(scenario, half_lowered, model="openai/gpt-6-sol", model_info={"max_input_tokens": stock_window // 3}) + _new_model(scenario, half_lowered, model="openai/gpt-6-sol") + _new_model(scenario, unset, model="openai/gpt-6-sol") + expected: Final = { + raised: ceiling, + lowered: stock_window // 2, + half_set: stock_window, + half_lowered: stock_window // 3, + unset: stock_window, + } + windows: Final = _settled( + lambda: { + slug: entry["context_window"] for slug, entry in _by_slug(_catalog(gateway)).items() if slug in names + }, + lambda value: value == expected, + ) + assert windows == expected, windows + + def test_stock_model_configured_tiers_override(gateway: Gateway) -> None: with gateway.scenario() as scenario: stock: Final = _bundled_stock()["gpt-6-sol"] diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths.py b/tests/integration/compatibility/test_openai_sdk_client_paths.py index f8d7449038a..d6b6d1d460d 100644 --- a/tests/integration/compatibility/test_openai_sdk_client_paths.py +++ b/tests/integration/compatibility/test_openai_sdk_client_paths.py @@ -630,19 +630,6 @@ def _spend_row(identity: str) -> Mapping[str, JsonValue]: 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), @@ -815,7 +802,7 @@ def test_r07_responses_stream_through_the_async_openai_sdk(gateway: Gateway) -> 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 + assert _spend_row(last.response.id)["model"] == DEPLOYMENT_MODEL def test_r08_completions_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: diff --git a/tests/integration/configuration/test_effective_settings.py b/tests/integration/configuration/test_effective_settings.py index 07477f1f048..1b0086c619a 100644 --- a/tests/integration/configuration/test_effective_settings.py +++ b/tests/integration/configuration/test_effective_settings.py @@ -20,7 +20,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter from tests.integration._support.client import Gateway, eventually, object_value, string_value from tests.integration._support.database import read_rows, scratch_database -from tests.integration._support.process import owned_proxy +from tests.integration._support.process import owned_proxy, owned_proxy_process from tests.integration._support.wire import Reply, Request, Wire, wire_server pytestmark: Final = pytest.mark.shares_worker_state @@ -1701,3 +1701,57 @@ def test_db_environment_variables_saved_from_a_session_outside_the_tables_schema assert isinstance(first, int) and isinstance(second, int) assert 0 < first < second + + +def test_a_stored_general_setting_no_settings_api_manages_stays_off_and_the_rest_of_the_row_applies( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + alias: Final = "stale-database-setting-" + uuid.uuid4().hex + with scratch_database() as database_url, wire_server(lambda _request: _completion_reply("pong")) as provider: + monkeypatch.setenv("DATABASE_URL", database_url) + _push_schema() + with psycopg.connect(database_url, autocommit=True) as writer: + writer.execute( + 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, %s)', + ("general_settings", Jsonb({"enforce_user_param": True, "store_prompts_in_spend_logs": True})), + ) + config: Final = tmp_path / "stale-database-setting.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": alias, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": provider.url + "/v1", + "api_key": "integration-provider-key", + }, + } + ], + "general_settings": {"master_key": gateway.key, "store_model_in_db": True}, + } + ) + ) + with owned_proxy_process( + gateway, + tmp_path / "stale-database-setting", + {"DATABASE_URL": database_url, "DISABLE_SCHEMA_UPDATE": "true"}, + config=config, + database_setup=(), + ) as owned: + answer: Final = owned.gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": [{"role": "user", "content": "ping"}]} + ) + listed: Final = owned.gateway.request("GET", "/config/list", params={"config_type": "general_settings"}) + log: Final = owned.log.read_text() + + assert answer.status_code == 200, answer.text + assert listed.status_code == 200, listed.text + stored_prompts: Final = next( + entry + for entry in TypeAdapter(list[dict[str, JsonValue]]).validate_json(listed.content) + if entry["field_name"] == "store_prompts_in_spend_logs" + ) + assert (stored_prompts["field_value"], stored_prompts["source"]) == (True, "db") + assert "general_settings.enforce_user_param is stored in the database, but the proxy does not apply" in log diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index d9f17c12bda..6495199b8bb 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -1588,6 +1588,7 @@ AZURE_PINNED_GPT_5_4_MINI_STREAM: Final = CostTrackingTestCase( content_type="text/event-stream", frames=( 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}\n\n', + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[],"usage":{"prompt_tokens":1840,"completion_tokens":412,"total_tokens":2252}}\n\n', "data: [DONE]\n\n", ), diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index bc7fb3486be..270bb6e105d 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -7952,6 +7952,7 @@ GPT_5_6_STREAM_USAGE_IN_EMPTY_CHOICES_CHUNK: Final = CostTrackingTestCase( frames=( 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}],"usage":null}', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{"content":"scripted answer"},"finish_reason":null}],"usage":null}', + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null}', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[],"usage":{"prompt_tokens":1840,"completion_tokens":412,"total_tokens":2252}}', "data: [DONE]", ), diff --git a/tests/integration/database/test_prisma_migration_lock.py b/tests/integration/database/test_prisma_migration_lock.py index 7482e96d3d3..f4a2a022668 100644 --- a/tests/integration/database/test_prisma_migration_lock.py +++ b/tests/integration/database/test_prisma_migration_lock.py @@ -19,7 +19,8 @@ def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object, def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]: rows: Final = connection.execute( - "SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted" + "SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted " + "AND database = (SELECT oid FROM pg_database WHERE datname = current_database())" ).fetchall() return frozenset(int(row[0]) for row in rows) diff --git a/tests/integration/management/test_agent_management.py b/tests/integration/management/test_agent_management.py index 5c0fdc68e09..a4c2650ffdd 100644 --- a/tests/integration/management/test_agent_management.py +++ b/tests/integration/management/test_agent_management.py @@ -5,12 +5,14 @@ from typing import Final from integration._support.client import Gateway, Scenario from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue _REDACTED: Final = "REDACTED_BY_LITELM" class _AgentRead(BaseModel): + model_config = ConfigDict(frozen=True) + agent_id: str agent_name: str litellm_params: dict[str, JsonValue] | None @@ -84,7 +86,8 @@ def test_agent_reads_redact_credentials_filter_by_key_permission_and_keep_the_st visible: Final = _register_agent(gateway, scenario, marker + "-visible", wire.url, secret) hidden: Final = _register_agent(gateway, scenario, marker + "-hidden", wire.url, secret) restricted: Final = scenario.key(object_permission={"agents": [visible]}) - redacted_params: Final = {"api_key": _REDACTED, "cost_per_query": 0.1, "is_public": False} + stored_params: Final = {"api_key": _REDACTED, "cost_per_query": 0.1} + redacted_params: Final = {**stored_params, "is_public": False} listed: Final = gateway.client.get("/v1/agents", headers={"Authorization": f"Bearer {restricted}"}) assert listed.status_code == 200, listed.text @@ -98,9 +101,8 @@ def test_agent_reads_redact_credentials_filter_by_key_permission_and_keep_the_st assert fetched.status_code == 200, fetched.text assert secret not in fetched.text, fetched.text read: Final = _AgentRead.model_validate_json(fetched.content) - assert (read.agent_id, read.litellm_params, read.static_headers) == (visible, redacted_params, None), ( - fetched.text - ) + assert (read.agent_id, read.static_headers) == (visible, None), fetched.text + assert read.litellm_params in (stored_params, redacted_params), fetched.text refused: Final = gateway.client.get(f"/v1/agents/{hidden}", headers={"Authorization": f"Bearer {restricted}"}) assert (refused.status_code, refused.json()) == ( diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index 8b3c6b06bba..2a8508ff3c9 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -24,6 +24,7 @@ import yaml from pydantic import JsonValue from redis import Redis +from litellm.constants import TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS from tests.integration._support.client import ( Gateway, Scenario, @@ -166,19 +167,11 @@ def _remove_team_by_sql(team_id: str) -> None: assert _team_rows(team_id) == [] -def _remove_team_if_present(gateway: Gateway, team_id: str) -> None: - """Cleanup for teams the test deletes itself: the API delete first, SQL for anything it leaves.""" - if not _team_rows(team_id): - return - gateway.request("POST", "/team/delete", {"team_ids": [team_id]}) - _remove_team_by_sql(team_id) - - def _create_team(scenario: Scenario, **fields: JsonValue) -> str: """A team the test deletes itself; cleanup only removes it if the test left it behind.""" created: Final = scenario.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}", **fields}) team_id: Final = string_value(created["team_id"]) - scenario.cleanups.callback(_remove_team_if_present, scenario.gateway, team_id) + scenario.cleanups.callback(_remove_team_by_sql, team_id) return team_id @@ -783,7 +776,9 @@ def test_deleting_a_team_with_more_keys_than_the_publish_cap_revokes_every_key_o keys: Final = _seeded_team_keys(team, KEYS_WELL_PAST_THE_PUBLISH_CAP) assert _statuses(peer, keys) == {200: len(keys)} - response: Final = _delete_teams(gateway, [team]) + response: Final = _post_with_timeout( + gateway, "/team/delete", {"team_ids": [team]}, timeout=2 * TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + ) assert response.status_code == 200, response.text assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,)) == [] diff --git a/tests/integration/management/test_team_delete_member_add_race.py b/tests/integration/management/test_team_delete_member_add_race.py index e1f2a933ec3..c558264a650 100644 --- a/tests/integration/management/test_team_delete_member_add_race.py +++ b/tests/integration/management/test_team_delete_member_add_race.py @@ -22,8 +22,12 @@ import time import uuid from contextlib import asynccontextmanager from datetime import timedelta +from typing import Final from unittest.mock import MagicMock +from psycopg import AsyncConnection +from psycopg.rows import dict_row + import pytest from fastapi import HTTPException @@ -48,7 +52,7 @@ _HELD_LOCK_KEY_SQL = ( _LOCK_WAITER_SQL = ( "SELECT count(*)::int AS waiters FROM pg_locks " "WHERE locktype = 'advisory' AND NOT granted " - "AND classid::bigint = $1 AND objid::bigint = $2" + "AND classid::bigint = %s AND objid::bigint = %s" ) _LOCK_WAIT_TIMEOUT_SECONDS = 20.0 _LOCK_POLL_SECONDS = 0.01 @@ -66,21 +70,17 @@ async def _hold_team_lock(held, team_id: str) -> tuple[int, int]: return rows[0]["classid"], rows[0]["objid"] -async def _await_lock_contention(watcher, lock_key: tuple[int, int], task, what: str) -> None: - """Block until Postgres reports `task` queued behind the held lock. - - This is the assertion that the endpoint serializes on the team's advisory lock, and it - is what a fixed sleep was standing in for: the endpoint is only provably waiting once a - non-granted advisory lock on the same key exists. `watcher` must be a connection that is - not itself blocked, so it can observe the queue.""" - classid, objid = lock_key - deadline = time.monotonic() + _LOCK_WAIT_TIMEOUT_SECONDS - while time.monotonic() < deadline: - if task.done(): - raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception() - if (await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid))[0]["waiters"]: - return - await asyncio.sleep(_LOCK_POLL_SECONDS) +async def _await_lock_contention(lock_key: tuple[int, int], task: asyncio.Task[object], what: str) -> None: + deadline: Final = time.monotonic() + _LOCK_WAIT_TIMEOUT_SECONDS + async with await AsyncConnection.connect(os.environ["DATABASE_URL"], autocommit=True, row_factory=dict_row) as watcher: + while time.monotonic() < deadline: + if task.done(): + raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception() + cursor: Final = await watcher.execute(_LOCK_WAITER_SQL, lock_key) + observed: Final = await cursor.fetchone() + if observed is not None and observed["waiters"]: + return + await asyncio.sleep(_LOCK_POLL_SECONDS) raise AssertionError(f"{what} never queued on the team's advisory lock within {_LOCK_WAIT_TIMEOUT_SECONDS}s") @@ -171,7 +171,7 @@ async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): async with blocker.tx(timeout=timedelta(seconds=30)) as held: lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(add_member()) - await _await_lock_contention(db, lock_key, task, "member_add") + await _await_lock_contention(lock_key, task, "member_add") await held.execute_raw(_DELETE_TEAM, team_id) @@ -234,7 +234,7 @@ async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster async with blocker.tx(timeout=timedelta(seconds=30)) as held: lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await _await_lock_contention(db, lock_key, task, "member_delete") + await _await_lock_contention(lock_key, task, "member_delete") await held.litellm_teamtable.update( where={"team_id": team_id}, @@ -311,7 +311,7 @@ async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): async with blocker.tx(timeout=timedelta(seconds=30)) as held: lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await _await_lock_contention(db, lock_key, task, "delete_team") + await _await_lock_contention(lock_key, task, "delete_team") await held.litellm_usertable.upsert( where={"user_id": user_id}, diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index 2b05d35d553..bb5766e5816 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -193,8 +193,10 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( timings["s6_customer_update"] = _timed_post( candidate, "/customer/update", {"user_id": customer_id, "max_budget": 6.0} ) - assert timings["s6_customer_update"] < _HANDLER_BUDGET_SECONDS, ( - f"/customer/update took {timings['s6_customer_update']:.3f}s with a wedged coordination Redis" + assert timings["s6_customer_update"] < _SHORTEST_REDIS_WAIT_SECONDS, ( + f"/customer/update took {timings['s6_customer_update']:.3f}s, reaching the " + f"{_SHORTEST_REDIS_WAIT_SECONDS}s shortest wait the proxy puts on Redis, so it waited on the " + "wedged coordination Redis" ) timings["s7_reset_spend"] = _timed_post(candidate, f"/key/{key}/reset_spend", {"reset_to": 0}) assert timings["s7_reset_spend"] < _HANDLER_BUDGET_SECONDS, ( diff --git a/tests/integration/management/test_vector_store_file_list_managed_ids.py b/tests/integration/management/test_vector_store_file_list_managed_ids.py index 16d33ee83af..5ef709d7365 100644 --- a/tests/integration/management/test_vector_store_file_list_managed_ids.py +++ b/tests/integration/management/test_vector_store_file_list_managed_ids.py @@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from functools import partial +from itertools import count from pathlib import Path from queue import SimpleQueue from types import MappingProxyType @@ -1042,6 +1043,77 @@ def test_multi_model_upload_maps_only_the_provider_id_the_managed_id_carries(gat assert second_ids == ((managed_a,) if carried == second.file_id("a.txt") else (second.file_id("a.txt"),)) +def _team_wildcard(scenario: Scenario, pattern: str, api_base: str, bearer: str, team: str) -> str: + identity: Final = f"wildcard-{uuid.uuid4().hex}" + scenario.gateway.post( + "/model/new", + { + "model_name": pattern, + "litellm_params": {"model": "openai/*", "api_key": bearer, "api_base": api_base}, + "model_info": {"id": identity, "team_id": team}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, identity) + eventually(lambda: _every_worker_has_deployment(scenario.gateway, identity), lambda served: served, seconds=90) + return identity + + +def test_two_targets_served_by_one_team_wildcard_keep_one_upload_that_delete_removes(gateway: Gateway) -> None: + store: Final = "vs_" + uuid.uuid4().hex + bearer: Final = "provider-key-" + uuid.uuid4().hex[:8] + prefix: Final = f"wc-{uuid.uuid4().hex}" + upload_numbers: Final = count(1) + first_upload: Final = f"file-{prefix}-1" + second_upload: Final = f"file-{prefix}-2" + served: Final = _provider(store, _constant_listing(store, first_upload)) + + def respond(request: Request) -> Reply: + if (request.method, urlsplit(request.target).path) != ("POST", "/v1/files"): + return served(request) + _assert_provider_upload(request, bearer, "a.txt") + return _json_reply(_file_object(f"file-{prefix}-{next(upload_numbers)}")) + + with gateway.scenario() as scenario, wire_server(respond) as wire: + team: Final = scenario.team(models=[]) + _team_wildcard(scenario, f"{prefix}/*", f"{wire.url}/v1", bearer, team) + key: Final = scenario.key(team_id=team, user_id=scenario.member(team)) + wire.drain() + + managed: Final = _upload(gateway, key, f"{prefix}/gpt-one,{prefix}/gpt-two", "a.txt") + assert _carried_provider_file_id(managed) == first_upload, _decoded(managed) + assert [row["flat_model_file_ids"] for row in read_rows(MANAGED_FILE_ROW, (managed,))] == [[first_upload]] + eventually(lambda: wire.received.qsize(), lambda received: received >= 3) + create_requests: Final = wire.drain() + assert [(request.method, request.target) for request in create_requests] == [ + ("POST", "/v1/files"), + ("POST", "/v1/files"), + ("DELETE", f"/v1/files/{second_upload}"), + ], create_requests + assert create_requests[2].headers["authorization"] == f"Bearer {bearer}", create_requests[2].headers + + attached: Final = gateway.request( + "POST", f"/v1/vector_stores/{store}/files", {"file_id": managed}, key=key, params={"model": f"{prefix}/*"} + ) + assert attached.status_code == 200, attached.text + listed: Final = gateway.request( + "GET", f"/v1/vector_stores/{store}/files", key=key, params={"model": f"{prefix}/*"} + ) + assert _ids(_listed(listed)) == (managed,), listed.text + attach_bodies: Final = [ + JSON_OBJECT.validate_json(request.body) for request in wire.drain() if request.method == "POST" + ] + assert attach_bodies == [{"file_id": first_upload}], attach_bodies + + deleted: Final = gateway.request("DELETE", f"/v1/files/{managed}", key=key) + assert deleted.status_code == 200, deleted.text + assert read_rows(MANAGED_FILE_ROW, (managed,)) == [] + delete_requests: Final = wire.drain() + assert [(request.method, request.target) for request in delete_requests] == [ + ("DELETE", f"/v1/files/{first_upload}") + ], delete_requests + assert delete_requests[0].headers["authorization"] == f"Bearer {bearer}", delete_requests[0].headers + + def test_deleting_the_managed_file_makes_its_listing_raw_again(gateway: Gateway) -> None: with _rig(gateway, "a.txt") as rig: member: Final = _member(rig.scenario, rig.model) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py index e159573e9fd..1f57584572b 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py @@ -97,9 +97,10 @@ def test_shared_gateway_route_is_disabled_without_reaching_provider(gateway: Gat headers=_cli_bearer_headers(gateway.key), ) assert response.status_code == 404, response.text - assert _JSON_OBJECT.validate_json(response.content) == {"detail": "Claude Code gateway is not enabled"}, ( - response.text - ) + assert _JSON_OBJECT.validate_json(response.content) == { + "type": "error", + "error": {"type": "not_found_error", "message": "Claude Code gateway is not enabled"}, + }, response.text assert wire.drain() == () @@ -268,15 +269,15 @@ def test_enabled_gateway_count_tokens_alias_matches_native_count_tokens_route(en ) assert restricted_response.status_code == 403, restricted_response.text assert _JSON_OBJECT.validate_json(restricted_response.content) == { + "type": "error", "error": { + "type": "permission_error", "message": ( f"The requested model '{restricted_model}' is not available for this API key, or the model name " "is invalid. Check the models available to you and try again." ), - "type": "key_model_access_denied", - "param": "model", - "code": "403", - } + "litellm_error_type": "key_model_access_denied", + }, }, restricted_response.text assert wire.drain() == () @@ -285,15 +286,15 @@ def test_enabled_gateway_count_tokens_alias_matches_native_count_tokens_route(en def test_enabled_gateway_refuses_model_outside_virtual_key(enabled_gateway: Gateway) -> None: result: Final = _restricted_model_responses(enabled_gateway) expected_body: Final = { + "type": "error", "error": { + "type": "permission_error", "message": ( f"The requested model '{result[4]}' is not available for this API key, or the model name is invalid. " "Check the models available to you and try again." ), - "type": "key_model_access_denied", - "param": "model", - "code": "403", - } + "litellm_error_type": "key_model_access_denied", + }, } assert result[0] == result[2] == 403, (result[1], result[3]) assert _JSON_OBJECT.validate_json(result[1]) == expected_body, result[1] diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_input_transformations_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_input_transformations_wire.py new file mode 100644 index 00000000000..6b9739def8c --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_input_transformations_wire.py @@ -0,0 +1,91 @@ +"""Anthropic's ``input_transformations`` report reaches the caller on chat completions and ``/v1/messages``.""" + +import json +from typing import Final + +import pytest +from integration._support import claude_code as cc +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-4-6" +_API_KEY: Final = "scripted-anthropic-key" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_CHOICES: Final = TypeAdapter(list[dict[str, JsonValue]]) +_INPUT_TRANSFORMATIONS: Final[list[JsonValue]] = [{"type": "prefix_binding_mismatch", "index": 1}] +_CONTENT: Final[tuple[dict[str, JsonValue], ...]] = ({"type": "text", "text": "Ok"},) +_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 8, "output_tokens": 2} +_IDENTITY: Final = "msg_input_transformations" + + +def _reply_body() -> bytes: + body: Final = _JSON_OBJECT.validate_json(cc.message_reply(_IDENTITY, _BACKEND, _CONTENT, _USAGE)) + return json.dumps({**body, "input_transformations": _INPUT_TRANSFORMATIONS}).encode() + + +def _stream_reply() -> tuple[bytes, ...]: + start, *rest = cc.message_stream(_IDENTITY, _BACKEND, _CONTENT, _USAGE) + event_name, data = cc.sse_events(start.decode())[0] + message: Final = _JSON_OBJECT.validate_python(data["message"]) + started: Final = cc.sse_frame( + event_name, {**data, "message": {**message, "input_transformations": _INPUT_TRANSFORMATIONS}} + ) + return (started, *rest) + + +def _anthropic_upstream(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == _API_KEY, request.headers + body: Final = _JSON_OBJECT.validate_json(request.body) + if body.get("stream") is True: + return Reply(content_type="text/event-stream", chunks=_stream_reply()) + return Reply(body=_reply_body()) + + +def test_chat_completion_message_carries_input_transformations(gateway: Gateway) -> None: + with wire_server(_anthropic_upstream) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "Say ok"}]}, + ) + assert response.status_code == 200, response.text + message: Final = response.json()["choices"][0]["message"] + assert message["provider_specific_fields"]["input_transformations"] == _INPUT_TRANSFORMATIONS, response.text + + +def test_first_chat_completion_chunk_carries_input_transformations(gateway: Gateway) -> None: + with wire_server(_anthropic_upstream) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 16, "stream": True, "messages": [{"role": "user", "content": "Say ok"}]}, + ) + assert response.status_code == 200, response.text + chunks: Final = [ + _JSON_OBJECT.validate_json(line[len("data: ") :]) + for line in response.text.splitlines() + if line.startswith("data: {") + ] + first_delta: Final = _JSON_OBJECT.validate_python(_CHOICES.validate_python(chunks[0]["choices"])[0]["delta"]) + assert first_delta["provider_specific_fields"] == {"input_transformations": _INPUT_TRANSFORMATIONS}, chunks[0] + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_messages_endpoint_keeps_input_transformations(gateway: Gateway, stream: bool) -> None: + with wire_server(_anthropic_upstream) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "stream": stream, "messages": [{"role": "user", "content": "Say ok"}]}, + ) + assert response.status_code == 200, response.text + if stream: + started: Final = next(data for name, data in cc.sse_events(response.text) if name == "message_start") + assert _JSON_OBJECT.validate_python(started["message"])["input_transformations"] == _INPUT_TRANSFORMATIONS + else: + assert response.json()["input_transformations"] == _INPUT_TRANSFORMATIONS, response.text diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py index 3ac1bdb0334..a1c5f6964bf 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py @@ -3,7 +3,7 @@ import json import re import time import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from pathlib import Path from types import MappingProxyType @@ -13,7 +13,7 @@ import anthropic import httpx import openai import pytest -from integration._support.client import Gateway, eventually, object_value +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.vertex import service_account_json @@ -417,14 +417,31 @@ def test_a_deployment_dropping_context_management_sends_neither_the_field_nor_it _spend_row(f"msg_{marker}") +@pytest.fixture(scope="module") +def local_beta_map_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with ( + gateway_from_environment() as shared, + owned_proxy( + shared, + tmp_path_factory.mktemp("safeguards-local-beta-map"), + {"LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "true"}, + ) as owned, + ): + yield owned + + @pytest.mark.parametrize("row", tuple(_FOUNDRY_SHAPES)) -def test_azure_ai_messages_forward_the_mapped_client_betas_beside_the_derived_one(gateway: Gateway, row: str) -> None: +def test_azure_ai_messages_forward_the_mapped_client_betas_beside_the_derived_one( + local_beta_map_gateway: Gateway, row: str +) -> None: case: Final = _FOUNDRY_SHAPES[row] marker: Final = uuid.uuid4().hex - with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + with wire_server(anthropic_peer) as wire, local_beta_map_gateway.scenario() as scenario: model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) response: Final = _post( - gateway, messages_body(model, marker, safeguards=case.sent, stream=case.stream), beta=case.client_beta + local_beta_map_gateway, + messages_body(model, marker, safeguards=case.sent, stream=case.stream), + beta=case.client_beta, ) if case.stream: _assert_streamed(response, marker) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py index 8eeaab66553..210b754f626 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py @@ -199,14 +199,14 @@ def test_utils_token_counter_adds_title_and_context_once_per_document( @pytest.mark.parametrize("label", list(_UNREADABLE)) -def test_utils_token_counter_prices_unreadable_pdf_bytes_like_one_image( +def test_utils_token_counter_uses_the_same_bounded_fallback_for_unreadable_pdf_bytes( gateway: Gateway, provider: SharedProvider, label: str ) -> None: data: Final = _UNREADABLE[label] as_pdf: Final = _local(gateway, provider, [document(base64_source(data))]) - as_png: Final = _local(gateway, provider, [document(base64_source(data, media_type="image/png"))]) + fallback: Final = _local(gateway, provider, [document(base64_source(encoded(b"unreadable-document")))]) without: Final = _local(gateway, provider, []) - assert as_pdf == as_png, (as_pdf, as_png) + assert as_pdf == fallback, (as_pdf, fallback) assert 0 <= as_pdf - without < rendered_tokens((LETTER,)), (as_pdf, without) diff --git a/tests/integration/messages_endpoint/test_beta_gated_fields_wire.py b/tests/integration/messages_endpoint/test_beta_gated_fields_wire.py index 26ad9f20897..fc1e91c79d7 100644 --- a/tests/integration/messages_endpoint/test_beta_gated_fields_wire.py +++ b/tests/integration/messages_endpoint/test_beta_gated_fields_wire.py @@ -16,7 +16,8 @@ pytestmark: Final = pytest.mark.shares_worker_state ANTHROPIC_MODEL: Final = "beta-gated-anthropic" FOUNDRY_MODEL: Final = "beta-gated-foundry" -FOUNDRY_DROP_PARAMS_MODEL: Final = "beta-gated-foundry-drop-params" +DATABRICKS_MODEL: Final = "beta-gated-databricks" +DATABRICKS_DROP_PARAMS_MODEL: Final = "beta-gated-databricks-drop-params" CLEAR_AT_BETA: Final = "mid-conversation-system-clear-at-2026-08-21" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _NONCE: Final = re.compile(rb"nonce-([0-9a-f]{32})") @@ -73,7 +74,15 @@ def write_config(directory: Path, wire_url: str) -> Path: deployment(ANTHROPIC_MODEL, "anthropic/claude-opus-4-8", f"{wire_url}/anthropic"), deployment(FOUNDRY_MODEL, "azure_ai/claude-opus-4-8", f"{wire_url}/foundry"), deployment( - FOUNDRY_DROP_PARAMS_MODEL, "azure_ai/claude-opus-4-8", f"{wire_url}/foundry", drop_params=True + DATABRICKS_MODEL, + "databricks/databricks-claude-opus-4-8", + f"{wire_url}/databricks/serving-endpoints", + ), + deployment( + DATABRICKS_DROP_PARAMS_MODEL, + "databricks/databricks-claude-opus-4-8", + f"{wire_url}/databricks/serving-endpoints", + drop_params=True, ), ], "general_settings": { @@ -150,7 +159,7 @@ def system_message(body: bytes) -> dict[str, JsonValue]: def test_a_drop_params_leg_whose_row_lacks_the_beta_sends_neither_the_beta_nor_clear_at(rig: Rig) -> None: request_nonce: Final = nonce() - status, raw = rig.send(FOUNDRY_DROP_PARAMS_MODEL, request_nonce) + status, raw = rig.send(DATABRICKS_DROP_PARAMS_MODEL, request_nonce) assert reply_text(status, raw) == f"answered {request_nonce.removeprefix('nonce-')}" upstream: Final = rig.received(request_nonce) assert CLEAR_AT_BETA not in beta_header(upstream.headers), upstream.headers @@ -159,7 +168,7 @@ def test_a_drop_params_leg_whose_row_lacks_the_beta_sends_neither_the_beta_nor_c def test_a_leg_whose_row_lacks_the_beta_forwards_clear_at_as_sent_without_drop_params(rig: Rig) -> None: request_nonce: Final = nonce() - status, raw = rig.send(FOUNDRY_MODEL, request_nonce) + status, raw = rig.send(DATABRICKS_MODEL, request_nonce) assert reply_text(status, raw) == f"answered {request_nonce.removeprefix('nonce-')}" upstream: Final = rig.received(request_nonce) assert CLEAR_AT_BETA not in beta_header(upstream.headers), upstream.headers @@ -170,9 +179,10 @@ def test_a_leg_whose_row_lacks_the_beta_forwards_clear_at_as_sent_without_drop_p } -def test_a_leg_that_carries_the_beta_forwards_clear_at_untouched(rig: Rig) -> None: +@pytest.mark.parametrize("model", (ANTHROPIC_MODEL, FOUNDRY_MODEL)) +def test_a_leg_that_carries_the_beta_forwards_clear_at_untouched(rig: Rig, model: str) -> None: request_nonce: Final = nonce() - status, raw = rig.send(ANTHROPIC_MODEL, request_nonce) + status, raw = rig.send(model, request_nonce) assert reply_text(status, raw) == f"answered {request_nonce.removeprefix('nonce-')}" upstream: Final = rig.received(request_nonce) assert CLEAR_AT_BETA in beta_header(upstream.headers), upstream.headers diff --git a/tests/integration/observability/_logging_only_scope_support.py b/tests/integration/observability/_logging_only_scope_support.py index 3d29f7ab8c5..3b14c502623 100644 --- a/tests/integration/observability/_logging_only_scope_support.py +++ b/tests/integration/observability/_logging_only_scope_support.py @@ -22,7 +22,7 @@ import yaml from anthropic import Anthropic, AsyncAnthropic from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa -from integration._support.client import Gateway, Scenario, eventually, object_value +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows, write_rows from integration._support.upstream import delete_scenario, register_scenario from integration._support.wire import Reply, Request, Wire @@ -798,6 +798,7 @@ def _insert_database_guardrail( @contextmanager def _database_guardrail( + gateway: Gateway, identity: str, policy_url: str, scope: str | None, @@ -809,11 +810,15 @@ def _database_guardrail( try: yield finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) -def _delete_database_guardrail(identity: str) -> None: - write_rows('DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', (identity,)) +def _delete_database_guardrail(gateway: Gateway, identity: str) -> None: + rows: Final = read_rows('SELECT guardrail_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', (identity,)) + responses: Final = tuple( + gateway.request("DELETE", f"/guardrails/{string_value(row['guardrail_id'])}") for row in rows + ) + assert all(response.status_code == 200 for response in responses), [response.text for response in responses] def _post_guardrail_body( diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index c52649c46df..3abb2a7e2be 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -253,6 +253,16 @@ def _count(samples: tuple, name: str, model_id: str) -> float: ) +def _deployment_states(rig: Rig) -> tuple[float, ...]: + return tuple( + sample.value + for sample in _metric_samples(rig.candidate, rig.model_name) + if sample.name == DEPLOYMENT_STATE + and sample.labels.get("model_id") == rig.deployment_id + and sample.labels.get("api_base") == rig.provider.url + "/v1" + ) + + def _populated_failures(samples: tuple, rig: Rig, api_provider: str) -> float: return float( sum( @@ -414,7 +424,7 @@ def test_cache_hit_post_call_reject_keeps_deployment_labels_anthropic(gateway: G raised = True assert raised, "cache-hit post_call guardrail did not reject /v1/messages" assert rig.provider.received.qsize() == 1, rig.provider.drain() - _expect_metrics(rig, 1, 0, api_provider="anthropic", pf_id="None") + _expect_metrics(rig, 1, 0, api_provider="anthropic") def test_cache_hit_post_call_reject_keeps_deployment_labels_responses(gateway: Gateway, tmp_path: Path) -> None: @@ -431,7 +441,7 @@ def test_cache_hit_post_call_reject_keeps_deployment_labels_responses(gateway: G ) assert reject.status_code == 400, reject.text assert rig.provider.received.qsize() == 1, rig.provider.drain() - _expect_metrics(rig, 1, 0, pf_id="None") + _expect_metrics(rig, 1, 0) _assert_spend(_call_id(reject), rig) @@ -558,7 +568,8 @@ def test_cache_hit_post_call_guardrail_outage_keeps_deployment_labels( ) assert reject.status_code >= 400, reject.text assert rig.provider.received.qsize() == 1, rig.provider.drain() - _expect_metrics(rig, 0, 0, pf_populated=1, pf_blank=0) + _expect_metrics(rig, 1, 0) + _assert_spend(_call_id(reject), rig) def test_two_identical_cache_hit_rejects_increment_populated_series(gateway: Gateway, tmp_path: Path) -> None: @@ -635,6 +646,7 @@ def test_cache_hit_reject_deployment_metric_set_diff(gateway: Gateway, tmp_path: text: Final = "cache hit control e4 " + marker warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) assert warm.status_code == 200, warm.text + eventually(lambda: _deployment_states(rig), lambda states: states == (0.0,)) reject: Final = rig.candidate.request( "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) ) @@ -642,11 +654,4 @@ def test_cache_hit_reject_deployment_metric_set_diff(gateway: Gateway, tmp_path: samples: Final = _expect_metrics(rig, 1, 0) blank: Final = tuple(sample for sample in samples if sample.labels.get("model_id") == "") assert blank == (), blank - states: Final = tuple( - sample.value - for sample in samples - if sample.name == DEPLOYMENT_STATE - and sample.labels.get("model_id") == rig.deployment_id - and sample.labels.get("api_base") == "" - ) - assert states == (1.0,), states + eventually(lambda: _deployment_states(rig), lambda states: states == (1.0,)) diff --git a/tests/integration/observability/test_langfuse_otel_metadata_chaos.py b/tests/integration/observability/test_langfuse_otel_metadata_chaos.py index bb386f6d41e..15bc076b8ba 100644 --- a/tests/integration/observability/test_langfuse_otel_metadata_chaos.py +++ b/tests/integration/observability/test_langfuse_otel_metadata_chaos.py @@ -107,9 +107,11 @@ def test_langfuse_otel_sink_outage_mid_burst_lands_every_response_id_once(gatewa ) payloads: Final = tuple(json.loads(line) for line in events) found: Final = next( - payload.get("response_id") or (payload.get("message") or {}).get("id") or payload.get("id") + (payload.get("response") or {}).get("id") + or (payload.get("message") or {}).get("id") + or payload.get("id") for payload in payloads - if payload.get("response_id") or payload.get("message") or payload.get("id") + if payload.get("response") or payload.get("message") or payload.get("id") ) return str(found) response: Final = rig.candidate.client.post(path, json=body, headers={"Authorization": f"Bearer {key}"}) diff --git a/tests/integration/observability/test_langfuse_otel_metadata_clients.py b/tests/integration/observability/test_langfuse_otel_metadata_clients.py index c2e0573b48a..fd2b9758449 100644 --- a/tests/integration/observability/test_langfuse_otel_metadata_clients.py +++ b/tests/integration/observability/test_langfuse_otel_metadata_clients.py @@ -162,8 +162,13 @@ def test_langfuse_otel_metadata_reaches_langfuse_for_responses_sse_stream(gatewa for line in stream.iter_lines() if line.startswith("data: ") and line != "data: [DONE]" ) - response_id: Final = next( - event["response_id"] for event in events if event.get("type") == "response.output_text.delta" + advertised: Final = frozenset( + event["response"]["id"] + for event in events + if event.get("type") in ("response.created", "response.completed") ) + assert len(advertised) == 1, events + response_id: Final = next(iter(advertised)) attrs: Final = exported_llm_span(rig.collector, response_id) + assert attrs["gen_ai.response.id"] == response_id, attrs assert_identity(attrs, chat_identity_fields(identity)) diff --git a/tests/integration/observability/test_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index 0841c91f8ca..5a37e779ef3 100644 --- a/tests/integration/observability/test_logging_only_scope_config.py +++ b/tests/integration/observability/test_logging_only_scope_config.py @@ -35,6 +35,7 @@ from _logging_only_scope_support import ( _reset_upstream, _response_body_without_ids, _spend_row_for_call_id, + _spend_row_for_response_id, _spend_rows, _upstream_requests, wire_server, @@ -115,7 +116,7 @@ def test_G7_database_invalid_scope_reads_preserve_pre_call_blocking(gateway: Gat assert blocked_response.status_code == 400, blocked_response.text assert _upstream_requests(gateway.upstream_url) == () finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) finally: delete_scenario(upstream_handle) @@ -404,6 +405,7 @@ def test_presidio_scope_matches_no_scope_behavior(gateway: Gateway, tmp_path: Pa bool, seconds=30, ) + _spend_row_for_response_id(no_scope_result.response_id) no_scope_analyzer_calls: Final = no_scope_analyzer.drain() no_scope_anonymizer_calls: Final = no_scope_anonymizer.drain() no_scope_call_counts: Final = ( @@ -431,6 +433,7 @@ def test_presidio_scope_matches_no_scope_behavior(gateway: Gateway, tmp_path: Pa bool, seconds=30, ) + _spend_row_for_response_id(scoped_result.response_id) scoped_analyzer_calls: Final = scoped_analyzer.drain() scoped_anonymizer_calls: Final = scoped_anonymizer.drain() scoped_call_counts: Final = ( @@ -678,7 +681,7 @@ def test_database_guardrail_load_keeps_pre_call_blocking( for entry in entries ) == ((identity, "pre_call", "guardrail_intervened"),), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) finally: delete_scenario(upstream_handle) @@ -799,7 +802,7 @@ def test_G6_database_guardrail_polling_normalizes_invalid_scope_without_reinitia ) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries) assert _upstream_requests(gateway.upstream_url) == () finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) finally: delete_scenario(upstream_handle) @@ -952,7 +955,7 @@ def test_management_post_validates_logging_only_scope( assert isinstance(after_list, list), after_list assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_H5_management_put_rejection_preserves_database_and_runtime(gateway: Gateway, tmp_path: Path) -> None: @@ -1074,7 +1077,7 @@ def test_H5_management_put_rejection_preserves_database_and_runtime(gateway: Gat for entry in entries ) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries) finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1153,7 +1156,7 @@ def test_H6_management_patch_clears_scope_when_switching_to_blocking_only_mode( for entry in entries ) == ((identity, ("pre_call",), "guardrail_intervened"),), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1240,7 +1243,7 @@ def test_H7_management_patch_rejection_preserves_pre_call_guardrail(gateway: Gat for entry in entries ) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries) finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1331,7 +1334,7 @@ def test_management_patch_default_on_heals_stored_scope( for entry in entries ) == ((identity, "pre_call", "guardrail_intervened"),), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) finally: delete_scenario(upstream_handle) @@ -1429,7 +1432,7 @@ def test_H10_management_patch_null_scope_restores_both_directions(gateway: Gatew (identity, "logging_only", "guardrail_intervened"), ), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1517,7 +1520,7 @@ def test_H11_management_patch_same_scope_is_idempotent_and_output_only(gateway: (entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"]) for entry in entries ) == tuple((identity, "logging_only", "guardrail_intervened") for _ in expected_directions), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1550,7 +1553,7 @@ def test_H12_management_reads_expose_typed_logging_only_scope(gateway: Gateway, assert object_value(info["litellm_params"])["logging_only_scope"] == "input", info assert object_value(matches[0]["litellm_params"])["logging_only_scope"] == "input", matches finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_H13_unauthenticated_management_and_chat_requests_do_not_scan(gateway: Gateway, tmp_path: Path) -> None: @@ -1629,7 +1632,7 @@ def test_K1_post_rejects_continue_flag_outside_logging_only(gateway: Gateway, tm assert after_list == before_list, (before_list, after_list) assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_path: Path) -> None: @@ -1675,7 +1678,7 @@ def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_ } == {key: after_params[key] for key in after_params if key != "api_key"}, (before_params, after_params) assert str(after_params["api_key"]).startswith("litellm_enc::"), after_params finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway, tmp_path: Path) -> None: @@ -1713,7 +1716,7 @@ def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway assert params.get("logging_only_continue_on_input_failure") is None, persisted assert params["mode"] == ["pre_call"], persisted finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Gateway, tmp_path: Path) -> None: @@ -1759,5 +1762,5 @@ def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Ga log_text: Final = owned.log.read_text() assert "Ignoring logging_only_continue_on_input_failure" in log_text, log_text finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) diff --git a/tests/integration/observability/test_logging_only_scope_runtime.py b/tests/integration/observability/test_logging_only_scope_runtime.py index 8da26f20eae..271d8d0e24d 100644 --- a/tests/integration/observability/test_logging_only_scope_runtime.py +++ b/tests/integration/observability/test_logging_only_scope_runtime.py @@ -1098,7 +1098,7 @@ def test_X2_scope_patch_toggles_during_concurrent_requests_keep_one_scan_per_res with wire_server(policy) as edge: config: Final = _empty_proxy_configuration(tmp_path, identity) with ( - _database_guardrail(identity, edge.url, "output", mode="logging_only"), + _database_guardrail(gateway, identity, edge.url, "output", mode="logging_only"), owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate, ): guardrails: Final = candidate.get("/v2/guardrails/list")["guardrails"] diff --git a/tests/integration/observability/test_otel_conversation_id.py b/tests/integration/observability/test_otel_conversation_id.py index d7f4de2961b..cd095991e53 100644 --- a/tests/integration/observability/test_otel_conversation_id.py +++ b/tests/integration/observability/test_otel_conversation_id.py @@ -1,5 +1,4 @@ import asyncio -import base64 import json import re import signal @@ -118,17 +117,6 @@ def _responses_reply(identity: str, stream: bool) -> Reply: ) -def _decoded_responses_id(identity: str) -> str: - try: - return base64.b64decode(identity.removeprefix("resp_").encode()).decode() - except (ValueError, UnicodeDecodeError): - return identity - - -def _canonical_id(identity: str) -> str: - return _decoded_responses_id(identity).rpartition("response_id:")[2] - - def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: return tuple( json.loads(line[6:]) for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]" @@ -170,8 +158,7 @@ class Collector: return tuple( attributes for attributes in self.attributes() - if isinstance(logged := attributes.get("gen_ai.response.id", {}).get("stringValue"), str) - and _canonical_id(logged) == _canonical_id(response_id) + if attributes.get("gen_ai.response.id", {}).get("stringValue") == response_id ) def conversation_ids(self, response_id: str) -> tuple[str | None, ...]: @@ -183,10 +170,6 @@ class Collector: def single_span(self, response_id: str) -> str | None: return eventually(lambda: self.conversation_ids(response_id), lambda values: len(values) == 1, seconds=30)[0] - def logged_id(self, response_id: str) -> str: - spans: Final = eventually(lambda: self.spans(response_id), lambda values: len(values) == 1, seconds=30) - return str(spans[0]["gen_ai.response.id"]["stringValue"]) - @pytest.fixture(scope="session") def collector() -> Iterator[Collector]: @@ -437,12 +420,13 @@ def test_responses_stream_raw_metadata_session_id_lands_as_conversation_id(rig: completed: Final = tuple(event for event in events if event["type"] == "response.completed") assert len(completed) == 1, events assert completed[0]["response"]["output"][0]["id"] == f"msg_resp_{marker}", completed - assert str(completed[0]["response"]["id"]).startswith("resp_"), completed + client_id: Final = str(completed[0]["response"]["id"]) + assert client_id.startswith("resp_"), completed assert rig.upstream_bodies(marker) == ( {"model": "gpt-4o-mini", "input": marker, "metadata": {"session_id": session}, "stream": True}, ) - assert rig.sink.single_span(f"resp_{marker}") == session - assert rig.spend_session(rig.sink.logged_id(f"resp_{marker}")) == session + assert rig.sink.single_span(client_id) == session + assert rig.spend_session(client_id) == session @pytest.mark.covers("other.observability.otel.conversation_id_from_metadata_chat_raw") @@ -689,15 +673,11 @@ def _burst(rig: Rig, count: int, session_for: Mapping[int, str]) -> tuple[tuple[ return tuple(pool.map(one, range(count))) -def _is_encrypted_responses_id(identity: str) -> bool: - return identity.startswith("resp_") and _decoded_responses_id(identity) == identity - - def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]]: spans: Final = rig.sink.attributes() return { session: tuple( - _canonical_id(str(attributes["gen_ai.response.id"]["stringValue"])) + str(attributes["gen_ai.response.id"]["stringValue"]) for attributes in spans if attributes.get(CONVERSATION, {}).get("stringValue") == session and "gen_ai.response.id" in attributes ) @@ -707,21 +687,13 @@ def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]] def _assert_exactly_once(rig: Rig, expected: Mapping[str, str], landed: Mapping[str, tuple[str, ...]]) -> None: spend: Final = eventually( - lambda: { - session: tuple(_canonical_id(identity) for identity in rig.spend_request_ids(session)) - for session in expected.values() - }, + lambda: {session: rig.spend_request_ids(session) for session in expected.values()}, lambda rows: all(len(values) >= 1 for values in rows.values()), seconds=70, ) assert landed == spend, (landed, spend) assert all(len(values) == 1 for values in landed.values()), landed - caller_visible: Final = { - session: (_canonical_id(identity),) - for identity, session in expected.items() - if not _is_encrypted_responses_id(identity) - } - assert {session: landed[session] for session in caller_visible} == caller_visible, landed + assert landed == {session: (identity,) for identity, session in expected.items()}, landed @pytest.mark.covers("other.observability.otel.conversation_id_sink_outage_recovers_exactly_once") diff --git a/tests/integration/observability/test_prometheus_provider_auth_failure.py b/tests/integration/observability/test_prometheus_provider_auth_failure.py index 13d264ca6b0..de1f19bc523 100644 --- a/tests/integration/observability/test_prometheus_provider_auth_failure.py +++ b/tests/integration/observability/test_prometheus_provider_auth_failure.py @@ -995,7 +995,7 @@ def _gated(gateway: Gateway, cell: str, key: str | None) -> httpx.Response: def recorded_before(cell: str) -> list[tuple[str, str, str, float]]: - if cell in ("invalid-key", "expired-key"): + if cell in ("invalid-key", "missing-key", "expired-key", "blocked-key"): return [] return sorted([(FAILURES, "", "", 1.0), (TOTALS, "", "", 1.0), (LLM_FAILED, identity(cell), "", 1.0)]) @@ -1020,7 +1020,9 @@ def test_a_request_the_proxy_rejects_itself_never_marks_the_deployment_down_and_ assert response.status_code == (403 if cell == "no-model-access" else 401), response.text assert rig.upstream.drain() == () samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert values(samples, PROXY_FAILED, requested_model=name(cell)) == [1.0], recorded(samples) assert deployment_rows(samples, cell) == recorded_before(cell), recorded(samples, requested_model=name(cell)) + assert values(samples, STATE, model_id=identity(cell)) == [], recorded(samples) def test_a_key_rate_limit_rejected_at_the_proxy_is_still_counted_as_before(rig: Rig) -> None: diff --git a/tests/integration/observability/test_signoz_delivery.py b/tests/integration/observability/test_signoz_delivery.py index e25126e4fa8..93539d7641b 100644 --- a/tests/integration/observability/test_signoz_delivery.py +++ b/tests/integration/observability/test_signoz_delivery.py @@ -1,5 +1,4 @@ import asyncio -import base64 import json import re import signal @@ -9,7 +8,6 @@ from collections import deque from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from functools import partial from pathlib import Path from typing import Final @@ -137,17 +135,6 @@ def _upstream(request: Request) -> Reply: return _chat_reply(f"chatcmpl-{marker}", stream) -def _decoded_responses_id(identity: str) -> str: - try: - return base64.b64decode(identity.removeprefix("resp_").encode()).decode() - except (ValueError, UnicodeDecodeError): - return identity - - -def _canonical_id(identity: str) -> str: - return _decoded_responses_id(identity).rpartition("response_id:")[2] - - def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: return tuple( object_value(JSON.validate_json(line[6:])) @@ -232,8 +219,7 @@ class Collector: return tuple( span for span in self.spans() - if RESPONSE_ID in span.attributes - and _canonical_id(span.attributes[RESPONSE_ID]) == _canonical_id(response_id) + if RESPONSE_ID in span.attributes and span.attributes[RESPONSE_ID] == response_id ) def single_span(self, response_id: str, *, elsewhere: "Collector | None" = None) -> Span: @@ -249,11 +235,8 @@ class Collector: def landed(self, response_ids: Sequence[str]) -> dict[str, int]: spans: Final = self.spans() return { - _canonical_id(identity): sum( - 1 - for span in spans - if RESPONSE_ID in span.attributes - and _canonical_id(span.attributes[RESPONSE_ID]) == _canonical_id(identity) + identity: sum( + 1 for span in spans if RESPONSE_ID in span.attributes and span.attributes[RESPONSE_ID] == identity ) for identity in response_ids } @@ -448,7 +431,7 @@ def _assert_operator_span(rig: Rig, response_id: str, marker: str) -> Span: span: Final = rig.sink.single_span(response_id) assert span.target == "/v1/traces", span assert span.ingestion_key == OPERATOR_KEY, span - assert rig.tenant_sink.landed((response_id,)) == {_canonical_id(response_id): 0} + assert rig.tenant_sink.landed((response_id,)) == {response_id: 0} bodies: Final = rig.upstream_bodies(marker) assert len(bodies) == 1, bodies assert "signoz" not in json.dumps(bodies[0]).replace(marker, ""), bodies[0] @@ -522,7 +505,7 @@ def test_responses_stream_raw_httpx_span_lands_once_after_the_stream_is_consumed marker: Final = _marker() response: Final = rig.proxy.request("POST", "/v1/responses", {"model": rig.model, "input": marker, "stream": True}) assert response.status_code == 200, response.text - _assert_operator_span(rig, _responses_id(response, marker), marker) + _assert_operator_span(rig, _responses_id(response), marker) def test_v2_flag_on_still_delivers_the_operator_span_with_the_ingestion_key(v2_rig: Rig) -> None: @@ -821,7 +804,7 @@ def _chat_id(response: httpx.Response) -> str: return next(iter(identities)) -def _responses_id(response: httpx.Response, marker: str) -> str: +def _responses_id(response: httpx.Response) -> str: if not _is_event_stream(response): return _body_id(response) completed: Final = tuple( @@ -830,7 +813,7 @@ def _responses_id(response: httpx.Response, marker: str) -> str: if event.get("type") == "response.completed" ) assert len(completed) == 1 and completed[0].startswith("resp_"), response.text - return f"resp_{marker}" + return completed[0] def _message_id(response: httpx.Response) -> str: @@ -852,7 +835,7 @@ def _burst(rig: Rig, count: int) -> tuple[tuple[int, str, str | None], ...]: stream: Final = index % 2 == 0 path, body, identity_of = ( ("/v1/chat/completions", {"model": rig.model, "messages": [{"role": "user", "content": marker}]}, _chat_id), - ("/v1/responses", {"model": rig.model, "input": marker}, partial(_responses_id, marker=marker)), + ("/v1/responses", {"model": rig.model, "input": marker}, _responses_id), ( "/v1/messages", {"model": rig.model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, @@ -874,8 +857,8 @@ def _assert_exactly_once(rig: Rig, identities: Sequence[str]) -> None: landed: Final = eventually( lambda: rig.sink.landed(identities), lambda seen: all(count >= 1 for count in seen.values()), seconds=80 ) - assert landed == {_canonical_id(identity): 1 for identity in identities}, landed - charged: Final = frozenset(identity for identity in identities if not identity.startswith("resp_")) + assert landed == {identity: 1 for identity in identities}, landed + charged: Final = frozenset(identities) spend: Final = eventually( lambda: read_rows( 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])', diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py index 0aa3a3bcba2..a91f5cba1ad 100644 --- a/tests/integration/pricing/test_price_data_reload_prices.py +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -284,7 +284,7 @@ def _charged_as_listed(proxy: Gateway, deployment: _Deployment, scale: float) -> assert _listed_rates(proxy, deployment.name) == _expected_listing(rates), deployment.name request_id, charged = _billed(proxy.client, proxy.key, deployment) expected: Final = _expected_cost(rates, deployment.usage) - assert charged == (None if rates is None else pytest.approx(expected, rel=1e-6)), deployment.name + assert charged == pytest.approx(expected, rel=1e-6), deployment.name return _Charge(request_id=request_id, expected=expected) diff --git a/tests/integration/providers/_mantle_gpt_prompt_cache_support.py b/tests/integration/providers/_mantle_gpt_prompt_cache_support.py index 8f497fea74f..96261bab8b2 100644 --- a/tests/integration/providers/_mantle_gpt_prompt_cache_support.py +++ b/tests/integration/providers/_mantle_gpt_prompt_cache_support.py @@ -475,10 +475,6 @@ def row_key(request_id: str) -> str: return upstream_id_of(request_id.split("_cache_hit", 1)[0]) -def caller_sees_upstream_id(endpoint: Endpoint, *, stream: bool) -> bool: - return (endpoint, stream) != ("messages", True) - - def usage_of(payload: Mapping[str, JsonValue]) -> dict[str, JsonValue]: usage: Final = payload.get("usage") return JSON.validate_python(usage) if isinstance(usage, dict) else {} @@ -733,10 +729,9 @@ def marked_cell(gateway: Gateway, endpoint: Endpoint, *, stream: bool) -> None: assert_answered(outcome, marker) assert_wire(received, expected_wire(GPT, prompt_text(marker), endpoint=endpoint, marked=True), streaming=stream) assert_usage(endpoint, outcome.usage) - row: Final = success_row(name, marker) + row: Final = success_row(name, marker, outcome.response_id) assert_priced_row(row, GPT) - if caller_sees_upstream_id(endpoint, stream=stream): - assert row_key(str(row["request_id"])) == outcome.response_id, (row, outcome.response_id) + assert row_key(str(row["request_id"])) == outcome.response_id, (row, outcome.response_id) if not stream: assert math.isclose(float(outcome.headers["x-litellm-response-cost"]), expected_spend(GPT), rel_tol=1e-9), ( outcome.headers @@ -901,7 +896,7 @@ def assert_marked_sdk_cell( only_received(wire), expected_wire(GPT, prompt_text(marker), endpoint=endpoint, marked=True), streaming=stream ) assert_usage(endpoint, outcome.usage) - assert_priced_row(success_row(name, marker), GPT) + assert_priced_row(success_row(name, marker, outcome.response_id), GPT) HOSTILE_OPTIONS: Final[dict[str, JsonValue]] = { diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py index c9a85545d8b..8774d590b37 100644 --- a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py @@ -201,20 +201,27 @@ def _marker() -> str: return uuid.uuid4().hex +def _bedrock_invoke_deployment( + scenario: Scenario, wire_url: str, model_info: Mapping[str, JsonValue] | None = None +) -> str: + return scenario.model( + model=f"bedrock/invoke/{BEDROCK_MODEL}", + model_info=model_info, + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + ) + + def _deployment(scenario: Scenario, provider: Provider, rig: Rig) -> str: wire_url: Final = rig.wire.url match provider: case "anthropic": return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key=_ANTHROPIC_KEY) case "bedrock_invoke": - return scenario.model( - model=f"bedrock/invoke/{BEDROCK_MODEL}", - api_base=wire_url, - aws_access_key_id="AKIASCRIPTEDPROVIDER", - aws_secret_access_key="scripted-secret", - aws_region_name="us-east-1", - aws_bedrock_runtime_endpoint=wire_url, - ) + return _bedrock_invoke_deployment(scenario, wire_url) case "claude_platform": return scenario.model( model=f"bedrock/claude_platform/{MODEL}", @@ -1184,7 +1191,7 @@ def test_stream_cache_hit_skips_cached_provider_fields_that_are_not_an_object(ri assert tuple(_message_fields(event["response"]) for event in stale_logs) == ((),) * len(stale_logs), stale_logs -JsonModeRoute: TypeAlias = Literal["anthropic_native", "anthropic_tool", "bedrock_invoke"] +JsonModeRoute: TypeAlias = Literal["anthropic_native", "anthropic_tool", "bedrock_invoke_native", "bedrock_invoke_tool"] _UNMAPPED_MODEL: Final = "claude-scripted-9" _JSON_ANSWER: Final = json.dumps({"answer": DOCS_PARTIAL_TEXT}) @@ -1216,11 +1223,20 @@ _JSON_MODE_CASES: Final[Mapping[str, JsonModeCase]] = { "anthropic_tool_enveloped": JsonModeCase( "anthropic_tool", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" ), + "bedrock_invoke_native": JsonModeCase( + "bedrock_invoke_native", + "partial", + DOCS_PARTIAL_TEXT, + DOCS_PARTIAL_TEXT, + "content_filter", + "content_filter", + "native", + ), "bedrock_invoke_declared": JsonModeCase( - "bedrock_invoke", "jsontool", DECLARED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + "bedrock_invoke_tool", "jsontool", DECLARED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" ), "bedrock_invoke_enveloped": JsonModeCase( - "bedrock_invoke", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + "bedrock_invoke_tool", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" ), } @@ -1234,8 +1250,15 @@ def _json_mode_deployment(scenario: Scenario, route: JsonModeRoute, rig: Rig) -> scenario.model(model=f"anthropic/{_UNMAPPED_MODEL}", api_base=rig.wire.url, api_key=_ANTHROPIC_KEY), "anthropic", ) - case "bedrock_invoke": + case "bedrock_invoke_native": return _deployment(scenario, "bedrock_invoke", rig), "bedrock_invoke" + case "bedrock_invoke_tool": + return ( + _bedrock_invoke_deployment( + scenario, rig.wire.url, {"supports_bedrock_native_structured_output": False} + ), + "bedrock_invoke", + ) def _requested_format(body: Mapping[str, JsonValue]) -> tuple[bool, bool]: diff --git a/tests/integration/providers/test_azure_ai_claude_messages_wire.py b/tests/integration/providers/test_azure_ai_claude_messages_wire.py index c150de7520b..d39c6d2d2c1 100644 --- a/tests/integration/providers/test_azure_ai_claude_messages_wire.py +++ b/tests/integration/providers/test_azure_ai_claude_messages_wire.py @@ -1,10 +1,12 @@ +from collections.abc import Iterator from typing import Final import anthropic import pytest from anthropic.types import Message from integration._support import claude_code as cc -from integration._support.client import Gateway +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -144,8 +146,23 @@ def test_azure_ai_claude_messages_preserves_target_auth_beta_and_typed_response( assert _JSON_OBJECT.validate_json(provider_request.body) == expected_provider_body, provider_request.body +@pytest.fixture(scope="module") +def local_beta_map_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with ( + gateway_from_environment() as shared, + owned_proxy( + shared, + tmp_path_factory.mktemp("azure-ai-local-beta-map"), + {"LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "true"}, + ) as owned, + ): + yield owned + + @pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(gateway: Gateway, stream: bool) -> None: +def test_azure_ai_claude_messages_forwards_the_auto_injected_betas_azure_accepts( + local_beta_map_gateway: Gateway, stream: bool +) -> None: identity: Final = f"msg_azure_claude_beta_filter_{stream}" api_path: Final = "/anthropic/v1/messages" reply_body: Final = cc.message_reply(identity, _BACKEND, _CONTENT, _USAGE) @@ -172,7 +189,7 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g "context_management": context_management, } expected_provider_body: Final = {**request_body_template, "model": _BACKEND} - expected_beta: Final = "structured-outputs-2025-11-13" + expected_beta: Final = "compact-2026-01-12,structured-outputs-2025-11-13" def respond(request: Request) -> Reply: assert request.method == "POST" and request.target == api_path, request.target @@ -185,7 +202,7 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g return Reply(content_type="text/event-stream", chunks=stream_reply) return Reply(body=reply_body) - with wire_server(respond) as wire, gateway.scenario() as scenario: + with wire_server(respond) as wire, local_beta_map_gateway.scenario() as scenario: model: Final = scenario.model( model=f"azure_ai/{_BACKEND}", api_base=wire.url, @@ -193,8 +210,8 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g num_retries=0, ) with anthropic.Anthropic( - api_key=gateway.key, - base_url=str(gateway.client.base_url), + api_key=local_beta_map_gateway.key, + base_url=str(local_beta_map_gateway.client.base_url), max_retries=0, ) as client: if stream: diff --git a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py index 3baeebcd74a..1f3bded035b 100644 --- a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py +++ b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py @@ -547,7 +547,7 @@ def test_r21_responses_stream_input_item_without_content_after_an_assistant_item assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_ANSWERED], received (row,) = _success_rows(marker, 1) assert row["call_type"] == "aresponses" and row["end_user"] == marker, row - assert row["request_id"] == _inner_response_id(str(completed["response"]["id"])), (row, completed) + assert row["request_id"] == str(completed["response"]["id"]), (row, completed) def test_r22_passthrough_converse_forwards_a_message_without_content_verbatim(gateway: Gateway) -> None: diff --git a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py index a1432ffd7a0..19aa11a8e36 100644 --- a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py +++ b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py @@ -1,7 +1,6 @@ import asyncio import base64 import json -import os import re import threading import uuid @@ -26,8 +25,6 @@ from integration._support.upstream import _aws_event_frame, aws_event_stream_fra from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with - pytestmark: Final = pytest.mark.shares_worker_state _MODEL_ID: Final = "global.moonshotai.kimi-k3" @@ -43,7 +40,6 @@ _REJECTION: Final = "structured output schema uses unsupported regex negative lo _THROTTLED: Final = "Too many requests, please wait before trying again" _UNKNOWN_TYPE: Final = "somethingBedrockAddedLater" _UNKNOWN_ONLY: Final = "none of its 1 events carried a known event type" -_SIGNING_KEY: Final = os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt") _JSON_HEADERS: Final = MappingProxyType({":content-type": "application/json", ":message-type": "event"}) _USAGE: Final[dict[str, JsonValue]] = {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}} _RESPONSE: Final = json.dumps( @@ -277,17 +273,9 @@ def _responses_text(events: Iterable[dict[str, JsonValue]]) -> str: ) -def _inner_response_id(identity: str) -> str: - managed: Final = decrypt_if_encrypted_with(identity.removeprefix("resp_"), _SIGNING_KEY) - assert managed is not None, identity - issued: Final = managed.split(";", 1)[0].rsplit("response_id:", 1)[1] - decoded: Final = base64.b64decode(issued.removeprefix("resp_")).decode() - return decoded.rsplit("response_id:", 1)[1] - - def _completed_response_id(events: Iterable[dict[str, JsonValue]]) -> str: (completed,) = tuple(event for event in events if event.get("type") == "response.completed") - return _inner_response_id(str(_JSON.validate_python(completed["response"])["id"])) + return str(_JSON.validate_python(completed["response"])["id"]) def _only_received(wire: Wire) -> tuple[str, dict[str, JsonValue]]: diff --git a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py index 3d6eb7fcaff..0788b9ac85a 100644 --- a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py +++ b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py @@ -62,12 +62,12 @@ def _body(request: Request) -> dict[str, JsonValue]: # ciphertext the caller received, because the spend row id is read from response_obj["id"] before the # ResponsesIDSecurity hook rewrites it in place; the row is looked up under both ids until that ordering is fixed on # main -def _spend_row(client_id: str, issued_id: str) -> dict[str, JsonValue]: +def _spend_row(*request_ids: str) -> dict[str, JsonValue]: rows: Final = eventually( lambda: read_rows( 'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" ' "WHERE request_id = ANY(%s)", - ([client_id, issued_id],), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array + (list(request_ids),), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array ), lambda found: len(found) == 1, seconds=70, @@ -120,4 +120,4 @@ async def test_async_openai_sdk_responses_stream_is_served_by_the_native_respons assert issued.upstream == f"resp_upstream_{marker}", completed.response.id request: Final = _native_request(wire) assert _body(request) == {"model": GPT, "input": _prompt(marker), "stream": True}, request.body - assert _spend_row(completed.response.id, issued.issued) == _success_row(model) + assert _spend_row(completed.response.id) == _success_row(model) diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 4441582fb6b..05e08c5fc27 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -456,7 +456,7 @@ def _inner_response_id(identity: str) -> str: def _completed_response_id(events: Iterable[dict[str, JsonValue]]) -> str: (completed,) = tuple(event for event in events if event.get("type") == "response.completed") - return _inner_response_id(str(_JSON.validate_python(completed["response"])["id"])) + return str(_JSON.validate_python(completed["response"])["id"]) def _chat_text(body: Mapping[str, JsonValue]) -> str: @@ -911,7 +911,7 @@ async def test_k06_openai_sdk_async_responses_stream_reaches_kimi_without_a_cach received: Final = _only_received(rig.wire) assert received.model == _KIMI_US and received.streaming, received assert received.cache_points == 0, received.body - _success_row(_inner_response_id(completed[0].id)) + _success_row(completed[0].id) _LOCATIONS: Final = ( @@ -1251,17 +1251,19 @@ def test_s01_to_s06_an_odd_typed_flag_is_read_as_false_never_as_true( @pytest.mark.timeout(600) -def test_s10_s11_two_deployments_of_one_arn_share_the_flag_and_a_delete_leaves_it_in_place(gateway: Gateway) -> None: +def test_s10_s11_two_deployments_of_one_arn_keep_their_own_flags_and_a_delete_changes_neither( + gateway: Gateway, +) -> None: arn: Final = _arn("shared") with wire_server(_peer()) as wire, gateway.scenario() as scenario: flagged: Final = _deployment( gateway, scenario, wire, f"bedrock/{arn}", {"supports_prompt_cache_breakpoint": False}, cleanup=False ) plain: Final = _deployment(gateway, scenario, wire, f"bedrock/{arn}") - for name in (flagged, plain): + for name, cache_points in ((flagged, 0), (plain, 2)): outcome, received = _observe_at(gateway, wire, "chat", _chat_body(name, _prompt(), _SYSTEM_AND_USER)) _assert_answered(outcome) - assert received.cache_points == 0, (name, received.body) + assert received.cache_points == cache_points, (name, received.body) _success_row(outcome.response_id) flagged_identity: Final = _model_id(flagged) gateway.post("/model/delete", {"id": flagged_identity}) @@ -1276,7 +1278,7 @@ def test_s10_s11_two_deployments_of_one_arn_share_the_flag_and_a_delete_leaves_i wire.drain() after, after_received = _observe_at(gateway, wire, "chat", _chat_body(plain, _prompt(), _SYSTEM_AND_USER)) _assert_answered(after) - assert after_received.cache_points == 0, after_received.body + assert after_received.cache_points == 2, after_received.body _success_row(after.response_id) diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index d4f9533fcdc..3652e49fe8a 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -70,10 +70,18 @@ def _chat_messages(marker: str) -> list[JsonValue]: return [{"role": "user", "content": _prompt(marker)}] -def _outbound(backend: str, messages: list[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: +def _default_max_tokens(gateway: Gateway, backend: str) -> int: + match gateway.get("/public/litellm_model_cost_map").get(f"bedrock_mantle/{backend}"): + case {"max_output_tokens": int() as mapped}: + return mapped + case _: + return DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + + +def _outbound(gateway: Gateway, backend: str, messages: list[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: return { "messages": messages, - "max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + "max_tokens": _default_max_tokens(gateway, backend), "anthropic_version": _BRIDGE_VERSION, "model": backend, **fields, @@ -290,7 +298,7 @@ def test_chat_completion_on_a_mantle_claude_id_posts_native_messages_and_spends_ gateway: Gateway, backend: str ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{backend}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -314,7 +322,7 @@ def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_ gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[JsonValue, ...]] ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -355,7 +363,7 @@ def test_chat_request_fields_are_translated_to_the_native_messages_shape( outbound_fields: dict[str, JsonValue], ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], **outbound_fields) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], **outbound_fields) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -377,7 +385,7 @@ def test_chat_tools_reach_mantle_as_input_schema_and_the_tool_use_block_returns_ marker: Final = uuid4().hex backend: Final = _mantle_claude_where(gateway, "supports_forced_tool_use", disabled=False) expected: Final = _outbound( - backend, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], tool_choice=outbound_choice + gateway, backend, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], tool_choice=outbound_choice ) reply: Final = _tool_reply(marker, backend, "get_weather") with wire_server(_bearer_peer(expected, reply)) as wire, gateway.scenario() as scenario: @@ -407,7 +415,7 @@ def test_chat_tools_reach_mantle_as_input_schema_and_the_tool_use_block_returns_ def test_chat_stream_relays_a_native_tool_use_block_as_tool_call_chunks(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], stream=True) with wire_server(_bearer_peer(expected, _tool_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -439,7 +447,7 @@ def test_chat_tool_result_turn_reaches_mantle_as_tool_use_and_tool_result_blocks marker: Final = uuid4().hex tool_call: Final = f"toolu_{marker}" expected: Final = _outbound( - _HAIKU, + gateway, _HAIKU, [ _user_turn(marker), { @@ -484,7 +492,7 @@ def test_chat_json_schema_response_format_is_sent_as_a_forced_tool_and_returned_ ) -> None: marker: Final = uuid4().hex expected: Final = _outbound( - _HAIKU, + gateway, _HAIKU, [_user_turn(marker)], tools=[{"name": "json_tool_call", "input_schema": _CITY_SCHEMA}], tool_choice={"name": "json_tool_call", "type": "tool"}, @@ -537,7 +545,7 @@ def _sigv4_peer(expected: Mapping[str, JsonValue], reply: Reply) -> Callable[[Re def test_chat_completion_signs_the_native_messages_request_with_the_deployment_aws_keys(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_sigv4_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model( model=f"bedrock_mantle/{_HAIKU}", @@ -556,7 +564,7 @@ def test_chat_completion_signs_the_native_messages_request_with_the_deployment_a def test_region_prefix_in_a_mantle_claude_id_is_routing_only_and_keeps_the_mantle_price(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/us-east-2/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -577,7 +585,7 @@ def test_region_prefix_in_a_mantle_claude_id_is_routing_only_and_keeps_the_mantl def test_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) with OpenAI( @@ -606,7 +614,7 @@ def test_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> N def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) with OpenAI( @@ -645,7 +653,7 @@ def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: async def test_async_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) async with AsyncOpenAI( @@ -674,7 +682,7 @@ async def test_async_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: G async def test_async_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) async with AsyncOpenAI( @@ -714,7 +722,7 @@ def test_identical_chat_request_is_served_from_the_response_cache_and_logged_as_ gateway: Gateway, ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) request: Final[dict[str, JsonValue]] = {"model": model, "messages": _chat_messages(marker)} @@ -802,7 +810,7 @@ def test_open_weight_mantle_id_keeps_the_openai_compatible_chat_route_and_its_sa ) def test_api_base_with_a_route_suffix_still_posts_to_the_native_messages_path(gateway: Gateway, suffix: str) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url + suffix, api_key=_API_KEY) response: Final = gateway.request( @@ -814,7 +822,7 @@ def test_api_base_with_a_route_suffix_still_posts_to_the_native_messages_path(ga def test_responses_api_on_a_mantle_claude_id_bridges_to_native_messages(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": _prompt(marker)}) @@ -839,7 +847,7 @@ def test_responses_api_on_a_mantle_claude_id_bridges_to_native_messages(gateway: def test_responses_api_stream_on_a_mantle_claude_id_bridges_to_the_native_stream(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -996,7 +1004,7 @@ def test_unsupported_claude_param_is_dropped_under_drop_params( ) -> None: marker: Final = uuid4().hex backend: Final = _backend_without(gateway, flag) - expected: Final = _outbound(backend, [_user_turn(marker)], **outbound_fields) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)], **outbound_fields) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{backend}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -1114,7 +1122,7 @@ def _wildcard_prefix(gateway: Gateway, scenario: Scenario, api_base: str) -> str def test_uppercase_claude_id_through_a_mantle_wildcard_is_sent_to_native_messages(gateway: Gateway) -> None: marker: Final = uuid4().hex backend: Final = _HAIKU.upper() - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = f"{_wildcard_prefix(gateway, scenario, wire.url)}/{backend}" response: Final = gateway.request( @@ -1129,7 +1137,7 @@ def test_five_kilobyte_claude_id_through_a_mantle_wildcard_reaches_mantle_and_it ) -> None: marker: Final = uuid4().hex backend: Final = f"{_HAIKU}-{'x' * 5000}" - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) unknown: Final = _error_reply(404, "not_found_error", "model: synthetic unknown model") with wire_server(_bearer_peer(expected, unknown)) as wire, gateway.scenario() as scenario: model: Final = f"{_wildcard_prefix(gateway, scenario, wire.url)}/{backend}" @@ -1203,6 +1211,6 @@ def test_non_claude_mantle_id_with_a_bare_bedrock_twin_keeps_its_route_and_spend cost: Final = _mantle_cost(gateway, _LUNA) assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(cost), response.text _assert_success_row( - _spend_row(_BY_CALL_ID, response.headers["x-litellm-call-id"]), model, _LUNA, "responses", cost + _spend_row(_BY_CALL_ID, response.headers["x-litellm-call-id"]), model, _LUNA, "acompletion", cost ) assert [(request.method, request.target) for request in wire.drain()] == [("POST", _RESPONSES_PATH)] diff --git a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py index 89ad226761e..5aede070896 100644 --- a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py @@ -516,7 +516,7 @@ def test_responses_max_output_tokens_is_forwarded_natively_on_both_legs(gateway: for attempt in attempts: assert attempt.body["max_output_tokens"] == 1000 and _limits(attempt.body) == (None, None), attempt.body _assert_logged(issued, model, "success") - _assert_logged(_managed_response_id(streamed_issued), model, "success") + _assert_logged(streamed_issued, model, "success") def test_region_prefixed_deployment_on_a_v1_base_sends_the_bare_model_with_max_completion_tokens( diff --git a/tests/integration/providers/test_bedrock_model_discovery_wire.py b/tests/integration/providers/test_bedrock_model_discovery_wire.py index 3bf4583e736..2871e51de5e 100644 --- a/tests/integration/providers/test_bedrock_model_discovery_wire.py +++ b/tests/integration/providers/test_bedrock_model_discovery_wire.py @@ -301,11 +301,7 @@ def _spend_row_id(rows: Sequence[Mapping[str, JsonValue]], tag: str) -> str: def _landed_once(served: tuple[tuple[str, str, bool], ...], tags: tuple[str, ...]) -> None: - """One spend row per tagged call, carrying the id the client saw. - - FIXME: a streamed /v1/responses row carries the managed `resp_` id LiteLLM built before the proxy encrypted the - advertised one, so that row is matched by tag alone until the spend log reads the id the client received. - """ + """One spend row per tagged call, carrying the id the client saw.""" rows: Final = eventually( lambda: read_rows( 'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE request_tags ?| %s', @@ -316,9 +312,6 @@ def _landed_once(served: tuple[tuple[str, str, bool], ...], tags: tuple[str, ... ) for (identity, path, stream), tag in zip(served, tags, strict=True): row_id: Final = _spend_row_id(rows, tag) - if path == "/v1/responses" and stream: - assert row_id.startswith("resp_"), (row_id, identity) - continue assert row_id == identity, (path, stream, row_id, identity) diff --git a/tests/integration/providers/test_bedrock_native_structured_output_wire.py b/tests/integration/providers/test_bedrock_native_structured_output_wire.py index 5803d1dee44..67edcf7f3e4 100644 --- a/tests/integration/providers/test_bedrock_native_structured_output_wire.py +++ b/tests/integration/providers/test_bedrock_native_structured_output_wire.py @@ -20,6 +20,7 @@ _FABLE: Final = "us.anthropic.claude-fable-5" _SONNET: Final = "us.anthropic.claude-sonnet-4-6" _OPUS: Final = "us.anthropic.claude-opus-5-5" _OPT_IN: Final = _OPUS +_UNFORCEABLE_TOOL: Final = _OPUS _OPUS_4_8: Final = "us.anthropic.claude-opus-4-8" _DEEPSEEK: Final = "deepseek.v3.2" _AWS: Final[dict[str, JsonValue]] = { @@ -536,6 +537,12 @@ def _resolves(schema: Mapping[str, JsonValue], ref: str) -> bool: return ref.startswith(prefix) and ref.removeprefix(prefix) in object_value(schema.get("$defs") or {}) +def _tool_choice(body: Mapping[str, JsonValue], invoke: bool) -> JsonValue: + if invoke: + return body.get("tool_choice") + return object_value(body.get("toolConfig") or {}).get("toolChoice") + + def _assert_native_without_tools(body: Mapping[str, JsonValue], invoke: bool, schema: Mapping[str, JsonValue]) -> None: assert _native_schema(body, invoke) == schema, body assert _JSON_TOOL not in _tools(body, invoke), body @@ -631,6 +638,40 @@ def test_bedrock_native_model_without_thinking_keeps_the_forced_json_tool(gatewa assert _tools(body, invoke)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], body +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_json_schema_without_thinking_on_a_model_that_cannot_force_the_json_tool_gets_native_output( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _UNFORCEABLE_TOOL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_SCHEMA)) + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + _assert_native_without_tools(body, route == "invoke", _SCHEMA) + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_a_model_that_cannot_force_the_json_tool_retries_a_rejected_native_output_with_an_unforced_json_tool( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker() + invoke: Final = route == "invoke" + with wire_server(_rejecting_native_output(_FABLE_REJECTION)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _UNFORCEABLE_TOOL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_SCHEMA)) + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + rejected, retried = (_JSON.validate_json(request.body) for request in wire.drain()) + assert _native_schema(rejected, invoke) == _SCHEMA, rejected + assert _native_schema(retried, invoke) is None, retried + assert _tools(retried, invoke)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], retried + assert _tool_choice(retried, invoke) is None, retried + + @pytest.mark.parametrize( ("backend", "model_info", "native"), [ @@ -797,6 +838,25 @@ def test_thinking_with_a_caller_tool_on_a_bedrock_native_model_sends_only_the_ca assert _thinking_requested(body, False), body +def test_a_caller_tool_next_to_json_schema_without_thinking_on_a_model_that_cannot_force_the_json_tool( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _UNFORCEABLE_TOOL) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body(model, marker, tools=[_USER_TOOL_SPEC], response_format=_response_format(_SCHEMA)), + ) + (call,) = list_value(_chat_message(response)["tool_calls"]) + function: Final = object_value(object_value(call)["function"]) + assert function["name"] == _USER_TOOL, response.text + _, body = _one_call(wire, marker_hex) + assert _tools(body, False).keys() == {_USER_TOOL}, body + assert _native_schema(body, False) == _SCHEMA, body + + @pytest.mark.parametrize( ("schema", "expected"), [(_UNCLEANED_SCHEMA, _UNCLEANED_SCHEMA), (_NULLABLE_SCHEMA, _CLOSED_NULLABLE_SCHEMA)], diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py index 3cadc1b5ccf..42950fb18c3 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py @@ -172,7 +172,11 @@ def _upstream_id_inside(row_id: str) -> str | None: def _row_belongs_to(row_id: str, served: _Served) -> bool: if row_id == _response_id(served): return True - return served.call.endpoint == "responses" and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}" + return ( + served.call.endpoint == "responses" + and not served.call.stream + and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}" + ) def _assert_each_success_landed_once(rows: list[dict[str, JsonValue]], served: tuple[_Served, ...]) -> None: diff --git a/tests/integration/providers/test_chat_template_tool_count_wire.py b/tests/integration/providers/test_chat_template_tool_count_wire.py index a8b97fa732e..552df3a87bc 100644 --- a/tests/integration/providers/test_chat_template_tool_count_wire.py +++ b/tests/integration/providers/test_chat_template_tool_count_wire.py @@ -39,7 +39,7 @@ _API_KEY: Final = "scripted-provider-key" _DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731" _DEEPSEEK_R1: Final = "azure_ai/DeepSeek-R1-0528" _QWEN3_5: Final = "databricks/databricks-qwen35-122b-a10b" -_GROK: Final = "azure_ai/grok-4-1" +_GROK_3: Final = "azure_ai/grok-3" _DEEPSEEK_V4_TOOL_PREAMBLE: Final = 222 _QWEN3_5_TOOL_PREAMBLE: Final = 196 _QWEN3_5_OPENED_SYSTEM_TURN: Final = 5 @@ -250,7 +250,7 @@ def test_messages_count_tokens_skips_web_search_and_counts_unmapped_tool_keys_fo def test_messages_count_tokens_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: with wire_server(_completion) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, "count", _GROK, wire.url) + model: Final = _deployment(scenario, "count", _GROK_3, wire.url) without_tools: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ())) plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) deferred: Final = _count( @@ -428,7 +428,7 @@ def test_project_itpm_reserves_deepseek_v4_tools_the_way_its_chat_template_rende def test_project_itpm_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: chat: Final = _ENDPOINTS["chat-completions"] with wire_server(_completion) as wire, gateway.scenario() as scenario: - group: Final = _deployment(scenario, "grok-4-1", _GROK, wire.url) + group: Final = _deployment(scenario, "grok-3", _GROK_3, wire.url) response: Final = gateway.request( "POST", chat.path, @@ -504,8 +504,8 @@ def _agent_body(group: str, marker: str, file_body: str, trigger: int) -> dict[s @pytest.mark.parametrize( ("model", "files_kept"), - ((_DEEPSEEK_V4, (False, False, True)), (_GROK, (True, True, True))), - ids=("deepseek-v4", "grok-4-1"), + ((_DEEPSEEK_V4, (False, False, True)), (_GROK_3, (True, True, True))), + ids=("deepseek-v4", "grok-3"), ) def test_messages_context_management_trigger_counts_tools_the_way_the_models_chat_template_renders_them( gateway: Gateway, model: str, files_kept: tuple[bool, ...] diff --git a/tests/integration/providers/test_custom_video_provider_wire.py b/tests/integration/providers/test_custom_video_provider_wire.py new file mode 100644 index 00000000000..a40ecae3ed4 --- /dev/null +++ b/tests/integration/providers/test_custom_video_provider_wire.py @@ -0,0 +1,113 @@ +"""A provider registered in `custom_provider_map` with a `BaseVideoConfig` serves /v1/videos create, status and content.""" + +import json +import shutil +from contextlib import ExitStack +from pathlib import Path +from typing import Final + +from integration._support.client import Gateway +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +_SUPPORT_MODULE: Final = Path(__file__).resolve().parents[1] / "_support" / "custom_video_provider.py" +_PROMPT: Final = "a cat walking across a bridge" +_MP4: Final = b"\x00\x00\x00\x18ftypmp42" + b"queue-video-bytes" + + +def _queue_backend(request: Request) -> Reply: + if (request.method, request.target) == ("POST", "/jobs"): + return Reply(body=json.dumps({"job_id": "job-7", "state": "queued"}).encode()) + if (request.method, request.target) == ("GET", "/jobs/job-7"): + return Reply(body=json.dumps({"job_id": "job-7", "state": "done", "seconds": 8}).encode()) + if (request.method, request.target) == ("GET", "/jobs/job-7/file"): + return Reply(body=_MP4, content_type="video/mp4") + return Reply(status=404, body=json.dumps({"error": f"no route {request.method} {request.target}"}).encode()) + + +def _write_config(path: Path, wire: Wire) -> Path: + shutil.copy(_SUPPORT_MODULE, path.parent / "custom_video_provider.py") + path.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": "queue-video", + "litellm_params": { + "model": "queue-video/seedance-1", + "api_base": wire.url, + "api_key": "queue-secret", + }, + }, + { + "model_name": "image-only-video", + "litellm_params": { + "model": "image-only/seedance-1", + "api_base": wire.url, + "api_key": "unused", + }, + }, + ], + "litellm_settings": { + "custom_provider_map": [ + {"provider": "queue-video", "custom_handler": "custom_video_provider.VIDEO_CONFIG"}, + {"provider": "image-only", "custom_handler": "custom_video_provider.IMAGE_ONLY_HANDLER"}, + ] + }, + } + ), + encoding="utf-8", + ) + return path + + +def _lower_headers(request: Request) -> dict[str, str]: + return {name.lower(): value for name, value in request.headers.items()} + + +def test_custom_provider_map_base_video_config_serves_create_status_and_content( + gateway: Gateway, tmp_path: Path +) -> None: + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_queue_backend)) + config: Final = _write_config(tmp_path / "custom-video.yaml", wire) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + + created: Final = owned.gateway.request( + "POST", "/v1/videos", {"model": "queue-video", "prompt": _PROMPT, "seconds": "8"} + ) + assert created.status_code == 200, created.text + created_body: Final = created.json() + assert created_body["status"] == "queued", created_body + video_id: Final = created_body["id"] + assert isinstance(video_id, str) and video_id.startswith("video_"), created_body + + submits: Final = wire.drain() + assert [(r.method, r.target) for r in submits] == [("POST", "/jobs")], submits + assert json.loads(submits[0].body) == {"model": "seedance-1", "prompt": _PROMPT, "seconds": "8"} + assert _lower_headers(submits[0])["x-api-key"] == "queue-secret" + + status: Final = owned.gateway.request("GET", f"/v1/videos/{video_id}") + assert status.status_code == 200, status.text + assert status.json()["status"] == "completed", status.text + assert status.json()["id"] == video_id, status.text + + content: Final = owned.gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 200, content.text + assert content.content == _MP4 + + polls: Final = wire.drain() + assert [(r.method, r.target) for r in polls] == [("GET", "/jobs/job-7"), ("GET", "/jobs/job-7/file")], polls + assert all(_lower_headers(r)["x-api-key"] == "queue-secret" for r in polls) + + +def test_custom_provider_without_a_video_config_is_refused_with_400(gateway: Gateway, tmp_path: Path) -> None: + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_queue_backend)) + config: Final = _write_config(tmp_path / "custom-video.yaml", wire) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + + refused: Final = owned.gateway.request("POST", "/v1/videos", {"model": "image-only-video", "prompt": _PROMPT}) + assert refused.status_code == 400, refused.text + assert "video generation is not supported for image-only" in refused.text + assert wire.drain() == () diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index b93cebdbc4c..3d9d749eb04 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -353,9 +353,172 @@ def _shape_reply(identity: str, stream: bool) -> Reply: "/chat/completions", id="A5-responses-list-input", ), + pytest.param( + "json-object-two-blocks", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A6-chat-json-object-two-blocks-joined", + ), + pytest.param( + "json-object-two-blocks-stream", + "/v1/chat/completions", + "databricks-qwen35-122b-a10b", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + "stream": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A7-chat-stream-json-object-two-blocks-joined", + ), + pytest.param( + "json-object-history", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ], + "response_format": {"type": "json_object"}, + }, + [ + {"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": "Reply in JSON"}, + ], + False, + "/chat/completions", + id="A8-chat-json-object-earlier-turns-kept", + ), + pytest.param( + "json-object-trailing-system", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Say hi"}, {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}], + False, + "/chat/completions", + id="A8b-chat-json-object-trailing-system-keeps-break", + ), + pytest.param( + "json-object-drop-params-cache-control", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Reply"}, + {"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}}, + ], + } + ], + "response_format": {"type": "json_object"}, + "drop_params": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8c-chat-json-object-drop-params-drops-cache-control", + ), + pytest.param( + "json-object-trailing-thinking-turn", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + { + "role": "assistant", + "content": "", + "thinking_blocks": [{"type": "thinking", "thinking": "Plan.", "signature": "sig"}], + }, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8d-chat-json-object-trailing-thinking-turn-dropped", + ), + pytest.param( + "json-object-external-endpoint", + "/v1/chat/completions", + "openai-gpt-external", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}], + False, + "/chat/completions", + id="A8e-chat-json-object-external-endpoint-kept-as-sent", + ), + pytest.param( + "json-object-responses-two-parts", + "/v1/responses", + "databricks-qwen35-122b-a10b", + { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}], + } + ], + "text": {"format": {"type": "json_object"}}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A9-responses-json-object-two-parts-joined", + ), + pytest.param( + "json-object-responses-trailing-reasoning", + "/v1/responses", + "databricks-qwen35-122b-a10b", + { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}], + }, + {"type": "reasoning", "id": "rs_1", "summary": [{"type": "summary_text", "text": "Plan."}]}, + ], + "text": {"format": {"type": "json_object"}}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A9b-responses-json-object-trailing-reasoning-dropped", + ), ], ) -def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( +def test_databricks_non_claude_plain_text_content_reaches_upstream_as_string( gateway: Gateway, case: str, endpoint: str, @@ -365,7 +528,7 @@ def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( responses_mode: bool, upstream_path: str, ) -> None: - identity: Final = f"databricks-single-text-{case}-{uuid.uuid4().hex}" + identity: Final = f"databricks-plain-text-{case}-{uuid.uuid4().hex}" streaming: Final = request_body.get("stream") is True def respond(request: Request) -> Reply: @@ -1453,7 +1616,7 @@ def test_databricks_non_string_text_block_is_rejected_before_reaching_upstream(g ) received: Final = wire.drain() assert response.status_code == 400, response.text - assert "'int' object has no attribute 'strip'" in response.text + assert "A text content block needs a string 'text', got int" in response.text assert received == () diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 5af1532e9ed..4e091460667 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -262,11 +262,28 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, return _calls_to(_observed_requests(gateway), handle) -def _sdk_inference_call(gateway: Gateway, handle: ScenarioHandle, provider: _Provider) -> dict[str, JsonValue]: +def _sdk_inference_call( + gateway: Gateway, + handle: ScenarioHandle, + provider: _Provider, + *, + metadata_authorization: str | None = None, +) -> dict[str, JsonValue]: observed: Final = _upstream_calls(gateway, handle) inference_path: Final = f"/{handle.scenario_id}/v1/systemone" (call,) = tuple(request for request in observed if request["path"] == inference_path) - metadata_calls: Final = tuple(request for request in observed if request["path"] != inference_path) + environment_metadata: Final = { + "path": f"/{handle.scenario_id}/api/show", + "authorization": metadata_authorization, + "body": {"name": provider.body_model}, + "method": "POST", + "api_key": "", + } + metadata_calls: Final = tuple( + request + for request in observed + if request["path"] != inference_path and (metadata_authorization is None or request != environment_metadata) + ) assert metadata_calls in ( (), ( @@ -556,7 +573,9 @@ def test_ollama_decisions_sdk_respects_environment_and_explicit_auth( api_key=api_key, ) assert response.model_dump(mode="json") == _answer_body(provider) - call: Final = _sdk_inference_call(gateway, handle, provider) + call: Final = _sdk_inference_call( + gateway, handle, provider, metadata_authorization=authorization if environment_base else None + ) assert (call["path"], call["authorization"], call["body"]) == ( f"/{handle.scenario_id}/v1/systemone", authorization, diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 8aba6841cf9..91f502179f3 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -28,7 +28,6 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = [pytest.mark.timeout(900), pytest.mark.shares_worker_state] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_DROPPED: Final = re.compile(r"original model: (\S+), ignoring and continuing") _FIXED: Final[Mapping[str, str]] = MappingProxyType( {"chatgpt-fixed": "chatgpt/gpt-5.6-terra", "copilot-fixed": "github_copilot/gpt-5.4"} ) @@ -41,8 +40,6 @@ _CONTROL_CALLS_BEFORE_THE_RELOAD: Final = 1 _SHAPES: Final[Mapping[str, bool]] = MappingProxyType( {"docs-shapes": False, "docs-shapes-beside-a-universal-wildcard": True} ) -_NOT_FOUND: Final = "Invalid model name passed in model=" -_NO_HEALTHY: Final = "There are no healthy deployments for this model" _CONTROL: Final = "device-login-boot-control" _EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0, "cache": {"no-cache": True}}) _CALL_ID: Final = "x-litellm-call-id" @@ -174,15 +171,11 @@ def _lists_login_models(listed: tuple[str, ...]) -> bool: ) -def _every_worker_lists( - owned: OwnedProxy, workers: tuple[int, ...], *, login_models: bool -) -> Mapping[int, tuple[str, ...]]: +def _every_worker_lists(owned: OwnedProxy, workers: tuple[int, ...]) -> Mapping[int, tuple[str, ...]]: def settled(by_worker: Mapping[int, tuple[str, ...]]) -> bool: if set(by_worker) != set(workers): return False - if login_models: - return all(_lists_login_models(listed) for listed in by_worker.values()) - return all(_login_names_listed(listed) == frozenset() for listed in by_worker.values()) + return all(_lists_login_models(listed) for listed in by_worker.values()) return eventually(lambda: _listing_by_worker(owned, workers), settled, seconds=90) @@ -212,16 +205,8 @@ def _reload(owned: OwnedProxy) -> None: assert response.status_code == 200, response.text -def _refusal_before_the_reload(name: str, universal: bool) -> str: - if name in _FIXED: - return _NO_HEALTHY - if not universal: - return _NOT_FOUND - return dl.CHATGPT_REFUSAL if name.startswith("chatgpt/") else dl.COPILOT_REFUSAL - - @pytest.mark.parametrize("universal", _SHAPES.values(), ids=_SHAPES.keys()) -def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_mounting_brings_them_back( +def test_booting_without_tokens_lists_login_deployments_and_serves_only_after_tokens_are_mounted( gateway: Gateway, tmp_path: Path, universal: bool ) -> None: dirs: Final = _token_dirs(tmp_path) @@ -240,20 +225,18 @@ def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_m workers=2, ) as owned: workers: Final = _workers(owned) - _every_worker_lists(owned, workers, login_models=False) + _every_worker_lists(owned, workers) boot_log: Final = owned.log.read_text() assert peers.auth_connections() == () assert dl.CHATGPT_USER_CODE not in boot_log and dl.COPILOT_USER_CODE not in boot_log, boot_log[-3000:] - dropped_models: Final = _DROPPED.findall(boot_log) - expected_drops: Final = frozenset((*_FIXED.values(), *_WILDCARDS)) - assert frozenset(dropped_models) == expected_drops, boot_log[-3000:] - assert all(dropped_models.count(model) >= len(workers) for model in expected_drops), boot_log[-3000:] with owned.gateway.scenario() as scenario: key: Final = scenario.key() for name in _LOGIN_NAMES: refused: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert refused.status == 400, refused.text - assert _refusal_before_the_reload(name, universal) in refused.message(), refused.text + refusal: Final = dl.CHATGPT_REFUSAL if name.startswith("chatgpt") else dl.COPILOT_REFUSAL + assert refusal in refused.message(), refused.text + assert peers.api.drain() == () control: Final = _chat(_base_url(owned), key, _CONTROL, uuid.uuid4().hex) assert control.status == 200, control.text assert peers.auth_connections() == () @@ -261,7 +244,7 @@ def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_m dirs.store_valid_tokens(peers.api.url) _reload(owned) - _every_worker_lists(owned, workers, login_models=True) + _every_worker_lists(owned, workers) for served_logins, name in enumerate(_LOGIN_NAMES, start=1): served: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert served.status == 200, served.text diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index b4044674ed5..39e36fd95fa 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -746,10 +746,7 @@ def test_a_recent_device_code_request_does_not_hold_the_worker_for_the_rest_of_t assert scene.rig.peers.auth_connections() == () -_NOT_LIVE: Final = re.compile(r"saved to the database, but the model id\(s\) \['([0-9a-f-]+)'\] are not live") - - -def test_model_created_without_a_token_is_saved_not_live_and_serves_once_the_token_is_mounted(scene: _Scene) -> None: +def test_model_created_without_a_token_refuses_requests_and_serves_once_the_token_is_mounted(scene: _Scene) -> None: gateway: Final = scene.rig.owned.gateway name: Final = f"device-login-db-{uuid.uuid4().hex}" identity: Final = str(uuid.uuid4()) @@ -763,14 +760,18 @@ def test_model_created_without_a_token_is_saved_not_live_and_serves_once_the_tok }, ) try: - assert created.status_code == 500, created.text - not_live: Final = _NOT_LIVE.search(created.text) - assert not_live is not None and not_live.group(1) == identity, created.text + assert created.status_code == 200, created.text + assert created.json()["model_info"]["id"] == identity, created.text assert scene.rig.peers.auth_connections() == () - scene.store_valid_token("chatgpt") call: Final = _Call( scene.rig.base_url(), scene.key, name, "responses", False, "httpx", uuid.uuid4().hex, SimpleQueue() ) + refused: Final = _serve(call) + assert refused.status == 400, refused.text + assert dl.CHATGPT_REFUSAL in _error_message(refused.text), refused.text + assert scene.rig.peers.auth_connections() == () + assert scene.rig.peers.api.drain() == () + scene.store_valid_token("chatgpt") served: Final = eventually(lambda: _serve(call), lambda answer: answer.status == 200, seconds=90) assert set(rv.MARKER.findall(served.text)) == {call.marker}, served.text assert scene.rig.peers.auth_connections() == () diff --git a/tests/integration/providers/test_gemini_long_decode_wire.py b/tests/integration/providers/test_gemini_long_decode_wire.py index 5ecb5e2a942..90e7616e26c 100644 --- a/tests/integration/providers/test_gemini_long_decode_wire.py +++ b/tests/integration/providers/test_gemini_long_decode_wire.py @@ -380,7 +380,9 @@ def test_cut_chat_stream_hands_its_token_on_a_chunk_with_empty_vertex_ai_lists_a chunks: Final = payloads(post(gateway, "chat", request_body("chat", model, True)).text, True) carriers: Final = [chunk for chunk in chunks if "continuation_token" in json.dumps(chunk)] assert [_vertex_lists(chunk) for chunk in carriers] == [dict(_EMPTY_VERTEX_LISTS)], chunks - assert [_vertex_lists(chunk) for chunk in chunks if chunk not in carriers] == [{}] * (len(chunks) - 1), chunks + assert [_vertex_lists(chunk) for chunk in chunks] == [ + dict(_EMPTY_VERTEX_LISTS) if at(chunk, "choices", 0, "delta", "content") else {} for chunk in chunks + ], chunks metered: Final = [chunk for chunk in chunks if isinstance(chunk.get("usage"), dict)] assert [at(chunk, "usage", "completion_tokens_details") for chunk in metered] == [ {"text_tokens": CUT.output_tokens} diff --git a/tests/integration/providers/test_gemini_thinking_replay_chaos.py b/tests/integration/providers/test_gemini_thinking_replay_chaos.py index f8bb52ef60a..fb3092adfc3 100644 --- a/tests/integration/providers/test_gemini_thinking_replay_chaos.py +++ b/tests/integration/providers/test_gemini_thinking_replay_chaos.py @@ -233,6 +233,18 @@ def _assert_answered_with_its_own_marker(served: _Served) -> None: assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text +def _spend_id(served: _Served) -> str: + if served.call.endpoint == "responses" and served.call.stream: + payloads: Final = ( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in served.text.splitlines() + if line.startswith("data: {") + ) + (completed,) = tuple(event for event in payloads if event.get("type") == "response.completed") + return str(_JSON_OBJECT.validate_python(completed["response"])["id"]) + return _response_id(served.call.marker) + + async def test_concurrent_thinking_replays_across_endpoints_reach_gemini_signature_free(gateway: Gateway) -> None: calls: Final = _calls(30, ("chat", "messages", "responses"), lambda index: index % 2 == 0) with wire_server(_echo) as wire, gateway.scenario() as scenario: @@ -244,7 +256,7 @@ async def test_concurrent_thinking_replays_across_endpoints_reach_gemini_signatu _assert_no_bleed(wire.drain(), frozenset(call.marker for call in calls)) landed: Final = _spend_request_ids(model, 30) assert landed >= { - _response_id(call.marker) for call in calls if not (call.endpoint == "messages" and call.stream) + _spend_id(item) for item in served if not (item.call.endpoint == "messages" and item.call.stream) } diff --git a/tests/integration/providers/test_gemini_thinking_replay_wire.py b/tests/integration/providers/test_gemini_thinking_replay_wire.py index 4f965f079ec..70e339ed8b7 100644 --- a/tests/integration/providers/test_gemini_thinking_replay_wire.py +++ b/tests/integration/providers/test_gemini_thinking_replay_wire.py @@ -2,6 +2,7 @@ import base64 import itertools import json import os +import re import uuid from collections.abc import Iterator, Mapping, Sequence from collections.abc import Set as AbstractSet @@ -679,9 +680,70 @@ def test_gemini_own_tool_call_signature_is_still_replayed_on_the_function_call_p assert _spend_row(rounds[1])["status"] == "success" -def test_responses_previous_response_id_replays_gemini_session_history( - gateway: Gateway, tmp_path: Path +_PLACEHOLDER_SIGNATURE: Final = base64.b64encode(b"skip_thought_signature_validator").decode() +_ISSUED_ID: Final = re.compile(r'"id":\s*"([^"]+)"') +_PLACEHOLDER_TOOL_TURN: Final = _model_turn( + _thought("I should count."), + {"function_call": {"name": "count_fruit", "args": {"kind": "apple"}}, "thoughtSignature": _PLACEHOLDER_SIGNATURE}, +) +# Wording AI Studio answers with when a Vertex-minted signature is replayed to it, captured live on 2026-10-27 +_FOREIGN_SIGNATURE_REJECTED: Final = Reply( + status=400, + body=json.dumps( + {"error": {"code": 400, "message": "Corrupted thought signature.", "status": "INVALID_ARGUMENT"}} + ).encode(), +) + + +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +def test_gemini_tool_call_signature_the_backend_rejects_is_resent_once_with_the_placeholder( + gateway: Gateway, endpoint: Endpoint, stream: bool ) -> None: + rounds: Final = (f"gemini-reply-{uuid.uuid4().hex}", f"gemini-reply-{uuid.uuid4().hex}") + calls: Final = itertools.count() + + def respond(request: Request) -> Reply: + match next(calls): + case 0: + return _reply(rounds[0], stream=False, parts=_GEMINI_TOOL_REPLY_PARTS) + case 1: + return _FOREIGN_SIGNATURE_REJECTED + case _: + return _reply(rounds[1], stream=stream, parts=({"text": "Twelve apples."},)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + first: Final = _tool_round_one(endpoint, model) + headers: Final = {"anthropic-version": "2023-06-01"} + opened: Final = gateway.request("POST", _path(endpoint), first, headers=headers) + assert opened.status_code == 200, opened.text + second: Final = { + **_tool_round_two(endpoint, first, _JSON_OBJECT.validate_json(opened.content)), + "stream": stream, + } + closed: Final = gateway.request("POST", _path(endpoint), second, headers=headers) + assert closed.status_code == 200, closed.text + assert "Twelve apples." in closed.text, closed.text + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [ + ("POST", _target("gemini", False)), + ("POST", _target("gemini", stream)), + ("POST", _target("gemini", stream)), + ], received + rejected, resent = (_JSON_OBJECT.validate_json(request.body) for request in received[1:]) + assert rejected == { + **_expected_body(endpoint, _user("Count the apples."), _GEMINI_TOOL_REPLAYED_TURN, _FUNCTION_RESPONSE_TURN), + **_DECLARATIONS, + }, rejected + assert resent == { + **_expected_body(endpoint, _user("Count the apples."), _PLACEHOLDER_TOOL_TURN, _FUNCTION_RESPONSE_TURN), + **_DECLARATIONS, + }, resent + assert _spend_row(rounds[1], *_ISSUED_ID.findall(closed.text))["status"] == "success" + + +def test_responses_previous_response_id_replays_gemini_session_history(gateway: Gateway, tmp_path: Path) -> None: first_prompt: Final = f"gemini-session-first-{uuid.uuid4().hex}" second_prompt: Final = f"gemini-session-second-{uuid.uuid4().hex}" first_answer: Final = f"gemini-answer-first-{uuid.uuid4().hex}" diff --git a/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py b/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py index ead83ed66aa..c185c968be9 100644 --- a/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py +++ b/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py @@ -301,6 +301,18 @@ def _assert_answered_with_its_own_marker(served: _Served) -> None: assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text +def _spend_id(served: _Served) -> str: + if served.call.endpoint == "responses" and served.call.stream: + payloads: Final = ( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in served.text.splitlines() + if line.startswith("data: {") + ) + (completed,) = tuple(event for event in payloads if event.get("type") == "response.completed") + return str(_JSON_OBJECT.validate_python(completed["response"])["id"]) + return _response_id(served.call.marker) + + async def test_a_concurrent_mixed_burst_reaches_gemini_with_each_tool_result_in_its_own_form(gateway: Gateway) -> None: calls: Final = _calls( 30, ("chat", "messages", "responses"), lambda index: index % 2 == 0, lambda index: (index // 2) % 2 == 0 @@ -314,7 +326,7 @@ async def test_a_concurrent_mixed_burst_reaches_gemini_with_each_tool_result_in_ _assert_each_in_its_own_form(wire.drain(), calls) landed: Final = _spend_rows_by_model(model, 30) assert landed >= { - _response_id(call.marker) for call in calls if not (call.endpoint == "messages" and call.stream) + _spend_id(item) for item in served if not (item.call.endpoint == "messages" and item.call.stream) } diff --git a/tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py index 807a7461675..858599829ad 100644 --- a/tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py @@ -41,6 +41,7 @@ class _Proxy: openai: Wire gpt_5_6_with_prompt_cache_options: str gpt_5_6_without_prompt_cache_options: str + gpt_5_6_dropping_params: str def openai_receives( self, model: str, messages: Sequence[JsonValue], *, stream: bool = False, tools: Sequence[JsonValue] = () @@ -79,6 +80,9 @@ def proxy() -> Iterator[_Proxy]: model="openai/gpt-5.6", api_base=api_base, prompt_cache_options={"mode": "explicit"} ), gpt_5_6_without_prompt_cache_options=scenario.model(model="openai/gpt-5.6", api_base=api_base), + gpt_5_6_dropping_params=scenario.model( + model="openai/gpt-5.6", api_base=api_base, prompt_cache_options={"mode": "explicit"}, drop_params=True + ), ) @@ -202,3 +206,56 @@ def test_custom_api_base_without_prompt_cache_options_forwards_cache_control_unc assert received["messages"] == _SYSTEM_AND_USER_WITH_CACHE_CONTROL assert "prompt_cache_options" not in received assert "prompt_cache_breakpoint" not in str(received) + + +_CLIENT_BREAKPOINTS_ON_PARTS: Final[list[JsonValue]] = [ + {"type": "text", "text": "Stable rules", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}}, + {"type": "text", "text": "More rules", "prompt_cache_breakpoint": {"mode": "implicit"}}, + {"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}}, +] +_CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM: Final[list[JsonValue]] = [ + {"type": "text", "text": "Stable rules", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "text", "text": "More rules"}, + {"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}}, +] +_CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT: Final = ( + "live gpt-6.1-sol Chat Completions answers 400 on a part breakpoint with any ttl or mode implicit (2026-10-28)" +) + + +@pytest.mark.parametrize("stream", [False, True], ids=["unary", "stream"]) +def test_drop_params_cuts_client_breakpoints_down_to_what_chat_completions_takes(proxy: _Proxy, stream: bool) -> None: + received: Final = proxy.openai_receives( + proxy.gpt_5_6_dropping_params, [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}], stream=stream + ) + + assert received["messages"] == [ + {"role": "user", "content": _CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM} + ], _CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT + + +def test_without_drop_params_client_breakpoints_reach_openai_unchanged(proxy: _Proxy) -> None: + received: Final = proxy.openai_receives( + proxy.gpt_5_6_with_prompt_cache_options, [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}] + ) + + assert received["messages"] == [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}] + + +def test_drop_params_cuts_client_breakpoints_on_v1_messages_routed_to_chat_completions(proxy: _Proxy) -> None: + _ = proxy.openai.drain() + response: Final = proxy.gateway.request( + "POST", + "/v1/messages", + { + "model": proxy.gpt_5_6_dropping_params, + "max_tokens": 16, + "stop_sequences": ["END"], + "messages": [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}], + }, + ) + assert response.status_code == 200, response.text + + assert proxy.received_by_openai()["messages"] == [ + {"role": "user", "content": _CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM} + ], _CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT diff --git a/tests/integration/providers/test_openai_chat_wire.py b/tests/integration/providers/test_openai_chat_wire.py index 69295485380..5ba631ac776 100644 --- a/tests/integration/providers/test_openai_chat_wire.py +++ b/tests/integration/providers/test_openai_chat_wire.py @@ -1,4 +1,3 @@ -import base64 import json import uuid from collections.abc import Callable @@ -547,14 +546,8 @@ def test_azure_gpt_6_bridged_no_cache_function_requests_each_reach_provider_and_ lambda found: len(found) == 2, seconds=70, ) - by_response_id: Final = { - (decoded := base64.b64decode(string_value(row["request_id"]).removeprefix("resp_")).decode()) - .rsplit("response_id:", 1)[1]: (decoded, row) - for row in rows - } - for response_id in response_ids: - decoded, row = by_response_id[response_id] - assert decoded.startswith("litellm:custom_llm_provider:azure;model_id:"), rows + assert sorted(string_value(row["request_id"]) for row in rows) == sorted(response_ids), rows + for row in rows: assert ( string_value(row["status"]), string_value(row["cache_hit"]), @@ -601,11 +594,7 @@ def test_azure_gpt_6_bridged_function_requests_without_cache_field_still_hit_cac seconds=70, ) priced, cached = rows - priced_request: Final = base64.b64decode( - string_value(priced["request_id"]).removeprefix("resp_") - ).decode() - assert priced_request.startswith("litellm:custom_llm_provider:azure;model_id:"), rows - assert priced_request.endswith(";response_id:resp_cached"), rows + assert string_value(priced["request_id"]) == "resp_cached", rows assert ( priced["status"], priced["cache_hit"], diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py index 33338cae95d..7ea02435369 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py @@ -27,8 +27,9 @@ _MALFORMED: Final[tuple[tuple[str, JsonValue], ...]] = ( ("5kb-string", "x" * 5000), ("empty-object", {}), ("bogus-mode", {"mode": "bogus"}), - ("bad-ttl", {"mode": "explicit", "ttl": "1h"}), + ("bogus-mode-with-ttl", {"mode": "bogus", "ttl": "30m"}), ) +_UNSUPPORTED_TTL: Final[Mapping[str, JsonValue]] = {"mode": "explicit", "ttl": "1h"} _MALFORMED_IDS: Final = tuple(name for name, _ in _MALFORMED) _MALFORMED_VALUES: Final = tuple(value for _, value in _MALFORMED) _CASES: Final[tuple[tuple[str, Mode, JsonValue, JsonValue], ...]] = ( @@ -192,6 +193,27 @@ def test_marker_with_an_unknown_key(bridge: _Bridge, mode: Mode, expected: JsonV bridge.spend.landed(bridge.model(mode), call_id, marker) +@pytest.mark.parametrize( + ("mode", "expected"), (("on", None), ("off", _UNSUPPORTED_TTL)), ids=("dropped-on", "verbatim-off") +) +def test_marker_with_an_unsupported_ttl(bridge: _Bridge, mode: Mode, expected: JsonValue) -> None: + marker: Final = uuid.uuid4().hex + call_id: Final = _completion(_chat(bridge, bridge.model(mode), [_user(marker, _UNSUPPORTED_TTL)]), marker) + pcb.assert_marker(_user_block_on_wire(bridge, marker), expected) + bridge.spend.landed(bridge.model(mode), call_id, marker) + + +def test_marker_with_an_unsupported_ttl_keeps_its_mode_on_a_host_opted_in_with_prompt_cache_options( + bridge: _Bridge, +) -> None: + marker: Final = uuid.uuid4().hex + with bridge.gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=bridge.api_base, drop_params=True, **_OPTED_IN) + call_id: Final = _completion(_chat(bridge, model, [_user(marker, _UNSUPPORTED_TTL)]), marker) + pcb.assert_marker(_user_block_on_wire(bridge, marker), pcb.EXPLICIT) + bridge.spend.landed(model, call_id, marker) + + def _second_block_on_wire(bridge: _Bridge, marker: str) -> dict[str, JsonValue]: request: Final = pcb.posted(bridge.wire, marker) _wire_body(request) diff --git a/tests/integration/providers/test_responses_core_wire.py b/tests/integration/providers/test_responses_core_wire.py index 9a24903a20d..f9ff26045bc 100644 --- a/tests/integration/providers/test_responses_core_wire.py +++ b/tests/integration/providers/test_responses_core_wire.py @@ -1,10 +1,12 @@ from __future__ import annotations +import base64 import json import os import uuid -from collections.abc import Callable -from typing import Final +from collections.abc import AsyncIterator, Callable, Iterator +from copy import deepcopy +from typing import Final, Protocol, runtime_checkable import openai import pytest @@ -14,6 +16,10 @@ from integration._support.worker_readiness import wait_for_worker_readiness from openai.types.responses import ResponseCompletedEvent from pydantic import JsonValue, TypeAdapter +import litellm +from litellm.types.llms.openai import ResponseCompletedEvent as LiteLLMResponseCompletedEvent +from litellm.types.llms.openai import ResponsesAPIResponse + _MODEL: Final = "gpt-5" _API_KEY: Final = "synthetic-responses-key" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -21,6 +27,16 @@ _PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", " _ALIASES: Final = ("/v1", "", "/openai/v1") +@runtime_checkable +class _AsyncResponseEvents(Protocol): + def __aiter__(self) -> AsyncIterator[object]: ... + + +@runtime_checkable +class _ResponseEvents(Protocol): + def __iter__(self) -> Iterator[object]: ... + + def _deployments_in_group(gateway: Gateway, group: str) -> int: entries: Final = gateway.get("/model/info")["data"] assert isinstance(entries, list), entries @@ -69,6 +85,175 @@ def _stream(identity: str) -> tuple[bytes, ...]: ) + (b"data: [DONE]\n\n",) +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("stream", (False, True), ids=("nonstream", "stream")) +@pytest.mark.parametrize("extra", (False, True), ids=("top-level", "extra-body-precedence")) +async def test_regression_sdk_comparison_response_id_is_raw_on_the_wire( + asynchronous: bool, stream: bool, extra: bool +) -> None: + raw_id: Final = f"resp_sdk_comparison_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/v1/responses") + body: Final = _JSON_OBJECT.validate_json(request.body) + if body["input"] == "seed comparison": + assert body == {"model": _MODEL, "input": "seed comparison"}, request.body + return Reply(body=_response(raw_id)) + assert body == { + "model": _MODEL, + "input": "compare response", + "stream": stream, + "previous_response_id": raw_id, + "prompt_cache_options": {"comparison_response_id": raw_id}, + }, request.body + return ( + Reply(content_type="text/event-stream", chunks=_stream(raw_id)) if stream else Reply(body=_response(raw_id)) + ) + + with wire_server(respond) as wire: + seeded: Final = litellm.responses( + model=f"openai/{_MODEL}", input="seed comparison", api_base=f"{wire.url}/v1", api_key=_API_KEY + ) + assert isinstance(seeded, ResponsesAPIResponse), seeded + assert seeded.id != raw_id, seeded + params: Final = { + "model": f"openai/{_MODEL}", + "input": "compare response", + "stream": stream, + "previous_response_id": seeded.id, + "api_base": f"{wire.url}/v1", + "api_key": _API_KEY, + "num_retries": 0, + **( + { + "prompt_cache_options": {"mode": "explicit", "ttl": "12h"}, + "extra_body": {"prompt_cache_options": {"comparison_response_id": seeded.id}}, + } + if extra + else {"prompt_cache_options": {"comparison_response_id": seeded.id}} + ), + } + original: Final = deepcopy(params) + result: Final = await litellm.aresponses(**params) if asynchronous else litellm.responses(**params) + if stream: + if asynchronous: + assert isinstance(result, _AsyncResponseEvents), result + async_events: Final = tuple([event async for event in result]) + async_completed: Final = async_events[-1] + assert isinstance(async_completed, LiteLLMResponseCompletedEvent), async_events + assert async_completed.response.status == "completed", async_events + else: + assert isinstance(result, _ResponseEvents), result + sync_events: Final = tuple(result) + sync_completed: Final = sync_events[-1] + assert isinstance(sync_completed, LiteLLMResponseCompletedEvent), sync_events + assert sync_completed.response.status == "completed", sync_events + else: + assert isinstance(result, ResponsesAPIResponse), result + assert result.status == "completed", result + assert params == original, "The caller's comparison options were mutated" + assert len(wire.drain()) == 2 + + +@pytest.mark.parametrize("extra", (False, True)) +@pytest.mark.parametrize( + "comparison_id", + ( + "resp_custom_vendor_id", + "resp_!!!not-base64", + "cmF3X3ZlbmRvcl9pZA==", + pytest.param( + "resp_" + + base64.b64encode(b"litellm:custom_llm_provider:xai;model_id:foreign;response_id:resp_foreign").decode(), + id="foreign-provider-wrapper", + ), + pytest.param("resp_" + base64.b64encode(b"custom;id;response").decode(), id="unmanaged-semicolons"), + pytest.param( + "resp_" + base64.b64encode(b"openai;foreign;resp_unmanaged").decode(), id="openai-without-envelope-markers" + ), + ), +) +def test_guard_sdk_custom_host_preserves_unmanaged_comparison_ids(extra: bool, comparison_id: str) -> None: + def respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/v1/responses") + assert _JSON_OBJECT.validate_json(request.body) == { + "model": _MODEL, + "input": "custom comparison", + "prompt_cache_options": {"comparison_response_id": comparison_id, "mode": "implicit", "ttl": "30m"}, + }, request.body + return Reply(body=_response("resp_custom_result")) + + with wire_server(respond) as wire: + options: Final = { + "comparison_response_id": comparison_id, + "mode": "implicit", + "ttl": "30m", + } + params: Final = ( + {"extra_body": {"prompt_cache_options": options}} if extra else {"prompt_cache_options": options} + ) + original: Final = deepcopy(params) + result: Final = litellm.responses( + model=f"openai/{_MODEL}", input="custom comparison", api_base=f"{wire.url}/v1", api_key=_API_KEY, **params + ) + assert isinstance(result, ResponsesAPIResponse), result + assert result.status == "completed", result + assert params == original + assert len(wire.drain()) == 1 + + +@pytest.mark.parametrize("provider", ("azure", "xai")) +@pytest.mark.parametrize("extra", (False, True)) +def test_guard_inherited_responses_providers_preserve_comparison_wrappers(provider: str, extra: bool) -> None: + raw_id: Final = f"resp_inherited_{uuid.uuid4().hex}" + + def seed(request: Request) -> Reply: + assert _JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": "seed inherited comparison"} + return Reply(body=_response(raw_id)) + + with wire_server(seed) as seed_wire: + seeded: Final = litellm.responses( + model=f"openai/{_MODEL}", + input="seed inherited comparison", + api_base=f"{seed_wire.url}/v1", + api_key=_API_KEY, + ) + assert isinstance(seeded, ResponsesAPIResponse), seeded + assert seeded.id != raw_id + assert len(seed_wire.drain()) == 1 + + def respond(request: Request) -> Reply: + assert request.method == "POST", request + expected_target: Final = ( + "/openai/responses?api-version=2025-03-01-preview" if provider == "azure" else "/v1/responses" + ) + assert request.target == expected_target, request.target + assert _JSON_OBJECT.validate_json(request.body) == { + "model": _MODEL, + "input": "inherited comparison", + "prompt_cache_options": {"comparison_response_id": seeded.id}, + }, request.body + return Reply(body=_response(raw_id)) + + with wire_server(respond) as wire: + options: Final = {"comparison_response_id": seeded.id} + params: Final = ( + {"extra_body": {"prompt_cache_options": options}} if extra else {"prompt_cache_options": options} + ) + result: Final = litellm.responses( + model=f"{provider}/{_MODEL}", + input="inherited comparison", + api_base=wire.url if provider == "azure" else f"{wire.url}/v1", + api_key=_API_KEY, + api_version="2025-03-01-preview" if provider == "azure" else None, + **params, + ) + assert isinstance(result, ResponsesAPIResponse), result + assert result.status == "completed", result + assert len(wire.drain()) == 1 + + def _register_in_group(scenario: Scenario, wire: Wire, group: str, order: int | None = None) -> str: existing: Final = _deployments_in_group(scenario.gateway, group) created: Final = scenario.gateway.post( diff --git a/tests/integration/providers/test_responses_lifecycle_wire.py b/tests/integration/providers/test_responses_lifecycle_wire.py index 02269885b44..ac5b4dc6d5c 100644 --- a/tests/integration/providers/test_responses_lifecycle_wire.py +++ b/tests/integration/providers/test_responses_lifecycle_wire.py @@ -203,9 +203,6 @@ def test_retrieve_routes_to_the_deployment_that_created_the_response( @pytest.mark.parametrize("prefix", _ALIASES) def test_retrieve_returns_the_id_the_client_holds(gateway: Gateway, prefix: str) -> None: - pytest.skip( - "BUG: GET /v1/responses/{id} returns a re-encrypted id that differs byte-for-byte from the id the client sent" - ) raw_id: Final = f"resp_retrieve_id_{uuid.uuid4().hex}" def respond(request: Request) -> Reply: @@ -697,7 +694,6 @@ def test_delete_returns_the_client_held_id_not_the_provider_id( gateway: Gateway, prefix: str, ) -> None: - pytest.skip("BUG: DELETE /v1/responses/{id} returns the raw provider id instead of the proxy-issued id") raw_id: Final = f"resp_delete_id_{uuid.uuid4().hex}" def respond(request: Request) -> Reply: diff --git a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py index 3175f2ec075..00f861678cb 100644 --- a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py @@ -25,7 +25,7 @@ from pydantic import JsonValue, TypeAdapter _INLINE_TOOLS: Final = "inline-tools-2026-09-15" _THINKING: Final = "interleaved-thinking-2025-05-14" -_VERTEX_UNSUPPORTED: Final = "effort-2025-11-24" +_VERTEX_UNSUPPORTED: Final = "skills-2025-10-02" _UNKNOWN: Final = "not-a-beta-2099-01-01" _BACKEND: Final = "claude-opus-5-5" _PROJECT: Final = "scripted-project" diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 4e31122440d..db81ff70263 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -1488,7 +1488,7 @@ def test_vertex_chat_json_schema_on_a_model_rejecting_forced_tools_maps_reasonin assert _reasoning_fields(structured_sent) == _reasoning_fields(plain_sent), (structured_sent.body, plain_sent.body) -def test_bedrock_invoke_json_schema_sends_the_json_tool_unforced_when_the_model_rejects_forced_tools( +def test_bedrock_invoke_json_schema_sends_native_structured_output_when_the_model_rejects_forced_tools( gateway: Gateway, ) -> None: marker, marker_hex = _marker("plain") @@ -1499,7 +1499,8 @@ def test_bedrock_invoke_json_schema_sends_the_json_tool_unforced_when_the_model_ assert _chat_content(response) == _payload(marker_hex), response.text sent: Final = _one_delivery(wire, marker_hex) assert sent.target == f"/model/{backend}/invoke", sent.target - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + assert object_value(_folded_format(sent)["schema"])["properties"] == _SCHEMA["properties"], sent.body + assert "tools" not in sent.body, sent.body assert "tool_choice" not in sent.body, sent.body diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index 4e2cab56d27..c552e11c22a 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -587,6 +587,11 @@ model_list: model: perplexity/pplx-decider-v1-27b api_base: http://127.0.0.1:8191 api_key: synthetic-perplexity-key + - model_name: parallel_ai/decision + litellm_params: + model: parallel_ai/decision + api_base: http://127.0.0.1:8191 + api_key: synthetic-parallel-key - model_name: typesafe/jev-1.13.0 litellm_params: model: typesafe/jev-1.13.0 diff --git a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py index e6ec3444369..ef2a9e5b72f 100644 --- a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py +++ b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py @@ -303,7 +303,7 @@ def scenario_api_base(scenario: Scenario, scenario_id: str, response: StoredResp def deployment( scenario: Scenario, group: str, - api_base: str, + api_base: str | None, *, order: int | None, model: str = RESPONSES_MODEL, @@ -316,7 +316,7 @@ def deployment( "litellm_params": { "model": model, "api_key": PROVIDER_KEY, - "api_base": api_base, + **({} if api_base is None else {"api_base": api_base}), **({} if order is None else {"order": order}), **extra_params, }, @@ -863,6 +863,38 @@ def test_hop_inside_one_encryption_boundary_keeps_the_reasoning(gateway: Gateway assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies +def test_hop_between_same_key_deployments_on_the_providers_default_endpoint_keeps_the_reasoning( + gateway: Gateway, tmp_path: Path +) -> None: + shared: Final = f"default-{uuid.uuid4().hex[:12]}" + with ( + scratch_database() as database, + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database, "OPENAI_BASE_URL": f"{gateway.upstream_url}/{shared}"}, + remove_environment=("DIRECT_URL", "DATABASE_URL_READ_REPLICA", "OPENAI_API_BASE"), + ) as owned, + owned.gateway.scenario() as scenario, + ): + owned.gateway.post("/config/update", {"router_settings": {"num_retries": 0}}) + one: Final = scenario_api_base(scenario, f"{shared}-o1", failing()).rsplit("/", 1)[-1] + two: Final = scenario_api_base(scenario, f"{shared}-o2", healthy_json()).rsplit("/", 1)[-1] + order_one: Final = deployment( + scenario, shared, None, order=1, extra_params={"extra_headers": {"x-scripted-scenario": one}} + ) + order_two: Final = deployment( + scenario, shared, None, order=2, extra_params={"extra_headers": {"x-scripted-scenario": two}} + ) + mark: Final = marker() + answer: Final = responses_answer( + owned.gateway, HTTPX, Target(shared, frozenset({order_one, order_two})), history(mark) + ) + assert_served_by(answer, order_two) + bodies: Final = bodies_for(observed(owned.gateway), shared) + assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies + + def assert_cross_group_hop_stripped(gateway: Gateway, source_scenario: str, target_scenario: str, mark: str) -> None: records: Final = observed(gateway) assert [body.get("input") for body in bodies_for(records, source_scenario)] == [history(mark)], records diff --git a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py index da091b1c12e..54a6df02edf 100644 --- a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py +++ b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py @@ -10,6 +10,7 @@ from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper _JSON_TOOL: Final = "json_tool_call" @@ -122,12 +123,10 @@ def _credentials(model: str) -> dict[str, str]: @pytest.mark.parametrize("model", [_ANTHROPIC_MODEL, _BEDROCK_INVOKE_MODEL], ids=["anthropic", "bedrock_invoke"]) @pytest.mark.parametrize( ("shape", "cut"), - [("plain", False), ("plain", True), ("wrapped", True)], - ids=["plain_complete", "plain_cut", "wrapped_cut"], + [("plain", False), ("wrapped", False), ("plain", True), ("wrapped", True)], + ids=["plain_complete", "wrapped_complete", "plain_cut", "wrapped_cut"], ) -def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whole_answer( - model: str, shape: str, cut: bool -) -> None: +def test_sync_json_tool_stream_requires_a_provider_finish_reason(model: str, shape: str, cut: bool) -> None: marker_hex: Final = uuid.uuid4().hex with wire_server(_peer(marker_hex, shape, cut)) as wire: response: Final = litellm.completion( @@ -141,6 +140,11 @@ def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whol **_credentials(model), ) assert isinstance(response, CustomStreamWrapper), type(response) + if cut: + with pytest.raises(MidStreamFallbackError, match="Stream ended before a finish_reason was received"): + tuple(response) + assert len(wire.drain()) == 1 + return text: Final = "".join(str(chunk.choices[0].delta.content or "") for chunk in tuple(response)) assert len(wire.drain()) == 1 assert json.loads(text) == _payload(marker_hex), text diff --git a/tests/integration/sdk/test_openai_realtime_http_ca_bundle_wire.py b/tests/integration/sdk/test_openai_realtime_http_ca_bundle_wire.py new file mode 100644 index 00000000000..358b36cd2b0 --- /dev/null +++ b/tests/integration/sdk/test_openai_realtime_http_ca_bundle_wire.py @@ -0,0 +1,95 @@ +import json +from collections.abc import Awaitable +from datetime import datetime, timezone +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Request, wire_server + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.realtime_api import main as realtime_main + +UPSTREAM_PATHS: Final = MappingProxyType( + { + "client_secrets": "/v1/realtime/client_secrets", + "translation_client_secrets": "/v1/realtime/translations/client_secrets", + "calls": "/v1/realtime/calls", + "translation_calls": "/v1/realtime/translations/calls", + } +) + + +def _logging(route: str) -> Logging: + return Logging( + model="gpt-realtime", + messages=[], + stream=False, + call_type="acreate_realtime_client_secret", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id=f"test-realtime-{route}", + function_id=f"test-realtime-{route}", + ) + + +def _route_call(route: str, api_base: str) -> Awaitable[httpx.Response]: + match route: + case "client_secrets": + return realtime_main.acreate_realtime_client_secret.__wrapped__( + model="gpt-realtime", + custom_llm_provider="openai", + api_key="sk-test", + api_base=api_base, + litellm_logging_obj=_logging(route), + ) + case "translation_client_secrets": + return realtime_main.acreate_realtime_translation_client_secret.__wrapped__( + model="gpt-realtime-translate", + custom_llm_provider="openai", + api_key="sk-test", + api_base=api_base, + litellm_logging_obj=_logging(route), + ) + case "calls": + return realtime_main.arealtime_calls.__wrapped__( + openai_ephemeral_key="ek_test", + sdp_body=b"v=0", + model="gpt-realtime", + custom_llm_provider="openai", + api_base=api_base, + litellm_logging_obj=_logging(route), + ) + case _: + return realtime_main.arealtime_translation_calls.__wrapped__( + openai_ephemeral_key="ek_test", + sdp_body=b"v=0", + model="gpt-realtime-translate", + custom_llm_provider="openai", + api_base=api_base, + litellm_logging_obj=_logging(route), + ) + + +def _ephemeral_key(_: Request) -> Reply: + return Reply(body=json.dumps({"value": "ek_test"}).encode()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", tuple(UPSTREAM_PATHS)) +async def test_openai_realtime_http_routes_trust_the_ca_bundle_set_in_litellm_ssl_verify( + route: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + certificate, key = write_self_signed_cert(tmp_path) + monkeypatch.setattr(litellm, "ssl_verify", str(certificate)) + + with wire_server(_ephemeral_key, tls=server_context(certificate, key)) as upstream: + response: Final = await _route_call(route, upstream.url) + targets: Final = tuple(request.target for request in upstream.drain()) + + assert response.status_code == 200 + assert response.json() == {"value": "ek_test"} + assert targets == (UPSTREAM_PATHS[route],) diff --git a/tests/integration/security/test_files_batches_access.py b/tests/integration/security/test_files_batches_access.py index 119055c64c7..93ccc742d58 100644 --- a/tests/integration/security/test_files_batches_access.py +++ b/tests/integration/security/test_files_batches_access.py @@ -25,7 +25,7 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa 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 +from integration._support.process import OwnedState, owned_proxy from integration._support.tls import server_context, write_self_signed_cert from integration._support.wire import Reply, Request, Wire, wire_server from openai import OpenAI, PermissionDeniedError @@ -250,15 +250,20 @@ def test_stranger_cannot_read_download_or_delete_another_teams_managed_file(gate assert _drained_other_than_model_list_probes(wire) == () -def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: Gateway) -> None: +@pytest.mark.timeout(180) +def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: Gateway, tmp_path: Path) -> None: key: Final = "provider-key-" + uuid.uuid4().hex[:8] - with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario: + with ( + owned_proxy(gateway, tmp_path, {"PROXY_BATCH_POLLING_ENABLED": "false"}, state=OwnedState.PRIVATE) as candidate, + wire_server(_openai_backend()) as wire, + candidate.scenario() as scenario, + ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key) - _wait_until_every_worker_serves(gateway, model) + _wait_until_every_worker_serves(candidate, model) owner: Final = _member(scenario, model) stranger: Final = _member(scenario, model) - uploaded: Final = gateway.request_multipart( + uploaded: Final = candidate.request_multipart( "/v1/files", {"purpose": "batch", "target_model_names": model}, {"file": (INPUT_FILENAME, _jsonl(model), "application/jsonl")}, @@ -266,7 +271,7 @@ def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: ) assert uploaded.status_code == 200, uploaded.text managed_input: Final = string_value(_json(uploaded)["id"]) - created: Final = gateway.request( + created: Final = candidate.request( "POST", "/v1/batches", {"input_file_id": managed_input, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, @@ -277,13 +282,13 @@ def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: wire.drain() expected_message: Final = f"User {stranger.user} does not have access to the object {unified_batch}" - stranger_get: Final = gateway.request("GET", f"/v1/batches/{unified_batch}", key=stranger.key) + stranger_get: Final = candidate.request("GET", f"/v1/batches/{unified_batch}", key=stranger.key) assert _error_message(stranger_get, 403) == expected_message, stranger_get.text - stranger_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=stranger.key) + stranger_cancel: Final = candidate.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=stranger.key) assert _error_message(stranger_cancel, 403) == expected_message, stranger_cancel.text assert _drained_other_than_model_list_probes(wire) == (), "stranger batch calls must not reach the provider" - owner_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=owner.key) + owner_cancel: Final = candidate.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=owner.key) assert owner_cancel.status_code == 200, owner_cancel.text assert string_value(_json(owner_cancel)["id"]) == unified_batch, owner_cancel.text (cancel_request,) = _drained_other_than_model_list_probes(wire) diff --git a/tests/integration/spend/test_end_user_budget_enforcement.py b/tests/integration/spend/test_end_user_budget_enforcement.py index 02c8aee8f5e..92259d463b9 100644 --- a/tests/integration/spend/test_end_user_budget_enforcement.py +++ b/tests/integration/spend/test_end_user_budget_enforcement.py @@ -258,9 +258,20 @@ def _observed(upstream: httpx.Client) -> list[JsonValue]: return requests -def _denial_error(denied: httpx.Response) -> dict[str, JsonValue]: +def _denial_error(denied: httpx.Response, route: str) -> dict[str, JsonValue]: assert denied.status_code == 422, denied.text - return object_value(object_value(denied.json())["error"]) + if route == "/v1/messages": + body: Final = object_value(denied.json()) + assert body["type"] == "error", denied.text + anthropic_error: Final = object_value(body["error"]) + assert anthropic_error["type"] == "invalid_request_error", denied.text + assert "param" not in anthropic_error and "code" not in anthropic_error, denied.text + return anthropic_error + openai_error: Final = object_value(object_value(denied.json())["error"]) + assert openai_error["type"] == "budget_exceeded", denied.text + assert openai_error["param"] is None + assert openai_error["code"] == "422" + return openai_error def _prove_end_user_budget_lockout( @@ -304,8 +315,7 @@ def _prove_end_user_budget_lockout( upstream.delete("/__observations").raise_for_status() denied: Final = _send(gateway, spelling, deployment, key, end_user_id, f"denied {end_user_id}") assert _observed(upstream) == [], denied.text - error: Final = _denial_error(denied) - assert error["type"] == "budget_exceeded", denied.text + error: Final = _denial_error(denied, spelling.route) matched: Final = re.fullmatch( rf"ExceededBudget: End User={re.escape(end_user_id)} over budget\. " rf"Spend=([-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?), Budget=0\.05", diff --git a/tests/integration/spend/test_end_user_spend_without_proxy_user.py b/tests/integration/spend/test_end_user_spend_without_proxy_user.py index 94884717092..725c809006a 100644 --- a/tests/integration/spend/test_end_user_spend_without_proxy_user.py +++ b/tests/integration/spend/test_end_user_spend_without_proxy_user.py @@ -120,6 +120,18 @@ def _assert_budget_refusal(response: httpx.Response, end_user: str) -> None: ] +def _assert_messages_budget_refusal(response: httpx.Response, end_user: str) -> None: + assert response.status_code == 422, response.text + body: Final = _json_object(response.content) + assert body["type"] == "error", response.text + error: Final = object_value(body["error"]) + assert error["type"] == "invalid_request_error" + assert "param" not in error and "code" not in error, response.text + assert str(error["message"]).startswith(f"ExceededBudget: End User={end_user} over budget. Spend="), error[ + "message" + ] + + def _json_object(raw: bytes) -> dict[str, JsonValue]: return _JSON_OBJECT.validate_json(raw) @@ -345,7 +357,7 @@ def test_anthropic_messages_metadata_user_id_is_charged_and_budget_enforced(gate json={**body, "messages": [{"role": "user", "content": f"again {uuid.uuid4().hex}"}]}, headers={"x-api-key": key, "anthropic-version": "2023-06-01"}, ) - _assert_budget_refusal(denied, end_user) + _assert_messages_budget_refusal(denied, end_user) assert _observed(upstream) == [] diff --git a/tests/integration/spend/test_normalized_error_long_message.py b/tests/integration/spend/test_normalized_error_long_message.py index a364327290f..23200c0450e 100644 --- a/tests/integration/spend/test_normalized_error_long_message.py +++ b/tests/integration/spend/test_normalized_error_long_message.py @@ -30,6 +30,7 @@ BENIGN_288KB_MODEL: Final = "x" * len(CRAFTED_MODEL) CRAFTED_SLOWDOWN_FACTOR: Final = 10 FAST_SECONDS: Final = 10.0 LIVELINESS_STALL_BUDGET_SECONDS: Final = 5.0 +BURST_STALL_BUDGET_SECONDS: Final = 5.0 ROW_SECONDS: Final = 70 CHAT: Final = "/v1/chat/completions" MESSAGES: Final = "/v1/messages" @@ -349,10 +350,16 @@ def _burst_call(client: httpx.Client, label: str, path: str, body: Mapping[str, def _burst( - client: httpx.Client, happy_model: str, happy_key: str, open_key: str, during: Callable[[], None] + client: httpx.Client, + happy_model: str, + happy_key: str, + open_key: str, + during: Callable[[], None], + *, + rejected_model: str = CRAFTED_MODEL, ) -> tuple[_BurstResult, ...]: crafted: Final = tuple( - (f"crafted-{path}-{index}", path, _body(path, CRAFTED_MODEL, uuid.uuid4().hex, stream=index % 2 == 1), open_key) + (f"crafted-{path}-{index}", path, _body(path, rejected_model, uuid.uuid4().hex, stream=index % 2 == 1), open_key) for path in (CHAT, MESSAGES, RESPONSES) for index in range(4) ) @@ -377,6 +384,18 @@ def _burst( return tuple(future.result() for future in first + second) +def _assert_burst_rejections_fast(results: tuple[_BurstResult, ...], control: tuple[_BurstResult, ...]) -> None: + crafted: Final = tuple(result for result in results if result.label.startswith("crafted") and result.status) + benign: Final = tuple(result for result in control if result.label.startswith("crafted")) + assert crafted and all(result.status == 400 for result in crafted), crafted + assert benign and all(result.status == 400 for result in benign), benign + assert_not_delayed_by( + (result.seconds for result in crafted), + control=(result.seconds for result in benign), + injected_delay=BURST_STALL_BUDGET_SECONDS, + ) + + def _assert_rows_land_exactly_once(results: tuple[_BurstResult, ...], prefix: str) -> None: landed: Final = tuple(result for result in results if result.label.startswith(prefix) and result.status == 200) assert landed, results @@ -405,20 +424,18 @@ def test_killing_one_worker_mid_burst_leaves_the_other_serving_crafted_and_happy ): model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) key: Final = scenario.key(models=[model]) + open_key: Final = scenario.key() + control: Final = _burst(client, model, key, open_key, lambda: None, rejected_model=BENIGN_288KB_MODEL) + assert all(result.status == (400 if result.label.startswith("crafted") else 200) for result in control), control workers: Final = owned.worker_pids() assert len(workers) == 2, workers victim: Final = min(workers) with owned.replacing_worker(victim) as kill: - results: Final = _burst(client, model, key, scenario.key(), kill) + results: Final = _burst(client, model, key, open_key, kill) dropped: Final = tuple(result for result in results if result.status is None) assert len(dropped) < len(results), results - crafted: Final = tuple( - result for result in results if result.label.startswith("crafted") and result.status - ) - assert crafted and all( - result.status == 400 and result.seconds < FAST_SECONDS for result in crafted - ), crafted + _assert_burst_rejections_fast(results, control) late: Final = tuple(result for result in results if result.label.startswith("late")) assert all(result.status == 200 for result in late), late _assert_rows_land_exactly_once(results, "late") @@ -446,7 +463,11 @@ def test_upstream_returning_503_mid_burst_logs_every_failure_with_its_own_cluste ): model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) key: Final = scenario.key(models=[model]) - results: Final = _burst(client, model, key, scenario.key(), lambda: None) + open_key: Final = scenario.key() + control: Final = _burst(client, model, key, open_key, lambda: None, rejected_model=BENIGN_288KB_MODEL) + assert all(result.status == (400 if result.label.startswith("crafted") else 503) for result in control), control + wire.drain() + results: Final = _burst(client, model, key, open_key, lambda: None) assert all(result.status is not None for result in results), results happy: Final = tuple(result for result in results if result.label.startswith("happy")) assert all(result.status == 503 for result in happy), happy @@ -456,8 +477,7 @@ def test_upstream_returning_503_mid_burst_logs_every_failure_with_its_own_cluste assert all(any(f"normalized error audit {result.label}" in body for body in seen) for result in happy + late), ( seen ) - crafted: Final = tuple(result for result in results if result.label.startswith("crafted")) - assert all(result.status == 400 and result.seconds < FAST_SECONDS for result in crafted), crafted + _assert_burst_rejections_fast(results, control) codes: Final = { result.label: _error_information(string_value(result.call_id))["normalized_error"] for result in results } diff --git a/tests/integration/spend/test_partial_stream_live_budget.py b/tests/integration/spend/test_partial_stream_live_budget.py index 89fb66c9bbf..cba51873e2d 100644 --- a/tests/integration/spend/test_partial_stream_live_budget.py +++ b/tests/integration/spend/test_partial_stream_live_budget.py @@ -346,5 +346,15 @@ def test_a_broken_messages_stream_counts_toward_the_budget(gateway: Gateway) -> _broken_messages_stream(gateway, model, key) after: Final = _messages(gateway, model, key, streamed=False) - assert_budget_refusal(after) + assert after.status_code == 422, after.text + body: Final = JSON_OBJECT.validate_json(after.content) + error: Final = JSON_OBJECT.validate_python(body["error"]) + assert body == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": error["message"], + "litellm_error_type": "budget_exceeded", + }, + }, after.text assert _rows(gateway, key, 2) == (metered, StatusRow(status="failure", spend=0.0)) diff --git a/tests/integration/spend/test_responses_bridge_spend_once.py b/tests/integration/spend/test_responses_bridge_spend_once.py index 9eddd832ca0..cc2a0af58a7 100644 --- a/tests/integration/spend/test_responses_bridge_spend_once.py +++ b/tests/integration/spend/test_responses_bridge_spend_once.py @@ -1,5 +1,4 @@ import json -from base64 import b64decode from collections.abc import Callable from hashlib import sha256 from typing import Final @@ -100,8 +99,7 @@ def test_bridged_chat_completion_above_log_offload_threshold_logs_and_charges_on call_rows: Final = [row for row in rows if row["litellm_call_id"] == call_id] assert len(call_rows) == 1, rows assert float(str(call_rows[0]["spend"])) == pytest.approx(EXPECTED_SPEND), rows - logged_id: Final = b64decode(str(call_rows[0]["request_id"]).removeprefix("resp_")).decode() - assert logged_id.startswith("litellm:") and f";response_id:{upstream_id}" in logged_id, rows + assert call_rows[0]["request_id"] == upstream_id, rows key_spend: Final = eventually( lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), diff --git a/tests/integration/spend/test_responses_stream_spend_row_id.py b/tests/integration/spend/test_responses_stream_spend_row_id.py new file mode 100644 index 00000000000..58ed9f7785f --- /dev/null +++ b/tests/integration/spend/test_responses_stream_spend_row_id.py @@ -0,0 +1,465 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final +from uuid import uuid4 + +import httpx +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, list_value, object_value, string_value +from integration._support.process import owned_proxy +from integration._support.responses_stream import frame, response_object +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_NATIVE_MODEL: Final = "openai/gpt-5.6" +_BRIDGED_MODEL: Final = "deepseek/deepseek-v4-pro" +_ANSWER: Final = "streamed answer" +_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 30, "output_tokens": 5, "total_tokens": 35} +_JSON_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def _sse(payload: dict[str, JsonValue]) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def _native_completed(upstream_id: str) -> dict[str, JsonValue]: + message: Final[dict[str, JsonValue]] = { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": _ANSWER, "annotations": []}], + } + return response_object(upstream_id, "completed", output=[message], usage=_USAGE) + + +def _native_stream(upstream_id: str) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=( + frame( + { + "type": "response.created", + "sequence_number": 0, + "response": response_object(upstream_id, "in_progress"), + } + ), + frame( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": _ANSWER, + } + ), + frame({"type": "response.completed", "sequence_number": 2, "response": _native_completed(upstream_id)}), + ), + ) + + +def _chat_stream(upstream_id: str, model: str) -> Reply: + def chunk(delta: dict[str, JsonValue], finish_reason: str | None) -> dict[str, JsonValue]: + return { + "id": upstream_id, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + + return Reply( + content_type="text/event-stream", + chunks=( + _sse(chunk({"role": "assistant", "content": _ANSWER}, None)), + _sse( + { + **chunk({}, "stop"), + "usage": {"prompt_tokens": 30, "completion_tokens": 5, "total_tokens": 35}, + } + ), + b"data: [DONE]\n\n", + ), + ) + + +def _respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + if request.method == "GET": + stored_id: Final = request.target.removeprefix("/v1/responses/") + assert stored_id.startswith("resp_upstream-"), request.target + return Reply(body=json.dumps(_native_completed(stored_id)).encode()) + if request.target.endswith("/cancel"): + cancelled_id: Final = request.target.removeprefix("/v1/responses/").removesuffix("/cancel") + assert cancelled_id.startswith("resp_upstream-"), request.target + return Reply(body=json.dumps(response_object(cancelled_id, "cancelled")).encode()) + body: Final = JSON_OBJECT.validate_json(request.body) + upstream_id: Final = f"upstream-{uuid4().hex}" + if request.target == "/v1/responses" and body.get("stream") is not True: + return Reply(body=json.dumps(_native_completed(f"resp_{upstream_id}")).encode()) + assert body.get("stream") is True, body + if request.target == "/v1/responses": + return _native_stream(f"resp_{upstream_id}") + assert request.target == "/v1/chat/completions", request.target + return _chat_stream(f"chatcmpl-{upstream_id}", string_value(body["model"])) + + +def _terminal_response_id( + gateway: Gateway, key: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None +) -> str: + response: Final = gateway.request("POST", "/v1/responses", body, key=key, headers=headers) + assert response.status_code == 200, response.text + frames: Final = tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert frames and frames[-1]["type"] == "response.completed", response.text + advertised: Final = tuple( + dict.fromkeys(string_value(object_value(event["response"])["id"]) for event in frames if "response" in event) + ) + assert len(advertised) == 1, advertised + return advertised[0] + + +def _created_response_id( + gateway: Gateway, key: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None +) -> str: + if body.get("stream") is True: + return _terminal_response_id(gateway, key, body, headers) + response: Final = gateway.request("POST", "/v1/responses", body, key=key, headers=headers) + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _read_and_cancel( + gateway: Gateway, key: str, response_id: str, headers: Mapping[str, str] | None = None +) -> tuple[httpx.Response, ...]: + responses: Final = tuple( + gateway.request(method, path, key=key, headers=headers) + for method, path in ( + ("GET", f"/v1/responses/{response_id}"), + ("GET", f"/v1/responses/{response_id}"), + ("POST", f"/v1/responses/{response_id}/cancel"), + ) + ) + assert [response.status_code for response in responses] == [200, 200, 200], [r.text for r in responses] + return responses + + +def _spend_logs_lookup(gateway: Gateway, request_id: str) -> list[JsonValue]: + response: Final = gateway.request("GET", "/spend/logs", params={"request_id": request_id}) + assert response.status_code == 200, response.text + return [row["request_id"] for row in _JSON_LIST.validate_json(response.content)] + + +def _request_ids_for_key(gateway: Gateway, key: str, count: int) -> list[str]: + rows: Final = eventually( + lambda: gateway.rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "endTime"', + (sha256(key.encode()).hexdigest(),), + ), + lambda found: len(found) >= count, + seconds=70, + ) + return [string_value(row["request_id"]) for row in rows] + + +def _logs_page_search(gateway: Gateway, client_id: str) -> list[str]: + now: Final = datetime.now(timezone.utc) + page: Final = gateway.get( + "/spend/logs/ui", + params={ + "search": client_id, + "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + }, + ) + return [string_value(object_value(row)["request_id"]) for row in list_value(page["data"])] + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("litellm_model", [_NATIVE_MODEL, _BRIDGED_MODEL], ids=["native", "bridge"]) +def test_streamed_responses_spend_row_is_filed_under_the_id_the_client_saw( + gateway: Gateway, litellm_model: str +) -> None: + with wire_server(_respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=litellm_model, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + + client_id: Final = _terminal_response_id( + gateway, key, {"model": model, "input": f"marker-{uuid4().hex}", "stream": True} + ) + + assert client_id.startswith("resp_"), client_id + assert _request_ids_for_key(gateway, key, 1) == [client_id] + assert _logs_page_search(gateway, client_id) == [client_id] + assert _spend_logs_lookup(gateway, client_id) == [client_id] + + +def _proxy_config( + directory: Path, + general_settings: Mapping[str, JsonValue], + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "responses-spend.yaml" + path.write_text( + yaml.safe_dump( + { + **stock, + "general_settings": {**object_value(stock["general_settings"]), **general_settings}, + "litellm_settings": {**object_value(stock["litellm_settings"]), **litellm_settings}, + } + ) + ) + return path + + +def _prompt_storing_config(directory: Path) -> Path: + return _proxy_config(directory, {"store_prompts_in_spend_logs": True}) + + +def test_a_follow_up_turn_on_a_streamed_bridged_response_replays_the_first_turn( + gateway: Gateway, tmp_path: Path +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + model: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + first_id: Final = _terminal_response_id(storing, key, {"model": model, "input": first_prompt, "stream": True}) + assert _request_ids_for_key(storing, key, 1) == [first_id] + wire.drain() + + second_id: Final = _terminal_response_id( + storing, + key, + {"model": model, "input": "second turn", "previous_response_id": first_id, "stream": True}, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed + assert _request_ids_for_key(storing, key, 2) == [first_id, second_id] + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize("stream_first_turn", [True, False], ids=["streamed-first-turn", "non-streamed-first-turn"]) +def test_a_follow_up_through_the_id_a_retrieve_returned_replays_the_first_turn_once( + gateway: Gateway, tmp_path: Path, stream_first_turn: bool +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id( + storing, key, {"model": native, "input": first_prompt, "stream": stream_first_turn} + ) + + retrieved: Final = storing.request("GET", f"/v1/responses/{first_id}", key=key) + + assert retrieved.status_code == 200, retrieved.text + retrieved_id: Final = string_value(JSON_OBJECT.validate_json(retrieved.content)["id"]) + assert retrieved_id == first_id + stored: Final = _request_ids_for_key(storing, key, 2) + assert stored.count(first_id) == 1, stored + assert len(set(stored)) == 2, stored + assert _spend_logs_lookup(storing, first_id) == [first_id] + wire.drain() + + _terminal_response_id( + storing, + key, + {"model": bridged, "input": "second turn", "previous_response_id": retrieved_id, "stream": True}, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed + + +def _replayed_turns(wire_body: bytes) -> list[tuple[JsonValue, JsonValue]]: + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire_body)["messages"]) + return [(message["role"], message["content"]) for message in replayed] + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize("reader", ["proxy-admin", "master-key"]) +def test_a_reader_other_than_the_creator_continues_through_the_id_its_retrieve_returned( + gateway: Gateway, tmp_path: Path, reader: str +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + creator: Final = scenario.key(user_id=scenario.user(), models=[native, bridged]) + reader_key: Final = ( + scenario.key(user_id=scenario.user(user_role="proxy_admin"), models=[native, bridged]) + if reader == "proxy-admin" + else storing.key + ) + first_id: Final = _terminal_response_id( + storing, creator, {"model": native, "input": first_prompt, "stream": True} + ) + assert _request_ids_for_key(storing, creator, 1) == [first_id] + + retrieved: Final = storing.request("GET", f"/v1/responses/{first_id}", key=reader_key) + + assert retrieved.status_code == 200, retrieved.text + retrieved_id: Final = string_value(JSON_OBJECT.validate_json(retrieved.content)["id"]) + wire.drain() + _terminal_response_id( + storing, + reader_key, + {"model": bridged, "input": "second turn", "previous_response_id": retrieved_id, "stream": True}, + ) + assert _replayed_turns(wire.drain()[-1].body) == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ] + assert retrieved_id == first_id + + +@pytest.mark.timeout(180) +def test_reads_and_cancels_sent_with_a_session_header_add_no_replayed_turns(gateway: Gateway, tmp_path: Path) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + session: Final = {"x-litellm-session-id": f"session-{uuid4().hex}"} + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id(storing, key, {"model": native, "input": first_prompt}, session) + + _read_and_cancel(storing, key, first_id, session) + + assert len(_request_ids_for_key(storing, key, 4)) == 4 + wire.drain() + _terminal_response_id( + storing, + key, + {"model": bridged, "input": "second turn", "previous_response_id": first_id, "stream": True}, + session, + ) + assert _replayed_turns(wire.drain()[-1].body) == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ] + + +@pytest.mark.timeout(180) +def test_each_read_and_cancel_reaches_callbacks_and_spend_logs_under_an_id_of_its_own( + gateway: Gateway, tmp_path: Path +) -> None: + logged_call_types: Final = frozenset({"aresponses", "aget_responses", "acancel_responses"}) + with ( + wire_server(_respond) as wire, + wire_server(lambda _: Reply()) as sink, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": sink.url}, + config=_proxy_config(tmp_path, {}, {"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}), + ) as logging_proxy, + logging_proxy.scenario() as scenario, + ): + model: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + first_id: Final = _created_response_id(logging_proxy, key, {"model": model, "input": f"marker-{uuid4().hex}"}) + + reads: Final = _read_and_cancel(logging_proxy, key, first_id) + + assert [string_value(JSON_OBJECT.validate_json(read.content)["id"]) for read in reads] == [first_id] * 3 + batches: Final[list[bytes]] = [] # mutable-ok: drain() consumes, later polls keep earlier batches + + def delivered() -> tuple[tuple[str, str], ...]: + batches.extend(batch.body for batch in sink.drain()) + return tuple( + (string_value(event["call_type"]), string_value(event["id"])) + for event in (object_value(item) for batch in batches for item in _JSON_LIST.validate_json(batch)) + if event["call_type"] in logged_call_types + ) + + events: Final = eventually(delivered, lambda found: len(found) >= 4, seconds=30) + read_ids: Final = [logged_id for call_type, logged_id in events if call_type != "aresponses"] + assert [logged_id for call_type, logged_id in events if call_type == "aresponses"] == [first_id], events + assert len(read_ids) == 3, events + assert first_id not in read_ids, events + assert len(set(read_ids)) == 3, events + assert sorted(_request_ids_for_key(logging_proxy, key, 4)) == sorted([first_id, *read_ids]) + + +@pytest.mark.timeout(240) +def test_with_id_security_off_reads_and_cancels_add_no_rows_and_no_replayed_turns( + gateway: Gateway, tmp_path: Path +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + session: Final = {"x-litellm-session-id": f"session-{uuid4().hex}"} + with ( + wire_server(_respond) as wire, + owned_proxy( + gateway, + tmp_path, + {}, + config=_proxy_config( + tmp_path, {"store_prompts_in_spend_logs": True, "disable_responses_id_security": True} + ), + ) as insecure, + insecure.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id(insecure, key, {"model": native, "input": first_prompt}, session) + + reads: Final = _read_and_cancel(insecure, key, first_id, session) + + assert [string_value(JSON_OBJECT.validate_json(read.content)["id"]) for read in reads] == [first_id] * 3 + later_id: Final = _created_response_id(insecure, key, {"model": native, "input": "a later request"}) + assert _request_ids_for_key(insecure, key, 2) == [first_id, later_id] + wire.drain() + + _terminal_response_id( + insecure, + key, + {"model": bridged, "input": "second turn", "previous_response_id": first_id, "stream": True}, + session, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py index 0ca62d2f014..5a7855102c0 100644 --- a/tests/integration/streaming/test_interception_converted_stream_options.py +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -5,6 +5,7 @@ import uuid from collections.abc import Callable, Generator, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass +from hashlib import sha256 from pathlib import Path from queue import SimpleQueue from types import MappingProxyType @@ -216,10 +217,13 @@ def scenario(rig: _Rig) -> Iterator[Scenario]: yield opened -def _deploy(rig: _Rig, scenario: Scenario, mode: iv.Mode, *, api_key: str = iv.AZURE_KEY) -> _Deployment: +def _deploy( + rig: _Rig, scenario: Scenario, mode: iv.Mode, *, api_key: str = iv.AZURE_KEY, priced: bool = False +) -> _Deployment: name: Final = iv.deployment(mode) + pricing: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} if priced else {} model: Final = scenario.model( - model=f"azure/{name}", api_base=rig.wire.url, api_key=api_key, api_version=iv.API_VERSION + model=f"azure/{name}", api_base=rig.wire.url, api_key=api_key, api_version=iv.API_VERSION, **pricing ) return _Deployment(model, name) @@ -325,6 +329,46 @@ def _assert_response_landed_once(model: str, identity: str) -> None: assert row["status"] == "success" and rv.same_response(string_value(row["request_id"]), identity), row +def _assert_single_model_spend_for_key(key: str, call_type: str, *, expected_searches: int = 1) -> None: + digest: Final = sha256(key.encode()).hexdigest() + observed: Final = eventually( + lambda: ( + read_rows( + "SELECT status, spend, request_id, call_type, model_group, session_id " + 'FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND status = %s', + (digest, "success"), + ), + read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (digest,)), + ), + lambda values: ( + bool(values[0]) + and ( + len(values[0]) > expected_searches + 1 + or ( + len(values[0]) == expected_searches + 1 + and len(values[1]) == 1 + and values[1][0]["spend"] is not None + and float(str(values[1][0]["spend"])) + == pytest.approx(sum(float(str(row["spend"])) for row in values[0])) + ) + ) + ), + seconds=70, + ) + rows, balance = observed + model_rows: Final = [row for row in rows if row["call_type"] == call_type] + search_rows: Final = [row for row in rows if row["call_type"] == "asearch"] + spend: Final = sum(float(str(row["spend"])) for row in rows) + assert ( + len(rows) == expected_searches + 1 + and len(model_rows) == 1 + and len(search_rows) == expected_searches + and float(str(model_rows[0]["spend"])) > 0 + and all(float(str(row["spend"])) > 0 for row in search_rows) + and float(str(balance[0]["spend"])) == pytest.approx(spend) + ), (rows, balance) + + def _openai(rig: _Rig) -> openai.OpenAI: return openai.OpenAI(base_url=f"{rig.base_url}/v1", api_key=rig.gateway.key, max_retries=0) @@ -402,6 +446,31 @@ def test_raw_chat_stream_without_stream_options_answers_without_a_usage_chunk(ri _assert_search_turn(rig, deployment, identity) +def test_chat_stream_through_web_search_has_single_spend_owner(rig: _Rig, scenario: Scenario) -> None: + deployment: Final = _deploy(rig, scenario, "search", priced=True) + key: Final = scenario.key(models=[deployment.model]) + stream: Final = _post(rig, "/v1/chat/completions", _chat_body(deployment), key=key) + identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=None) + _assert_search_turn(rig, deployment, identity) + _assert_single_model_spend_for_key(key, "acompletion") + + +def test_multi_round_chat_stream_through_web_search_bills_each_search( + rig: _Rig, scenario: Scenario +) -> None: + deployment: Final = _deploy(rig, scenario, "multi-search", priced=True) + key: Final = scenario.key(models=[deployment.model]) + stream: Final = _post(rig, "/v1/chat/completions", _chat_body(deployment), key=key) + identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=None) + received: Final = _upstream(rig, deployment) + _assert_upstream(received, model_calls=3, searches=2) + assert [search["query"] for search in received.searches] == [ + iv.query(deployment.name), + iv.query(deployment.name, turn=2), + ], received.searches + _assert_single_model_spend_for_key(key, "acompletion", expected_searches=2) + + def test_openai_sdk_responses_stream_through_web_search_keeps_stream_options_off_azure( rig: _Rig, scenario: Scenario ) -> None: @@ -454,6 +523,31 @@ def test_raw_responses_stream_without_stream_options_answers(rig: _Rig, scenario _assert_response_landed_once(deployment.model, identity) +def test_responses_stream_through_web_search_has_single_spend_owner(rig: _Rig, scenario: Scenario) -> None: + deployment: Final = _deploy(rig, scenario, "search", priced=True) + key: Final = scenario.key(models=[deployment.model]) + stream: Final = _post(rig, "/v1/responses", _responses_body(deployment), key=key) + _assert_responses_restreamed(stream, iv.answer(deployment.name)) + _assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1) + _assert_single_model_spend_for_key(key, "aresponses") + + +def test_multi_round_responses_stream_through_web_search_bills_each_search( + rig: _Rig, scenario: Scenario +) -> None: + deployment: Final = _deploy(rig, scenario, "multi-search", priced=True) + key: Final = scenario.key(models=[deployment.model]) + stream: Final = _post(rig, "/v1/responses", _responses_body(deployment), key=key) + _assert_responses_restreamed(stream, iv.answer(deployment.name)) + received: Final = _upstream(rig, deployment) + _assert_upstream(received, model_calls=3, searches=2) + assert [search["query"] for search in received.searches] == [ + iv.query(deployment.name), + iv.query(deployment.name, turn=2), + ], received.searches + _assert_single_model_spend_for_key(key, "aresponses", expected_searches=2) + + 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( diff --git a/tests/integration/translation/decisions/bases/parallel_ai.py b/tests/integration/translation/decisions/bases/parallel_ai.py new file mode 100644 index 00000000000..14e9cd6fc43 --- /dev/null +++ b/tests/integration/translation/decisions/bases/parallel_ai.py @@ -0,0 +1,257 @@ +from typing import Final + +from integration.translation.case import TranslationTestCase +from pydantic import JsonValue + +"""Provider request and reply shape from https://docs.parallel.ai/api-reference/decisions/make-web-grounded-typed-decisions. Mock reply captured live on 2026-10-31. +""" +_CITATIONS: Final[list[JsonValue]] = [ + { + "url": "https://access.redhat.com/solutions/7108129", + "title": 'NFS client hangs and "nfs4_handle_delegation_recall_error" found in dmesg. - Red Hat Customer Portal', + "excerpts": [ + 'NFS client hangs and "nfs4_handle_delegation_recall_error" found in dmesg.\n\nSolution Verified - Updated 2026-06-18T01:07:40+00:00 -\nEnglish\n* No translations currently exist.\n\nIssue\n\n* NFS client hangs while returning delegations.\n* Below delegation recall errors may be reported in dmesg:\nNFS: nfs4_handle_delegation_recall_error: unhandled error -512\nNFS: nfs4_handle_delegation_recall_error: unhandled error -10024\n* We have a hanging process, it\'s using data ontop of a mounted NFS export. The fcntl() system call issued by a user process has not returned, The user process processing does not finish - usually it finishes within 2 minutes.\n\nEnvironment\n\n* Red Hat Enterprise Linux (RHEL) 8, all minor releases\n* kernel:\n* rhel8.4: kernels before 4.18.0-305.194.1.el8_4 affected\n* rhel8.10: kernels before 4.18.0-553.37.1.el8_10 affected\n* NFSv4' + ], + }, + { + "url": "https://github.com/moby/buildkit/issues/2804", + "title": "GitHub cache exporter (`type=gha`) hangs reletively low number of cache reads/writes, satirating api limits/triggering throttling. · Issue #2804 · moby/buildkit · GitHub", + "excerpts": [ + 'Copy link\nOpen\nOpen\nGitHub cache exporter ( type=gha ) hangs reletively low number of cache reads/writes, satirating api limits/triggering throttling. #2804\nCopy link\n\nDescription\n\n@bendavies\nbendavies\nopened on Apr 13, 2022\nIssue body actions\n\nBehaviour\n\nCache operations using type=gha will start hanging after a relavatively low number of cache reads/writes.\nWe are a small team (5 devs) and are hitting these limits constantly on our main repository. Not exactly high throughput.\n\n...\n\nCache reads and writes go towards saturating the limit, and both are affected by this throttling.\n@dhadka mentions there is a limit of about "150 cache requests in 5 minutes", and seems about right: #2276 (comment)\n\nExpected behaviour\n\nI don\'t know, but the current behaviour of hanging is not great. there is no feedback as to what is hanging resulting in the user just being confused.' + ], + }, + { + "url": "https://github.com/huginn/huginn/issues/2963", + "title": "WebsiteAgent hangs on a specific url, never finishes job or times out · Issue #2963 · huginn/huginn · GitHub", + "excerpts": [ + 'WebsiteAgent hangs on a specific url, never finishes job or times out\n\n* Page: GitHub issue\n* URL: https://github.com/huginn/huginn/issues/2963\n* State: open\n* Author: BoKKeR\n* Created: 2021-03-14T16:42:01Z\n* Updated: 2021-03-16T19:09:24Z\n* Repository: huginn/huginn\n* Number: #2963\n* Comments: 1\nI know it might be controversial but I am trying to get a new GPU for a while. Wanted to setup a monitor for when the stock on amd.com changes. But the site just won\'t load. The WebsiteAgent gets stuck in\nFetching https://www.amd.com/de/direct-buy/se\nThe only way to stop it is by stopping the container. As I am running the service in a container I cant find any of the debug log locations. Can anyone else try if they are encountering the same error?\n{\n "debug": "true",\n "expected_update_period_in_days": "2",\n "url": "https://www.amd.com/de/direct-buy/se",\n "type": "html",' + ], + }, + { + "url": "https://github.com/mozilla-firefox/firefox-crb/blob/main/docs/bug-mgmt/guides/severity.md", + "title": "firefox-crb/docs/bug-mgmt/guides/severity.md at main · mozilla-firefox/firefox-crb · GitHub", + "excerpts": [ + "We use the severity field in Bugzilla to indicate the scope of a bug's effect on Firefox.\nThe field is display alongside the bug's priority.\n\nValues\n\nSeverity levels and their definitions are enumerated at https://wiki.mozilla.org/BMO/UserGuide/BugFields#bug_severity.\nBy default, new bugs have a severity of --.\n\nExamples of S1 bugs\n\n* WebExtensions disabled for all users\n* Web search not working from URL bar\n* Crashes with data loss\n\nExamples of S2 bugs\n\nBugs that could reasonably be expected to cause a Firefox user to switch browsers, either because the severity is bad enough, or the frequency of occurrence is high enough.\n* Bug 1640441 - Slack hangs indefinitely in a onResize loop" + ], + }, + { + "url": "https://github.com/n0-computer/iroh/issues/4604", + "title": "Endpoint::close hangs forever if the tokio runtime gets dropped while the connection is open · Issue #4604 · n0-computer/iroh · GitHub", + "excerpts": [ + "@alpergundogdu\nalpergundogdu\nopened on Oct 14, 2026\nIssue body actions\nDescribe the bug\nAfter opening a connection from inside a short-lived tokio runtime, and dropping the runtime while the connection is still open, Endpoint::close() never returns.\nIn my repro it hangs reliably; when I keep the runtime open, it closes instantly.\nLocal repro/test (has the control tests too): https://github.com/decdn/iroh-tests/blob/main/tests/bug_endpoint_close_after_runtime_drop.rs\n\n...\n\nIf spawning on the caller's runtime is intended, a note on Endpoint::connect saying the driver lives there would have saved us a lot of time. (Maybe it's an expectation/documentation issue?)\nIroh\nVersion: iroh 1.3.0, noq 1.3.0\nEndpoint configuration: bound to 127.0.0.1, no relays.\nPlatform(s)\nDesktop:\n\n...\n\nSorry @flub , that was unclear: the endpoint's runtime stays alive, and close() is called from it. Only a separate short-lived runtime that runs connect() gets dropped." + ], + }, + { + "url": "https://bugzilla.mozilla.org/show_bug.cgi?id=1640441", + "title": "1640441 - Slack hangs indefinitely in a onResize loop", + "excerpts": [ + "Closed Bug 1640441 Opened 6 years ago Closed 6 years ago\n\nSlack hangs indefinitely in a on Resize loop\n\nCategories\n\n(Web Compatibility :: Site Reports, defect, P1)\n\nProduct:\nWeb Compatibility ▾\nComponent:\nSite Reports ▾\nType:\ndefect\nPriority:\nP1 Severity:\nS2\n\nTracking\n\n(Not tracked)\n\nStatus:\nRESOLVED FIXED\n\nPeople\n\n(Reporter: Harald, Assigned: emilio)\n\nReferences\n\n(Depends on 1 open bug)\n\nDetails\n\n(Keywords: enterprise, perf, webcompat:site-wait)\n\nAttachments\n\n(1 file, 1 obsolete file)\n\n|relevant unminified code.\n6 years ago Emilio Cobos Álvarez [:emilio]\n* 13.50 KB,\n\n...\n\n* Hover over the link\n* ... tab immediately becomes unresponsive ...\nThis may well be a bug in slack, but IMO it's still on us to try and fix what's going on. When it happens, it affects not only Slack but my entire Firefox session. This is a serious usability issue -- I'm now at the point of restarting Firefox several times a day because of this bug, if" + ], + }, + { + "url": "https://www.ibm.com/support/pages/cognos-business-intelligence-1022-fix-lists", + "title": "Cognos Business Intelligence 10.2.2 Fix Lists", + "excerpts": [ + 'IBM Support \n\nCognos Business Intelligence 10.2.2 Fix Lists\n\nRelease Notes\n\nAbstract\n\nA comprehensive list of defect corrections for major releases, refresh packs and fix packs of Cognos Business Intelligence 10.2.2.\nDetails of the APARs listed below can be accessed by clicking the link for the APAR number. If you have questions about a particular defect, please contact Customer Support.\n\n...\n\n* PI05143 | MULTIPLE VALUES IN CONTEXT FILTERS CAUSE A CROSSTAB REPORT TO SHOW DASH "-" INSTEAD OF CORRECT VALUES\n* PI05166 | QE-DEF-0359 The query contains a reference to at least one object that does not exist\n* PI05233 | Export deployments from 8.4.1 with export folder or package tasks in Content Administration contain unknown content steps\n\n...\n\n* PI08523 | Progressive Chart - Incorrect Colours when "Reverse Catagory Order" = True\n* PI08585 | Cognos hangs when starting against a Content Store from previous version (upgrade) and using SSL gateway UR' + ], + }, + { + "url": "https://josm.eu/ticket/8175", + "title": "\n #8175 ([PATCH] Cadastre grab query sometimes hangs completely)\n – JOSM\n\n ", + "excerpts": [ + 'Opened 14 years ago\nClosed 14 years ago\nLast modified 14 years ago\n\n#8175 closed defect ( fixed )\n\n[PATCH] Cadastre grab query sometimes hangs completely\n\n* Reported by: | Don-vip | Owned by: | pieren\n* Priority: | normal | Milestone: | \n* Component: | Plugin cadastre-fr | Version: | \n* Keywords: | | Cc: | \n\nDescription\n\nDuring cadastre grab, the HTTP connection with the WMS server sometimes completely hangs (the progress dialog is frozen and cancel button does not work).\nThe only solution is to restart JOSM to have a clean state and grab the cadastre again. We had a similar problem with history request ( #7640 ) when dealing with the API and solved it in [r5324](https://josm.eu/changeset/5324/josm "fix #7640 - History query will not cancel or timeout\n\n...\n\nby Don-vip , 14 years ago\n\n* Attachment: | 8175.patch \u200b added\n\ncomment:1 by Don-vip , 14 years ago' + ], + }, + { + "url": "https://bugzilla.mozilla.org/show_bug.cgi?id=1749908", + "title": "1749908 - Infinite loop in HTTP3 hangs socket thread - Bugzilla", + "excerpts": [ + "Closed Bug 1749908 Opened 4 years ago Closed 4 years ago\n\nInfinite loop in HTTP3 hangs socket thread\n\nCategories\n\n(Core :: Networking, defect)\n\nProduct:\nCore ▾\nComponent:\nNetworking ▾\nPlatform:\nDesktop\nAll\nType:\ndefect\nPriority:\nNot set Severity:\n--\n\nTracking\n\n( )\n\nStatus:\nVERIFIED DUPLICATE of bug 1749910\n\nPeople\n\n(Reporter: heftig, Unassigned)\n\nDetails\n\nBottom ↓\nTags ▾\nTimeline ▾\n| |Jan Alexander Steffens [:heftig]\n* Reporter | \n* ### Description\n•\n* 4 years ago\n* \nWhile browsing, pages suddenly stopped loading.\n\n...\n\nIt's not related to a specific version, we're getting reports ESR is even affected. Suspicion is around some long existing HTTP3 bug that's being triggered by an external service updating.\n\n...\n\nI agree it is not the root cause but it certainly caused a massive problem for many users of Firefox, the majority of whom are opted-in by default. If they had not been, they would not have been affected." + ], + }, + { + "url": "https://bugzilla.mozilla.org/show_bug.cgi?id=1748414", + "title": "1748414 - Firefox 94 sometimes hangs on TLSv1.3 requests", + "excerpts": [ + "Closed Bug 1748414 Opened 4 years ago Closed 4 years ago\n\nFirefox 94 sometimes hangs on TLSv1 .3 requests\n\nCategories\n\n(Core :: Networking: HTTP, defect)\n\nProduct:\nCore ▾\nComponent:\nNetworking: HTTP ▾\nVersion:\nFirefox 94\nType:\ndefect\nPriority:\nNot set Severity:\n--\n\nTracking\n\n( )\n\nStatus:\nRESOLVED DUPLICATE of bug 1740856\n\nPeople\n\n(Reporter: code, Unassigned)\n\nDetails\n\nBottom ↓\n\n...\n\nSteps to reproduce:\nIn the beginning of November, colleagues started to experience hangs of 1 to 1.5 minutes when working with our own web app. We use latest official releases for Windows, Ubuntu and FreeBSD. The problems exist for all plattforms and very likely started with Firefox 94.\nActual results:\nSometimes while interacting normally with our web app, a request hangs for a long time (up to 1.5 minutes; but usually succeeds). This typically happens on POST requests." + ], + }, +] +DECISION_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "parallel_ai/decision", + "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={"x-api-key": "synthetic-parallel-key", "content-type": "application/json"}, + expected_provider_request={ + "mode": "fast", + "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={ + "mode": "fast", + "answers": { + "defect": {"type": "noul", "noul": 0.6442559258537284}, + "severity": { + "type": "choice", + "choice": "high", + "probabilities": {"low": 0.1616964510265985, "high": 0.8383035489734014}, + "confidence": 0.6766070979468028, + }, + "confidence": { + "type": "score", + "score": 0.25195637565890455, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.7480436243410955, "1": 0.25195637565890455}, + "confidence": 0.4960872486821909, + }, + }, + "usage": {"input_tokens": 3166, "output_tokens": 0}, + "citations": _CITATIONS, + "decision_id": "decisions_9fd99def328ed16256c11df297dea899", + }, + expected_litellm_response={ + "model": "decision", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.6442559258537284}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [ + {"value": "low", "probability": 0.1616964510265985}, + {"value": "high", "probability": 0.8383035489734014}, + ], + "confidence": 0.6766070979468028, + }, + { + "type": "score", + "name": "confidence", + "score": 0.25195637565890455, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.7480436243410955}, + {"value": 1, "label": "sure", "probability": 0.25195637565890455}, + ], + "confidence": 0.4960872486821909, + }, + ], + "usage": { + "input_tokens": 3166, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 3166, + }, + "mode": "fast", + "citations": _CITATIONS, + "decision_id": "decisions_9fd99def328ed16256c11df297dea899", + }, +) +DECISION_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", + litellm_endpoint="/v1/systemone", + litellm_request={ + "model": "parallel_ai/decision", + "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"]}, + }, + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/decisions", + expected_provider_headers={"x-api-key": "synthetic-parallel-key", "content-type": "application/json"}, + expected_provider_request={ + "mode": "fast", + "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={ + "mode": "fast", + "answers": { + "defect": {"type": "noul", "noul": 0.6442559258537284}, + "severity": { + "type": "choice", + "choice": "high", + "probabilities": {"low": 0.1616964510265985, "high": 0.8383035489734014}, + "confidence": 0.6766070979468028, + }, + "confidence": { + "type": "score", + "score": 0.25195637565890455, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.7480436243410955, "1": 0.25195637565890455}, + "confidence": 0.4960872486821909, + }, + }, + "usage": {"input_tokens": 3166, "output_tokens": 0}, + "citations": _CITATIONS, + "decision_id": "decisions_9fd99def328ed16256c11df297dea899", + }, + expected_litellm_response={ + "model": "decision", + "answers": { + "defect": {"type": "noul", "noul": 0.6442559258537284}, + "severity": { + "type": "choice", + "choice": "high", + "probabilities": {"low": 0.1616964510265985, "high": 0.8383035489734014}, + "confidence": 0.6766070979468028, + }, + "confidence": { + "type": "score", + "score": 0.25195637565890455, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.7480436243410955, "1": 0.25195637565890455}, + "confidence": 0.4960872486821909, + }, + }, + "usage": {"input_tokens": 3166, "output_tokens": 0}, + "mode": "fast", + "citations": _CITATIONS, + "decision_id": "decisions_9fd99def328ed16256c11df297dea899", + }, +) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_parallel_ai.py b/tests/integration/translation/decisions/basic/test_decisions_basic_parallel_ai.py new file mode 100644 index 00000000000..47eb057323b --- /dev/null +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_parallel_ai.py @@ -0,0 +1,11 @@ +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.parallel_ai import DECISION_SYSTEMONE_TEST_CASE, DECISION_TEST_CASE +from integration.translation.runner import assert_translation + + +@pytest.mark.parametrize("case", [DECISION_TEST_CASE, DECISION_SYSTEMONE_TEST_CASE], ids=lambda case: case.id) +def test_decisions_basic_parallel_ai(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/rust-python-harness/strategies/e2e_parity/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/__init__.py index 298c75f2af1..a2ea98f19ab 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/__init__.py +++ b/tests/rust-python-harness/strategies/e2e_parity/__init__.py @@ -20,7 +20,7 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( CaseDefinition( "messages", NotImplementedCaseSpec( - reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered." + reason="Parity runs as the pytest suite in sdk/messages/test_native_parity.py, not as a runner case." ), surface="sdk", ), diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/conftest.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/conftest.py new file mode 100644 index 00000000000..4bb20713c59 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/conftest.py @@ -0,0 +1,34 @@ +from collections.abc import AsyncIterator, Generator +from importlib import import_module +from pathlib import Path +from typing import Final + +import pytest +import pytest_asyncio + +from tests.test_litellm_rust.support.isolation import isolated_test_state +from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service + +PARITY_ROOT: Final = Path(__file__).parent + + +@pytest_asyncio.fixture(autouse=True, loop_scope="function") +async def isolate_test_state() -> AsyncIterator[None]: + async with isolated_test_state(): + yield + + +@pytest.fixture +def recording_server() -> Generator[RecordingServer]: + with recording_service() as server: + yield server + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + try: + import_module("litellm.rust_bridge._native") + except ImportError: + skip: Final = pytest.mark.skip(reason="requires a compiled litellm.rust_bridge._native extension") + for item in items: + if PARITY_ROOT in item.path.parents: + item.add_marker(skip) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py new file mode 100644 index 00000000000..41008f0d57f --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py @@ -0,0 +1,1079 @@ +"""Native Anthropic Messages calls must be indistinguishable from the Python implementation. + +Each case sends one `/v1/messages` SDK call through the Python handler and through the native route against the same +scripted upstream, then compares one aspect of what happened: + +- dispatch: the native leg was served natively instead of falling back to Python +- request: every request that reached the upstream (method, path, body, headers) +- outcome: what the caller got back (response body, stream bytes, or the raised exception) +- hidden_params: the `_hidden_params` keys the proxy builds its response headers from +- logged: which terminal callbacks fired and the stable standard logging payload fields +- logged_metadata, logged_model_info, logged_response_headers, logged_api_key, logged_call: single callback kwargs + fields that proxy integrations read + +Both legs run once per case and every aspect of that case reads the same pair + +The native leg prepends an unconstrained Rust rule, so feature detectors that keep a call on Python today do not +apply: a case for a detector-covered feature shows what Rust still lacks before that detector can go. Every known +difference is a strict xfail naming the gap, with an id from the Rust vs Python Messages inventory (A routing, +R request, S response, C callbacks, F covered by a fallback detector today) +""" + +import asyncio +import copy +import json +import os +import re +from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Iterable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import asdict, dataclass, field, replace +from types import MappingProxyType +from typing import Final, cast + +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.messages import dispatch +from litellm.rust_bridge import catalog +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.configuration import Rollout +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + +SONNET: Final = "anthropic/claude-sonnet-4-5" +OPUS_ADAPTIVE: Final = "anthropic/claude-opus-4-7" +OPUS_FAST: Final = "anthropic/claude-opus-5-5" +API_KEY: Final = "sk-ant-parity-test" +USER: Final[Mapping[str, JsonValue]] = MappingProxyType({"role": "user", "content": "What is the capital of France?"}) + +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_LOGGED: Final = TypeAdapter(dict[str, object]) +_TRANSPORT_HEADERS: Final = frozenset( + {"accept", "accept-encoding", "connection", "content-length", "host", "user-agent"} +) +_COMPARED_LOG_FIELDS: Final = ( + "call_type", + "custom_llm_provider", + "model", + "model_group", + "model_id", + "api_base", + "stream", + "cache_hit", + "prompt_tokens", + "completion_tokens", + "total_tokens", + "response_cost", + "status", +) +_TERMINAL_EVENTS: Final = ( + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", +) +_VOLATILE: Final = ( + (re.compile(r"time taken=[0-9.]+"), "time taken="), + (re.compile(r"Timeout passed=[0-9.]+"), "Timeout passed="), + (re.compile(r"127\.0\.0\.1:[0-9]+"), "127.0.0.1:"), +) + +TEXT_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "id": "msg_01ParityText", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "Paris"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 14, "output_tokens": 4}, + } +) + +TOOL_USE_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + **TEXT_RESPONSE, + "id": "msg_01ParityTool", + "content": [ + {"type": "text", "text": "Checking."}, + {"type": "tool_use", "id": "toolu_01Parity", "name": "get_weather", "input": {"city": "Paris"}}, + ], + "stop_reason": "tool_use", + } +) + +THINKING_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + **TEXT_RESPONSE, + "id": "msg_01ParityThinking", + "content": [ + {"type": "thinking", "thinking": "The capital of France is Paris.", "signature": "EqQBCkYIARgCKkB"}, + {"type": "text", "text": "Paris"}, + ], + } +) + +CACHED_USAGE_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + **TEXT_RESPONSE, + "usage": { + "input_tokens": 14, + "output_tokens": 4, + "cache_creation_input_tokens": 2048, + "cache_read_input_tokens": 1024, + "service_tier": "standard", + }, + } +) + +WEATHER_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "name": "get_weather", + "description": "Current weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } +) + +TOOL_TURNS: Final[tuple[JsonValue, ...]] = ( + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01Parity", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01Parity", "content": "18C, sunny"}]}, +) + +THINKING_REPLAY: Final[tuple[JsonValue, ...]] = ( + {"role": "user", "content": "What is the capital of France?"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Paris.", "signature": "EqQBCkYIARgCKkB"}, + {"type": "text", "text": "Paris"}, + ], + }, + {"role": "user", "content": "And Germany?"}, +) + +NO_INPUT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +PNG_PIXEL: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + + +def _events(*events: tuple[str, JsonValue]) -> tuple[tuple[str, object], ...]: + return tuple((name, data) for name, data in events) + + +def _message_start(message_id: str) -> tuple[str, JsonValue]: + return ( + "message_start", + { + "type": "message_start", + "message": { + **TEXT_RESPONSE, + "id": message_id, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 14, "output_tokens": 1}, + }, + }, + ) + + +def _text_block(index: int, *pieces: str) -> tuple[tuple[str, JsonValue], ...]: + return ( + ( + "content_block_start", + {"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}}, + ), + *( + ( + "content_block_delta", + {"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": piece}}, + ) + for piece in pieces + ), + ("content_block_stop", {"type": "content_block_stop", "index": index}), + ) + + +def _message_end(stop_reason: str, output_tokens: int) -> tuple[tuple[str, JsonValue], ...]: + return ( + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": {"output_tokens": output_tokens}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + + +TEXT_STREAM: Final = _events( + _message_start("msg_01ParityStream"), *_text_block(0, "Pa", "ris"), *_message_end("end_turn", 4) +) + +TOOL_STREAM: Final = _events( + _message_start("msg_01ParityToolStream"), + *_text_block(0, "Checking."), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_01Parity", "name": "get_weather", "input": {}}, + }, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": '{"city": '}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": '"Paris"}'}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 1}), + *_message_end("tool_use", 22), +) + +THINKING_STREAM: Final = _events( + _message_start("msg_01ParityThinkingStream"), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + }, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "Paris."}}, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "EqQBCkYIARgCKkB"}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + *_text_block(1, "Paris"), + *_message_end("end_turn", 12), +) + +ERROR_MID_STREAM: Final = _events( + _message_start("msg_01ParityOverloaded"), + *_text_block(0, "Pa"), + ("error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}), +) + +SLOW_STREAM: Final = _events( + _message_start("msg_01ParitySlow"), + *_text_block(0, "a", "b", "c", "d", "e", "f"), + *_message_end("end_turn", 6), +) + + +def _upstream_error(kind: str, message: str) -> Mapping[str, JsonValue]: + return MappingProxyType({"type": "error", "error": {"type": kind, "message": message}}) + + +def _ok(body: Mapping[str, JsonValue] = TEXT_RESPONSE, **headers: str) -> ResponseSpec: + return ResponseSpec(body=dict(body), headers=dict(headers)) + + +def _stream(events: tuple[tuple[str, object], ...], chunk_delay: float = 0) -> ResponseSpec: + return ResponseSpec(body={}, events=events, chunk_delay=chunk_delay) + + +def _failure(status: int, kind: str, message: str, **headers: str) -> ResponseSpec: + return ResponseSpec(body=dict(_upstream_error(kind, message)), status=status, headers=dict(headers)) + + +CALLBACK_ASPECTS: Final = ( + "logged_metadata", + "logged_model_info", + "logged_response_headers", + "logged_api_key", + "logged_call", +) +ASPECTS: Final = ("dispatch", "request", "outcome", "hidden_params", "logged", *CALLBACK_ASPECTS) + + +@dataclass(frozen=True, slots=True) +class Case: + name: str + kwargs: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + upstream: tuple[ResponseSpec, ...] = (ResponseSpec(body=dict(TEXT_RESPONSE)),) + model: str = SONNET + sync: bool = False + env: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + settings: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + private_keys: bool = False + gaps: Mapping[str, str | None] = field(default_factory=lambda: MappingProxyType({})) + + def params(self) -> Iterator[object]: + for aspect in ASPECTS: + gap = {**KNOWN_GAPS.get(self.name, {}), **self.gaps}.get(aspect) + marks = () if not gap else (pytest.mark.xfail(reason=gap, strict=True),) + yield pytest.param(self, aspect, id=f"{self.name}-{aspect}", marks=marks) + + +def _gaps(**gaps: str | None) -> Mapping[str, str | None]: + return MappingProxyType(gaps) + + +def _kw(**kwargs: object) -> Mapping[str, object]: + return MappingProxyType(kwargs) + + +A10_NO_MAX_TOKENS: Final = "A10: native admission hands a call without an int max_tokens to Python, whose signature raises the TypeError; only dispatch differs" +S6_PRIVATE_KEYS: Final = "S6: native writes _hidden_params into the response dict" +S15_COMPLETE_DEADLINE: Final = ( + "S15: by design, native bounds a complete response by one overall deadline where Python bounds each read, " + "so a body trickling in slower than the timeout succeeds in Python and times out natively" +) +F_MOCK: Final = "F: mock_response is Python-only (detector mock_response)" +F_THINKING_RETRY: Final = ( + "F: retry without thinking after a rejected signature is Python-only (detector thinking_signature_retry)" +) +F_CACHE_INJECTION: Final = "F: prompt cache_control injection is Python-only (detector cache_control_injection)" +F_TOOL_CUSTOM: Final = "F: hoisting a tool's custom field is Python-only (detector tool_custom_field)" +F_STREAM_COST: Final = "F: cost injection into streamed usage is Python-only (detector stream_cost_injection)" + +KNOWN_GAPS: Final[Mapping[str, Mapping[str, str | None]]] = MappingProxyType( + { + "response_private_keys": _gaps(outcome=S6_PRIVATE_KEYS), + "missing_max_tokens": _gaps(dispatch=A10_NO_MAX_TOKENS), + "mock_response": _gaps( + logged_response_headers=F_MOCK, + request=F_MOCK, + outcome=F_MOCK, + hidden_params=F_MOCK, + logged=F_MOCK, + logged_metadata=F_MOCK, + logged_api_key=F_MOCK, + logged_call=F_MOCK, + ), + "mock_response_rate_limit": _gaps( + logged_response_headers=F_MOCK, + request=F_MOCK, + outcome=F_MOCK, + logged=F_MOCK, + logged_api_key=F_MOCK, + logged_call=F_MOCK, + ), + "thinking_signature_rejected_then_retried": _gaps( + request=F_THINKING_RETRY, + outcome=F_THINKING_RETRY, + hidden_params=F_THINKING_RETRY, + logged=F_THINKING_RETRY, + logged_metadata=F_THINKING_RETRY, + logged_response_headers=F_THINKING_RETRY, + ), + "cache_control_injection_points": _gaps(request=F_CACHE_INJECTION), + "tool_custom_field": _gaps(request=F_TOOL_CUSTOM), + "stream_cost_injection": _gaps(outcome=F_STREAM_COST), + "enable_prompt_caching": _gaps(request=F_CACHE_INJECTION), + } +) + + +REQUEST_CASES: Final = ( + Case("plain_text"), + Case("system_string", _kw(system="Answer in one word.")), + Case( + "system_blocks_with_cache_control", + _kw(system=[{"type": "text", "text": "Answer in one word.", "cache_control": {"type": "ephemeral"}}]), + ), + Case("multi_turn_tool_result", _kw(messages=list(TOOL_TURNS), tools=[dict(WEATHER_TOOL)])), + Case( + "tools_with_tool_choice", + _kw(tools=[dict(WEATHER_TOOL)], tool_choice={"type": "tool", "name": "get_weather"}), + upstream=(_ok(TOOL_USE_RESPONSE),), + ), + Case("tool_choice_none", _kw(tools=[dict(WEATHER_TOOL)], tool_choice={"type": "none"})), + Case("tool_input_examples", _kw(tools=[{**WEATHER_TOOL, "input_examples": [{"city": "Paris"}]}])), + Case("sampling_params", _kw(temperature=0.3, top_p=0.9, top_k=40)), + Case("stop_sequences", _kw(stop_sequences=["\n\nHuman:"])), + Case("metadata_user_id", _kw(metadata={"user_id": "user-123"})), + Case("metadata_extra_keys_are_filtered", _kw(metadata={"user_id": "user-123", "session": "s-1"})), + Case("metadata_user_id_not_a_string", _kw(metadata={"user_id": 123})), + Case("service_tier", _kw(service_tier="auto")), + Case("thinking_enabled", _kw(thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096)), + Case("thinking_disabled", _kw(thinking={"type": "disabled"})), + Case( + "legacy_thinking_on_adaptive_model_keeps_display", + _kw(thinking={"type": "enabled", "budget_tokens": 2048, "display": "summarized"}, max_tokens=4096), + model=OPUS_ADAPTIVE, + ), + Case( + "legacy_thinking_display_updates_on_adaptive_model", + _kw(thinking={"type": "enabled", "budget_tokens": 2048, "display": "updates"}, max_tokens=4096), + model=OPUS_ADAPTIVE, + ), + Case( + "legacy_thinking_string_budget_on_adaptive_model", + _kw(thinking={"type": "enabled", "budget_tokens": " 4_096 "}, max_tokens=8192), + model=OPUS_ADAPTIVE, + ), + Case( + "legacy_thinking_non_numeric_budget_on_adaptive_model", + _kw(thinking={"type": "enabled", "budget_tokens": "lots"}, max_tokens=8192), + model=OPUS_ADAPTIVE, + ), + Case( + "legacy_thinking_null_budget_on_adaptive_model", + _kw(thinking={"type": "enabled", "budget_tokens": None}, max_tokens=8192), + model=OPUS_ADAPTIVE, + ), + Case( + "legacy_thinking_non_numeric_budget_dropped", + _kw(thinking={"type": "enabled", "budget_tokens": [1]}, max_tokens=8192, drop_params=True), + model=OPUS_ADAPTIVE, + ), + Case( + "legacy_thinking_carries_extra_keys", + _kw( + thinking={"type": "enabled", "budget_tokens": 2048, "block_binding": "strict", "note": "kept"}, + max_tokens=4096, + ), + model=OPUS_ADAPTIVE, + ), + Case( + "legacy_thinking_drops_extra_keys_with_drop_params", + _kw( + thinking={"type": "enabled", "budget_tokens": 2048, "block_binding": "strict", "note": "dropped"}, + max_tokens=4096, + drop_params=True, + ), + model=OPUS_ADAPTIVE, + ), + Case( + "legacy_thinking_float_budget_on_adaptive_model", + _kw(thinking={"type": "enabled", "budget_tokens": 8192.0}, max_tokens=16000), + model=OPUS_ADAPTIVE, + ), + Case( + "adaptive_thinking_with_block_binding", + _kw(thinking={"type": "adaptive", "block_binding": "strict"}, max_tokens=4096), + model=OPUS_ADAPTIVE, + ), + Case("reasoning_effort", _kw(reasoning_effort="high", max_tokens=8192)), + Case("reasoning_effort_on_adaptive_model", _kw(reasoning_effort="low", max_tokens=8192), model=OPUS_ADAPTIVE), + Case("output_config_effort", _kw(output_config={"effort": "medium"}), model=OPUS_ADAPTIVE), + Case( + "output_format_json_schema", + _kw( + output_format={ + "type": "json_schema", + "schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ), + ), + Case("speed_on_model_without_speed", _kw(speed="fast")), + Case("speed_on_model_without_speed_drop_params", _kw(speed="fast", drop_params=True)), + Case("speed_on_fast_model", _kw(speed="fast"), model=OPUS_FAST), + Case("temperature_on_adaptive_model", _kw(temperature=0.5), model=OPUS_ADAPTIVE), + Case("top_k_on_adaptive_model", _kw(top_k=5), model=OPUS_ADAPTIVE), + Case("top_k_on_adaptive_model_drop_params", _kw(top_k=5, drop_params=True), model=OPUS_ADAPTIVE), + Case( + "additional_drop_params_hides_rejected_temperature", + _kw(temperature=0.5, additional_drop_params=["temperature"]), + model=OPUS_ADAPTIVE, + ), + Case( + "additional_drop_params_nested_path", + _kw( + tools=[{**WEATHER_TOOL, "input_examples": [{"city": "Paris"}]}], + additional_drop_params=["tools[*].input_examples", "metadata"], + metadata={"user_id": "user-123"}, + ), + ), + Case( + "context_management_edits", + _kw(context_management={"edits": [{"type": "clear_tool_uses_20250919"}]}), + ), + Case( + "container_skills", _kw(container={"skills": [{"type": "anthropic", "skill_id": "pptx", "version": "latest"}]}) + ), + Case( + "remote_mcp_servers", + _kw(mcp_servers=[{"type": "url", "url": "https://mcp.example.com/sse", "name": "example"}]), + ), + Case("inference_geo", _kw(inference_geo="us")), + Case("top_level_cache_control", _kw(cache_control={"type": "ephemeral"})), + Case("null_safeguards_are_stripped", _kw(safeguards=None)), + Case( + "empty_text_blocks_are_stripped", + _kw( + messages=[ + {"role": "user", "content": [{"type": "text", "text": ""}, {"type": "text", "text": "hi"}]}, + ] + ), + ), + Case( + "tool_use_ids_are_sanitized", + _kw( + messages=[ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "call:abc/1", "name": "get_weather", "input": dict(NO_INPUT)} + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call:abc/1", "content": "sunny"}]}, + ], + tools=[dict(WEATHER_TOOL)], + ), + ), + Case( + "provider_specific_fields_are_stripped", + _kw(messages=[{"role": "user", "content": "hi", "provider_specific_fields": {"x": 1}}]), + ), + Case( + "image_block", + _kw( + messages=[ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_PIXEL}}, + {"type": "text", "text": "Describe it"}, + ], + } + ] + ), + ), + Case( + "document_block", + _kw( + messages=[ + { + "role": "user", + "content": [ + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "Paris."}}, + {"type": "text", "text": "Summarize"}, + ], + } + ] + ), + ), + Case("server_web_search_tool", _kw(tools=[{"type": "web_search_20250305", "name": "web_search", "max_uses": 2}])), + Case("extra_headers", _kw(extra_headers={"x-trace-id": "trace-1", "anthropic-beta": "files-api-2025-04-14"})), + Case( + "anthropic_beta_from_headers_and_extra_headers", + _kw(headers={"anthropic-beta": "client-beta-1"}, extra_headers={"anthropic-beta": "deploy-beta-2"}), + ), + Case("caller_anthropic_version_capitalised", _kw(extra_headers={"Anthropic-Version": "2023-06-01"})), + Case( + "provider_specific_header", + _kw(provider_specific_header={"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "yes"}}), + ), + Case("explicit_custom_llm_provider", _kw(custom_llm_provider="anthropic"), model="claude-sonnet-4-5"), + Case("bare_model_name", model="claude-sonnet-4-5"), + Case("api_base_with_messages_path", _kw(api_base_suffix="/v1/messages")), + Case( + "api_base_with_messages_path_and_query", + _kw(api_base_suffix="/v1/messages?beta=true"), + ), + Case("api_key_from_environment", _kw(api_key=None), env=MappingProxyType({"ANTHROPIC_API_KEY": API_KEY})), + Case( + "auth_token_from_environment", + _kw(api_key=None), + env=MappingProxyType({"ANTHROPIC_AUTH_TOKEN": "oauth-token"}), + ), + Case( + "reasoning_auto_summary", + _kw(thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + env=MappingProxyType({"LITELLM_REASONING_AUTO_SUMMARY": "true"}), + ), + Case("global_drop_params", _kw(top_k=5), model=OPUS_ADAPTIVE, settings=MappingProxyType({"drop_params": True})), + Case("request_timeout_setting", settings=MappingProxyType({"request_timeout": 30})), + Case( + "litellm_metadata_and_metadata", + _kw( + metadata={"user_id": "user-123"}, + litellm_metadata={"user_api_key": "hashed-key", "user_api_key_team_id": "t1"}, + ), + ), + Case("model_info_with_id", _kw(model_info={"id": "deployment-1", "db_model": False})), + Case( + "deployment_effort_levels", + _kw( + output_config={"effort": "max"}, + model_info={"id": "deployment-1", "reasoning_effort_levels": ["low", "high"]}, + ), + model=OPUS_ADAPTIVE, + ), + Case( + "deployment_effort_levels_from_legacy_thinking", + _kw( + thinking={"type": "enabled", "budget_tokens": 30000}, + max_tokens=40000, + model_info={"id": "deployment-1", "reasoning_effort_levels": ["low", "medium"]}, + ), + model=OPUS_ADAPTIVE, + ), + Case( + "deployment_effort_levels_below_request_with_drop_params", + _kw( + output_config={"effort": "high"}, + drop_params=True, + model_info={"id": "deployment-1", "reasoning_effort_levels": ["low"]}, + ), + model=OPUS_ADAPTIVE, + ), +) + +RESPONSE_CASES: Final = ( + Case("response_with_tool_use", _kw(tools=[dict(WEATHER_TOOL)]), upstream=(_ok(TOOL_USE_RESPONSE),)), + Case("response_with_thinking", upstream=(_ok(THINKING_RESPONSE),)), + Case("response_with_cache_usage", upstream=(_ok(CACHED_USAGE_RESPONSE),)), + Case( + "response_with_null_usage_and_container", + upstream=(_ok({**TEXT_RESPONSE, "usage": None, "container": None}),), + ), + Case( + "response_without_stop_sequence", + upstream=(_ok({k: v for k, v in TEXT_RESPONSE.items() if k != "stop_sequence"}),), + ), + Case("response_with_unknown_field", upstream=(_ok({**TEXT_RESPONSE, "future_field": {"a": 1}}),)), + Case("response_key_order", upstream=(_ok({"usage": TEXT_RESPONSE["usage"], **TEXT_RESPONSE}),)), + Case( + "response_upstream_headers", + upstream=( + _ok(TEXT_RESPONSE, **{"request-id": "req_01Parity", "anthropic-ratelimit-requests-remaining": "99"}), + ), + ), + Case("response_private_keys", private_keys=True), +) + +FAILURE_CASES: Final = ( + Case("upstream_400", upstream=(_failure(400, "invalid_request_error", "max_tokens: too large"),)), + Case("upstream_401", upstream=(_failure(401, "authentication_error", "invalid x-api-key"),)), + Case("upstream_403", upstream=(_failure(403, "permission_error", "not allowed"),)), + Case("upstream_404", upstream=(_failure(404, "not_found_error", "model not found"),)), + Case("upstream_413", upstream=(_failure(413, "request_too_large", "too large"),)), + Case( + "upstream_429_with_retry_after", + upstream=( + _failure(429, "rate_limit_error", "slow down", **{"retry-after": "7", "request-id": "req_01Parity"}), + ), + ), + Case("upstream_500", upstream=(_failure(500, "api_error", "internal"),)), + Case("upstream_529_overloaded", upstream=(_failure(529, "overloaded_error", "Overloaded"),)), + Case( + "upstream_400_context_window", + upstream=(_failure(400, "invalid_request_error", "prompt is too long: 250000 tokens > 200000 maximum"),), + ), + Case("upstream_non_json_error", upstream=(ResponseSpec(body=b"bad gateway", status=502),)), + Case("upstream_invalid_json_success", upstream=(ResponseSpec(body=b"not json", status=200),)), + Case("upstream_timeout", _kw(timeout=0.5), upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), delay=1.5),)), + Case( + "response_trickles_slower_than_timeout", + _kw(timeout=1), + upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), chunk_delay=0.5, body_chunks=4),), + gaps=_gaps( + outcome=S15_COMPLETE_DEADLINE, + hidden_params=S15_COMPLETE_DEADLINE, + logged=S15_COMPLETE_DEADLINE, + logged_metadata=S15_COMPLETE_DEADLINE, + logged_response_headers=S15_COMPLETE_DEADLINE, + ), + ), + Case("missing_max_tokens", _kw(max_tokens=None), upstream=()), + Case("missing_api_key", _kw(api_key=None), upstream=()), +) + +STREAM_CASES: Final = ( + Case("stream_text", _kw(stream=True), upstream=(_stream(TEXT_STREAM),)), + Case("stream_tool_use", _kw(stream=True, tools=[dict(WEATHER_TOOL)]), upstream=(_stream(TOOL_STREAM),)), + Case( + "stream_thinking", + _kw(stream=True, thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + upstream=(_stream(THINKING_STREAM),), + ), + Case("stream_error_event_mid_stream", _kw(stream=True), upstream=(_stream(ERROR_MID_STREAM),)), + Case( + "stream_longer_than_timeout", + _kw(stream=True, timeout=1), + upstream=(_stream(SLOW_STREAM, chunk_delay=0.25),), + ), + Case( + "stream_read_stalls_past_timeout", + _kw(stream=True, timeout=0.5), + upstream=(_stream(SLOW_STREAM, chunk_delay=1.0),), + ), + Case( + "stream_upstream_429", + _kw(stream=True), + upstream=(_failure(429, "rate_limit_error", "slow down", **{"retry-after": "3"}),), + ), + Case("stream_false_explicit", _kw(stream=False)), +) + +SYNC_CASES: Final = ( + Case("sync_text", sync=True), + Case("sync_stream", _kw(stream=True), upstream=(_stream(TEXT_STREAM),), sync=True), + Case("sync_upstream_429", upstream=(_failure(429, "rate_limit_error", "slow down"),), sync=True), +) + +FALLBACK_CASES: Final = ( + Case("mock_response", _kw(mock_response="Hi from the mock"), upstream=()), + Case("mock_response_rate_limit", _kw(mock_response="litellm.RateLimitError"), upstream=()), + Case( + "thinking_replay", + _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + ), + Case( + "thinking_signature_rejected_then_retried", + _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + upstream=( + _failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"), + _ok(), + ), + ), + Case( + "cache_control_injection_points", + _kw(system="Answer in one word.", cache_control_injection_points=[{"location": "message", "role": "system"}]), + ), + Case("tool_custom_field", _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": True}}])), + Case( + "strict_tools_rejected_then_retried", + _kw(tools=[{**WEATHER_TOOL, "strict": True}], drop_params=True), + upstream=( + _failure(400, "invalid_request_error", "tools.0.strict: too many strict tools"), + _ok(), + ), + ), + Case("modify_params", _kw(max_tokens=1), settings=MappingProxyType({"modify_params": True})), + Case( + "stream_cost_injection", + _kw(stream=True), + upstream=(_stream(TEXT_STREAM),), + settings=MappingProxyType({"include_cost_in_streaming_usage": True}), + ), + Case("enable_prompt_caching", _kw(system="Answer in one word.", enable_prompt_caching=True)), + Case("request_timeout_kwarg", _kw(request_timeout=30)), + Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)), +) + +CASES: Final = (*REQUEST_CASES, *RESPONSE_CASES, *FAILURE_CASES, *STREAM_CASES, *SYNC_CASES, *FALLBACK_CASES) + + +@dataclass(frozen=True, slots=True) +class Leg: + fell_back: bool + requests: tuple[Mapping[str, object], ...] + outcome: Mapping[str, object] + hidden_params: Mapping[str, object] + logged: Mapping[str, object] + callback_kwargs: Mapping[str, object] + + +@dataclass(slots=True) +class FallbackProbe: + calls: int = 0 + + +_PYTHON_AMESSAGES: Final = cast(Callable[..., Awaitable[object]], dispatch._PYTHON_AMESSAGES) # pyright: ignore[reportPrivateUsage] # the Python baseline under comparison +_PYTHON_MESSAGES: Final = cast(Callable[..., object], dispatch._PYTHON_MESSAGES) # pyright: ignore[reportPrivateUsage] # the Python baseline under comparison +_PUBLIC_AMESSAGES: Final = cast(Callable[..., Awaitable[object]], dispatch.anthropic_messages) +_PUBLIC_MESSAGES: Final = cast(Callable[..., object], dispatch.anthropic_messages_handler) + + +@pytest.fixture +def native_messages() -> Iterator[FallbackProbe]: + probe: Final = FallbackProbe() + + async def python_afallback(*args: object, **kwargs: object) -> object: + probe.calls += 1 + return await _PYTHON_AMESSAGES(*args, **kwargs) + + def python_fallback(*args: object, **kwargs: object) -> object: + probe.calls += 1 + return _PYTHON_MESSAGES(*args, **kwargs) + + with ( + rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT), *catalog.RULES)), + rebound(dispatch, "_PYTHON_AMESSAGES", python_afallback), + rebound(dispatch, "_PYTHON_MESSAGES", python_fallback), + ): + yield probe + + +@contextmanager +def _environment(env: Mapping[str, str]) -> Generator[None]: + # LITELLM_RUST overrides the catalog rule, so a shell that sets it would route the native leg to Python. + names: Final = ( + *env, + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "LITELLM_REASONING_AUTO_SUMMARY", + "LITELLM_RUST", + ) + with pytest.MonkeyPatch.context() as patch: + for name in names: + patch.delenv(name, raising=False) + for name, value in env.items(): + patch.setenv(name, value) + yield + + +@contextmanager +def _settings(settings: Mapping[str, object]) -> Generator[None]: + saved: Final[dict[str, object]] = {name: cast(object, getattr(litellm, name)) for name in settings} + for name, value in settings.items(): + setattr(litellm, name, value) + try: + yield + finally: + for name, value in saved.items(): + setattr(litellm, name, value) + + +def _kwargs(case: Case, server: RecordingServer) -> dict[str, object]: + # Deep: Python writes into the caller's own `litellm_metadata`, which must not leak into the + # other leg. + supplied: Final = copy.deepcopy(dict(case.kwargs)) + suffix: Final = supplied.pop("api_base_suffix", "") + kwargs: Final[dict[str, object]] = { + "model": case.model, + "max_tokens": 64, + "messages": [dict(USER)], + "api_key": API_KEY, + "api_base": f"{server.base_url}{suffix}", + **supplied, + } + return {name: value for name, value in kwargs.items() if value is not None or name in ("safeguards",)} + + +def _mapping(value: object) -> Mapping[str, object]: + return _LOGGED.validate_python(value) if isinstance(value, dict) else {} + + +def _normalized(text: str) -> str: + for pattern, replacement in _VOLATILE: + text = pattern.sub(replacement, text) # rebind-ok: successive substitutions + return text + + +def _error(error: BaseException, phase: str, chunks: Iterable[object] = ()) -> Mapping[str, object]: + header_items: Final = { + name.lower(): value for name, value in _mapping(getattr(error, "litellm_response_headers", None)).items() + } + return { + "kind": "error", + "phase": phase, + "class": type(error).__name__, + "status_code": getattr(error, "status_code", None), + "llm_provider": getattr(error, "llm_provider", None), + "model": getattr(error, "model", None), + "message": _normalized(str(error)), + "headers": {name: header_items.get(name) for name in ("retry-after", "request-id") if name in header_items}, + "partial": _stream_text(chunks), + } + + +def _stream_text(chunks: Iterable[object]) -> str: + return b"".join(chunk if isinstance(chunk, bytes) else str(chunk).encode() for chunk in chunks).decode() + + +def _response(result: object, private_keys: bool) -> Mapping[str, object]: + assert isinstance(result, dict), f"expected a dict response, got {type(result).__name__}" + body: Final = { + name: value + for name, value in cast(dict[str, object], result).items() + if private_keys or not name.startswith("_") + } + return {"kind": "response", "body": _OBJECT.validate_python(body)} + + +def _hidden_params(result: object) -> Mapping[str, object]: + hidden: Final = dict(_mapping(get_hidden_params_dict(result))) + headers: Final = _mapping(hidden.pop("additional_headers", None)) + header_names: Final = sorted(name for name in headers if name not in ("x-litellm-rust", "x-litellm-cache-key")) + return {"keys": sorted(hidden), "additional_headers": header_names} + + +async def _consume_async(stream: AsyncIterator[object]) -> tuple[list[object], BaseException | None]: + chunks: Final[list[object]] = [] + try: + async for chunk in stream: + chunks.append(chunk) + except Exception as error: # noqa: BLE001 # the raised error is part of the compared outcome + return chunks, error + return chunks, None + + +def _consume_sync(stream: Iterator[object]) -> tuple[list[object], BaseException | None]: + chunks: Final[list[object]] = [] + try: + for chunk in stream: + chunks.append(chunk) + except Exception as error: # noqa: BLE001 # the raised error is part of the compared outcome + return chunks, error + return chunks, None + + +async def _invoke( + case: Case, kwargs: Mapping[str, object], native: bool +) -> tuple[Mapping[str, object], Mapping[str, object]]: + if case.sync: + call = _PUBLIC_MESSAGES if native else _PYTHON_MESSAGES + try: + result = await asyncio.to_thread(lambda: call(**kwargs)) + except Exception as error: # noqa: BLE001 # the raised error is part of the compared outcome + return _error(error, "call"), {} + if isinstance(result, dict): + body = cast(Mapping[str, object], result) + return _response(body, case.private_keys), _hidden_params(body) + chunks, failure = await asyncio.to_thread(_consume_sync, cast(Iterator[object], result)) + else: + acall = _PUBLIC_AMESSAGES if native else _PYTHON_AMESSAGES + try: + result = await acall(**kwargs) + except Exception as error: # noqa: BLE001 # the raised error is part of the compared outcome + return _error(error, "call"), {} + if isinstance(result, dict): + body = cast(Mapping[str, object], result) + return _response(body, case.private_keys), _hidden_params(body) + chunks, failure = await _consume_async(cast(AsyncIterator[object], result)) + hidden: Final = _hidden_params(result) + if failure is not None: + return _error(failure, "stream", chunks), hidden + return { + "kind": "stream", + "chunk_types": sorted({type(chunk).__name__ for chunk in chunks}), + "text": _stream_text(chunks), + }, hidden + + +async def _terminal_events(recorder: RecordingLogger) -> tuple[str, ...]: + for _ in range(100): + await drain_logging() + if any(name in _TERMINAL_EVENTS for name in recorder.names): + break + await asyncio.sleep(0.05) + await asyncio.sleep(0.2) + await drain_logging() + return tuple(name for name in recorder.names if name in _TERMINAL_EVENTS) + + +def _logged(recorder: RecordingLogger, events: tuple[str, ...]) -> tuple[Mapping[str, object], Mapping[str, object]]: + terminal: Final = [event for event in recorder.events if event.name in _TERMINAL_EVENTS] + if not terminal: + return {"events": events}, {} + details: Final = _mapping(terminal[-1].kwargs) + standard: Final = _mapping(details.get("standard_logging_object")) + error_class: Final = _mapping(standard.get("error_information")).get("error_class") + litellm_params: Final = _mapping(details.get("litellm_params")) + metadata: Final = litellm_params.get("metadata") + metadata_keys: Final = _mapping(metadata).keys() + logged: Final[Mapping[str, object]] = { + "events": events, + **{name: standard.get(name) for name in _COMPARED_LOG_FIELDS}, + "error_class": error_class, + } + callback_kwargs: Final[Mapping[str, object]] = { + "logged_metadata": sorted(metadata_keys) if isinstance(metadata, dict) else metadata, + "logged_model_info": litellm_params.get("model_info"), + "logged_response_headers": bool(details.get("response_headers")), + "logged_api_key": details.get("api_key"), + "logged_call": { + "custom_llm_provider": details.get("custom_llm_provider"), + "model": details.get("model"), + "stream": details.get("stream"), + }, + } + return logged, callback_kwargs + + +async def _leg(case: Case, server: RecordingServer, probe: FallbackProbe | None) -> Leg: + with isolated_callback_registries(), _environment(case.env), _settings(case.settings): + recorder: Final = RecordingLogger() + litellm.callbacks.append(recorder) + for spec in case.upstream: + server.enqueue(replace(spec)) + request_count: Final = len(server.requests) + fallbacks_before: Final = 0 if probe is None else probe.calls + outcome, hidden = await _invoke(case, _kwargs(case, server), native=probe is not None) + events: Final = await _terminal_events(recorder) + server.responses.clear() + logged, callback_kwargs = _logged(recorder, events) + return Leg( + fell_back=probe is not None and probe.calls > fallbacks_before, + requests=tuple( + { + "method": sent.method, + "path": sent.path, + "body": sent.body if sent.body is not None else sent.raw_body.decode(errors="replace"), + "headers": { + name: value for name, value in sorted(sent.headers.items()) if name not in _TRANSPORT_HEADERS + }, + } + for sent in server.requests[request_count:] + ), + outcome=outcome, + hidden_params=hidden, + logged=logged, + callback_kwargs=callback_kwargs, + ) + + +@pytest.fixture +def parity_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.expected_requests = None + return recording_server + + +_LEGS: Final[dict[str, tuple[Leg, Leg]]] = {} + + +async def _legs(case: Case, server: RecordingServer, probe: FallbackProbe) -> tuple[Leg, Leg]: + if case.name not in _LEGS: + _LEGS[case.name] = (await _leg(case, server, probe=None), await _leg(case, server, probe=probe)) + _dump(case.name, *_LEGS[case.name]) + return _LEGS[case.name] + + +def _dump(name: str, python: Leg, native: Leg) -> None: + """Writes both legs to ``$MESSAGES_PARITY_DUMP/.json`` when that variable is set.""" + dump: Final = os.environ.get("MESSAGES_PARITY_DUMP") + if not dump: + return + os.makedirs(dump, exist_ok=True) + with open(os.path.join(dump, f"{name}.json"), "w") as out: + json.dump({"python": asdict(python), "native": asdict(native)}, out, indent=1, default=repr) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("case", "aspect"), tuple(param for case in CASES for param in case.params())) +async def test_native_messages_match_python( + case: Case, aspect: str, parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python, native = await _legs(case, parity_server, native_messages) + match aspect: + case "dispatch": + assert not native.fell_back + case "request": + assert native.requests == python.requests + case "outcome": + assert native.outcome == python.outcome + case "hidden_params": + assert native.hidden_params == python.hidden_params + case "logged": + assert native.logged == python.logged + case _: + assert native.callback_kwargs.get(aspect) == python.callback_kwargs.get(aspect) diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index 722cdbcbaca..beb818c7162 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -18,10 +18,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( ) from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache -from litellm.rust_bridge import runtime -from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule -from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.public_call import NativeCall +from litellm.rust_bridge.catalog import Route from litellm.types.caching import CachingSupportedCallTypes from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE @@ -430,49 +427,6 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task( assert len(recording_server.requests) == 1 -def test_sync_rust_messages_calls_python_cache(recording_server: RecordingServer) -> None: - from litellm.rust_bridge.messages.entrypoints import NATIVE_MESSAGES - - litellm.cache = Cache() - recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) - arguments: Final = { - "model": MESSAGES_MODEL, - "messages": list(MESSAGES), - "max_tokens": 32, - "api_key": "test-key", - "api_base": recording_server.base_url, - } - request: Final = NativeCall( - args=(), - kwargs=arguments, - base={ - "model": MESSAGES_MODEL, - "messages": list(MESSAGES), - "max_tokens": 32, - "stream": None, - "api_key": "test-key", - "api_base": recording_server.base_url, - "custom_llm_provider": "anthropic", - **arguments, - }, - ) - - def call() -> object: - return runtime.run( - RouteContext(Route.MESSAGES), - binding=NATIVE_MESSAGES, - native=lambda hook: hook(request), - python=runtime.NO_PYTHON, - rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),), - ) - - first: Final = call() - second: Final = call() - assert payload(first) == payload(second) - assert cache_key(second) - assert len(recording_server.requests) == 1 - - @pytest.mark.asyncio @pytest.mark.parametrize("namespace_source", ("cache", "metadata")) async def test_rust_messages_legacy_cache_honors_request_namespaces( diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index cd5478dfc3a..a3c03898a12 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -14,7 +14,6 @@ from litellm.types.integrations.argilla import ArgillaItem from litellm.types.utils import ModelResponse, StandardLoggingPayload from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, call from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body -from tests.test_litellm_rust.support.gaps import gap from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed from tests.test_litellm_rust.support.results import response_id, with_id @@ -22,7 +21,6 @@ from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension -SYNC_HANDLERS_ON_ASYNC_CALLS: Final = gap("native async calls also run the sync success and failure handlers") _OBJECT: Final = TypeAdapter(dict[str, object]) _HEADERS: Final = TypeAdapter(dict[str, str]) @@ -211,7 +209,7 @@ async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(se failures: Final = await recorder.wait_for_async("async_log_failure_event") assert raised.value is rejection - assert server.requests == [] + assert not server.requests assert len(failures) == 1 assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names) @@ -341,8 +339,8 @@ async def test_a_failing_logger_does_not_fail_the_call(server: RecordingServer) pytest.param("input", False, False, 1, id="global-input"), pytest.param("success", True, False, 1, id="global-async-success"), pytest.param("success", True, True, 1, id="request-async-success"), - pytest.param("success", False, False, 0, id="global-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), - pytest.param("success", False, True, 0, id="request-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("success", False, False, 0, id="global-sync-success"), + pytest.param("success", False, True, 0, id="request-sync-success"), pytest.param("failure", True, False, 1, id="global-async-failure"), pytest.param("failure", False, False, 1, id="global-sync-failure"), pytest.param("failure", False, True, 1, id="request-sync-failure"), @@ -384,12 +382,11 @@ async def test_pre_call_rule_rejects_before_provider_io(server: RecordingServer) checked.append(text) return False - with rebound(litellm, "pre_call_rules", [reject]): - with pytest.raises(litellm.APIResponseValidationError): - await call(server) + with rebound(litellm, "pre_call_rules", [reject]), pytest.raises(litellm.APIResponseValidationError): + await call(server) assert checked == ["Hello"] - assert server.requests == [] + assert not server.requests @pytest.mark.asyncio @@ -419,11 +416,11 @@ async def test_dataset_hook_returning_none_drops_the_record_from_the_success_sin @pytest.mark.asyncio -async def test_sync_call_returns_the_provider_message(server: RecordingServer) -> None: +async def test_sync_call_raises_like_python_without_logging(server: RecordingServer) -> None: recorder: Final = RecordingLogger() - response: Final = await call(server, asynchronous=False, callbacks=[recorder]) + with pytest.raises(ValueError, match="anthropic_messages_handler is not implemented for sync calls"): + await call(server, asynchronous=False, callbacks=[recorder]) - assert_served_natively(server, response) - assert response_id(response) == MESSAGES_RESPONSE["id"] - assert len(await recorder.wait_for_async("log_success_event")) == 1 + assert not server.requests + assert not recorder.names diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 5b11e628b8d..6d65f289524 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -569,6 +569,12 @@ async def test_strict_tools_stay_native_unless_python_could_drop_strict( ASYNC_ONLY_FEATURES: Final = ( pytest.param(Feature.PRE_REQUEST_HOOK, {}, {"callbacks": [_Rewrite()]}, id="pre-request-hook"), + pytest.param( + Feature.FOUNDRY_BETA_ROW, + {"api_base": "https://team.services.ai.azure.com/anthropic"}, + {}, + id="foundry-beta-row", + ), pytest.param( Feature.WEBSEARCH_INTERCEPTION, {}, {"callbacks": [WebSearchInterceptionLogger()]}, id="websearch-interception" ), diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index c700c093d1a..e57631f48b9 100644 --- a/tests/test_litellm_rust/messages/test_streaming.py +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -6,7 +6,6 @@ import pytest from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_RESPONSE, call from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed -from tests.test_litellm_rust.support.results import consume from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -46,16 +45,14 @@ async def test_stream_relays_provider_events_and_logs_success_once_after_the_las @pytest.mark.asyncio -async def test_sync_stream_relays_provider_events_and_logs_success_once(server: RecordingServer) -> None: - server.enqueue(STREAM) +async def test_sync_stream_raises_like_python_before_sending(server: RecordingServer) -> None: recorder: Final = RecordingLogger() - stream: Final = await call(server, asynchronous=False, stream=True, callbacks=[recorder]) - chunks: Final = await consume(stream) + with pytest.raises(ValueError, match="anthropic_messages_handler is not implemented for sync calls"): + await call(server, asynchronous=False, stream=True, callbacks=[recorder]) - assert_served_natively(server) - assert b"".join(chunks) == b"".join(STREAM.payloads()) - assert len(await recorder.wait_for_async("async_log_success_event")) == 1 + assert not server.requests + assert not recorder.names @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index d7b82767b97..5743658abdf 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -25,13 +25,16 @@ class ResponseSpec: status: int = 200 headers: dict[str, str] = field(default_factory=dict) delay: float = 0 + chunk_delay: float = 0 events: tuple[tuple[str, object], ...] = () + # Splits a non-stream body into this many writes, `chunk_delay` apart, so it trickles in. + body_chunks: int = 1 def payloads(self) -> tuple[bytes, ...]: - if isinstance(self.body, bytes): - return (self.body,) if not self.events: - return (json.dumps(self.body).encode(),) + body: Final = self.body if isinstance(self.body, bytes) else json.dumps(self.body).encode() + size: Final = max(-(-len(body) // max(self.body_chunks, 1)), 1) + return tuple(body[start : start + size] for start in range(0, len(body), size)) or (body,) return tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in self.events) @@ -89,7 +92,9 @@ def recording_service() -> Iterator[RecordingServer]: self.send_header(name, value) self.end_headers() try: - for payload in payloads: + for index, payload in enumerate(payloads): + if index and response.chunk_delay: + time.sleep(response.chunk_delay) self.wfile.write(payload) self.wfile.flush() except (BrokenPipeError, ConnectionResetError): diff --git a/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py index 0d8e7674e7d..210ab6717aa 100644 --- a/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py +++ b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py @@ -63,10 +63,16 @@ class TestCreateErrorResponse: ) assert response["request_id"] == "req_123" - def test_unknown_status_defaults_to_api_error(self): - """Test unknown status code defaults to api_error.""" - response = AnthropicExceptionMapping.create_error_response(418, "I'm a teapot") + @pytest.mark.parametrize("status_code", [402, 405, 409, 418, 422, 504]) + def test_a_status_outside_the_table_stays_an_api_error_for_the_sdk_and_stream_frames(self, status_code: int): + response: Final = AnthropicExceptionMapping.create_error_response(status_code, "scripted status") + frame: Final = anthropic_error_sse_frame(status_code=status_code, raw_message="scripted status") assert response["error"]["type"] == "api_error" + assert frame.error_response["error"]["type"] == "api_error" + + def test_an_explicit_error_type_wins_over_the_status_table(self): + response: Final = AnthropicExceptionMapping.create_error_response(402, "Budget exceeded", error_type="billing_error") + assert response["error"]["type"] == "billing_error" class TestExtractErrorMessage: diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 0f6ef2304fb..0911d5ef888 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -472,6 +472,30 @@ def test_cache_key_tells_apart_requests_that_differ_only_in_continuation_token() assert cache.get_cache_key(**request, continuation_token="") == first_segment_key +def test_cache_key_tells_apart_requests_that_ask_for_token_ids() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]} + + plain_key: Final = cache.get_cache_key(**request) + + assert cache.get_cache_key(**request, return_token_ids=True) != plain_key + assert cache.get_cache_key(**request, return_token_ids=False) == plain_key + assert cache.get_cache_key(**request, return_token_ids=None) == plain_key + + +def test_cache_key_tells_apart_requests_that_ask_for_token_ids_through_extra_body() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]} + + plain_key: Final = cache.get_cache_key(**request) + other_extra_body_key: Final = cache.get_cache_key(**request, extra_body={"top_k": 5}) + + assert cache.get_cache_key(**request, extra_body={"return_token_ids": True}) != plain_key + assert cache.get_cache_key(**request, extra_body={"top_k": 5, "return_token_ids": True}) != other_extra_body_key + assert cache.get_cache_key(**request, extra_body={"return_token_ids": False}) == plain_key + assert cache.get_cache_key(**request, extra_body={"continuation_token": "abc"}) != plain_key + + def test_cache_key_tells_apart_requests_that_differ_only_in_tool_call_id_shape() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) request: Final = { diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 4e6ed16e360..2d02468fd8c 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1,9 +1,13 @@ +import asyncio import copy import datetime import json import os +import queue import unittest -from collections.abc import Mapping, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple, cast, get_args from unittest.mock import ANY, MagicMock, Mock, patch @@ -24,15 +28,15 @@ import litellm from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, ) -from litellm.types.llms.openai import AllMessageValues, REASONING_EFFORT +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.llms.openai import AllMessageValues, REASONING_EFFORT, ResponsesAPIResponse +from litellm.types.utils import ModelResponse, StandardLoggingPayload, TextCompletionResponse if TYPE_CHECKING: from openai.types.responses import ResponseOutputItem from openai.types.responses.response_reasoning_item import ResponseReasoningItem - from litellm.types.llms.openai import ResponsesAPIResponse - from litellm.types.utils import ModelResponse - def test_completion_over_responses_maps_anthropic_search_tool_to_openai_tool_search() -> None: response_body: Final = { @@ -191,6 +195,417 @@ def test_auto_summary_sends_the_level_the_target_accepts( assert request_body["reasoning"] == {"effort": "low", "summary": expected_summary} +@dataclass(frozen=True, slots=True) +class _LoggedSuccess: + call_type: str + payload: StandardLoggingPayload + response_obj: object + complete_streaming_response: object + + +class _LoggedCallTypes(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen_by_caller: Final = MappingProxyType( + { + "completion": queue.SimpleQueue[_LoggedSuccess](), + "acompletion": queue.SimpleQueue[_LoggedSuccess](), + } + ) + + @staticmethod + def _logged(kwargs: Mapping[str, object], response_obj: object, streaming_key: str) -> _LoggedSuccess: + return _LoggedSuccess( + call_type=str(kwargs["call_type"]), + payload=cast(StandardLoggingPayload, kwargs["standard_logging_object"]), + response_obj=response_obj, + complete_streaming_response=kwargs.get(streaming_key), + ) + + def log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + self.seen_by_caller["completion"].put(self._logged(kwargs, response_obj, "complete_streaming_response")) + + async def async_log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + self.seen_by_caller["acompletion"].put( + self._logged(kwargs, response_obj, "async_complete_streaming_response") + ) + + +_GET_WEATHER_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + +_BRIDGED_PROVIDERS: Final = pytest.mark.parametrize( + ("model", "url", "provider_kwargs"), + [ + ("openai/gpt-5.3-codex", "https://api.openai.com/v1/responses", {}), + ( + "azure/gpt-6-sol", + "https://example.invalid/openai/responses", + {"api_base": "https://example.invalid", "api_version": "2025-04-01-preview"}, + ), + ("bedrock_mantle/openai.gpt-5.6-terra", _MANTLE_RESPONSES_URL, {"aws_region_name": "us-east-1"}), + ], + ids=["openai-responses-only", "azure-tools", "bedrock-mantle-tools"], +) + + +_PONG_MESSAGE: Final[Mapping[str, object]] = { + "type": "message", + "id": "msg_logged_call_type", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], +} + + +def _searched_responses_body(model: str) -> Mapping[str, object]: + return { + "id": "resp_logged_call_type", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": model.rsplit("/", 1)[-1], + "output": [ + { + "type": "web_search_call", + "id": "ws_logged_call_type", + "status": "completed", + "action": {"type": "search", "query": "weather"}, + }, + _PONG_MESSAGE, + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + + +def _responses_sse(response_body: Mapping[str, object]) -> str: + events: Final = ( + {"type": "response.created", "response": {**response_body, "status": "in_progress", "output": [], "usage": None}}, + { + "type": "response.output_text.delta", + "item_id": "msg_logged_call_type", + "output_index": 1, + "content_index": 0, + "delta": "pong", + }, + {"type": "response.completed", "response": response_body}, + ) + return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + + +def _bridged_request(model: str, provider_kwargs: Mapping[str, object], stream: bool) -> Mapping[str, object]: + return { + "model": model, + "messages": [{"role": "user", "content": "ping"}], + "tools": [_GET_WEATHER_TOOL], + "api_key": "test-key", + "cache": {"no-cache": True}, + "stream": stream, + **provider_kwargs, + } + + +async def _served_chunk_ids(caller: str, request: Mapping[str, object]) -> frozenset[str]: + if caller == "acompletion": + async_stream: Final = await litellm.acompletion(**request) + assert isinstance(async_stream, CustomStreamWrapper) + return frozenset([chunk.id async for chunk in async_stream]) + sync_stream: Final = litellm.completion(**request) + assert isinstance(sync_stream, CustomStreamWrapper) + return frozenset(chunk.id for chunk in sync_stream) + + +@pytest.mark.parametrize("caller", ["completion", "acompletion"]) +@_BRIDGED_PROVIDERS +@pytest.mark.asyncio +async def test_bridged_chat_call_is_logged_with_the_callers_call_type( + monkeypatch: pytest.MonkeyPatch, caller: str, model: str, url: str, provider_kwargs: Mapping[str, object] +) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _LoggedCallTypes() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + response_body: Final = _searched_responses_body(model) + request: Final = _bridged_request(model, provider_kwargs, stream=False) + + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith=url).mock(return_value=httpx.Response(200, json=response_body)) + if caller == "acompletion": + await litellm.acompletion(**request) + else: + litellm.completion(**request) + logged: Final = await asyncio.to_thread(recorder.seen_by_caller[caller].get, timeout=10) + + assert (logged.call_type, logged.payload["call_type"]) == (caller, caller) + assert isinstance(logged.response_obj, ModelResponse) + assert logged.response_obj.id == response_body["id"] + assert logged.response_obj.choices[0].message.content == "pong" + assert logged.payload["response"]["choices"][0]["message"]["content"] == "pong" + assert logged.payload["response_cost"] == litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate(response_body), model=model + ) + + +@pytest.mark.parametrize("caller", ["completion", "acompletion"]) +@_BRIDGED_PROVIDERS +@pytest.mark.asyncio +async def test_bridged_chat_stream_is_logged_as_the_chat_completion_it_streamed( + monkeypatch: pytest.MonkeyPatch, caller: str, model: str, url: str, provider_kwargs: Mapping[str, object] +) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _LoggedCallTypes() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + response_body: Final = _searched_responses_body(model) + request: Final = _bridged_request(model, provider_kwargs, stream=True) + + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith=url).mock( + return_value=httpx.Response( + 200, text=_responses_sse(response_body), headers={"content-type": "text/event-stream"} + ) + ) + served_ids: Final = await _served_chunk_ids(caller, request) + logged: Final = await asyncio.to_thread(recorder.seen_by_caller[caller].get, timeout=10) + + assert isinstance(logged.complete_streaming_response, ModelResponse) + assert logged.response_obj is logged.complete_streaming_response + assert served_ids == {logged.complete_streaming_response.id} + assert logged.complete_streaming_response.choices[0].message.content == "pong" + assert logged.payload["response"]["choices"][0]["message"]["content"] == "pong" + assert logged.payload["response_cost"] == litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate(response_body), model=model + ) + + +def _json_message(item_id: str, document: Mapping[str, object]) -> Mapping[str, object]: + return { + "type": "message", + "id": item_id, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": json.dumps(document), "annotations": []}], + } + + +@pytest.mark.parametrize("caller", ["completion", "acompletion"]) +@pytest.mark.asyncio +async def test_bridged_structured_output_call_is_logged_with_the_json_document_the_client_got( + monkeypatch: pytest.MonkeyPatch, caller: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _LoggedCallTypes() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + response_body: Final = { + **_searched_responses_body("openai/gpt-5.3-codex"), + "output": [_json_message("msg_first", {"answer": "pong"}), _json_message("msg_second", {"answer": "again"})], + } + request: Final = { + **_bridged_request("openai/gpt-5.3-codex", {}, stream=False), + "response_format": {"type": "json_object"}, + } + + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith="https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=response_body) + ) + served: Final = await litellm.acompletion(**request) if caller == "acompletion" else litellm.completion(**request) + logged: Final = await asyncio.to_thread(recorder.seen_by_caller[caller].get, timeout=10) + + assert isinstance(served, ModelResponse) + assert isinstance(logged.response_obj, ModelResponse) + served_message: Final = served.choices[0].message + logged_message: Final = logged.response_obj.choices[0].message + assert json.loads(served_message.content) == {"answer": "pong"} + assert logged_message.content == served_message.content + assert logged_message.provider_specific_fields == served_message.provider_specific_fields + assert logged.payload["response"]["choices"][0]["message"]["content"] == served_message.content + + +async def _logged_text_completion( + monkeypatch: pytest.MonkeyPatch, caller: str +) -> tuple[TextCompletionResponse, _LoggedSuccess, Mapping[str, object]]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _LoggedCallTypes() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + response_body: Final = _searched_responses_body("openai/gpt-5.3-codex") + request: Final = {"model": "openai/gpt-5.3-codex", "prompt": "ping", "api_key": "test-key"} + + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith="https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=response_body) + ) + served: Final = ( + await litellm.atext_completion(**request) if caller == "acompletion" else litellm.text_completion(**request) + ) + logged: Final = await asyncio.to_thread(recorder.seen_by_caller[caller].get, timeout=10) + assert isinstance(served, TextCompletionResponse) + return served, logged, response_body + + +@pytest.mark.parametrize(("caller", "call_type"), [("completion", "text_completion"), ("acompletion", "atext_completion")]) +@pytest.mark.asyncio +async def test_bridged_text_completion_is_logged_and_priced_under_the_callers_call_type( + monkeypatch: pytest.MonkeyPatch, caller: str, call_type: str +) -> None: + served, logged, response_body = await _logged_text_completion(monkeypatch, caller) + + assert (logged.call_type, logged.payload["call_type"]) == (call_type, call_type) + assert getattr(logged.response_obj, "id") == served.id == response_body["id"] + assert logged.payload["response_cost"] == litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate(response_body), model="openai/gpt-5.3-codex" + ) + + +@pytest.mark.parametrize("caller", ["completion", "acompletion"]) +@pytest.mark.asyncio +async def test_bridged_text_completion_is_logged_as_the_text_completion_it_served( + monkeypatch: pytest.MonkeyPatch, caller: str +) -> None: + served, logged, _ = await _logged_text_completion(monkeypatch, caller) + + assert isinstance(logged.response_obj, TextCompletionResponse) + assert logged.response_obj.choices[0].text == served.choices[0].text == "pong" + assert logged.payload["response"]["object"] == served.object == "text_completion" + assert logged.payload["response"]["choices"][0]["text"] == "pong" + + +@pytest.mark.asyncio +async def test_a_bridged_chat_response_waits_on_the_usage_ack_like_any_chat_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response + from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings + from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + database: Final = CommitBarrier() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.write, UsageCaptureSettings(batch_size=2)) + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) + logging_obj: Final = Logging( + model="openai/gpt-5.3-codex", + messages=[], + stream=False, + call_type="acompletion", + start_time=datetime.datetime.now(), + litellm_call_id="bridged-chat", + function_id="bridged-chat", + ) + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith="https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_searched_responses_body("openai/gpt-5.3-codex")) + ) + served: Final = await litellm.acompletion( + model="openai/gpt-5.3-codex", + messages=[{"role": "user", "content": "ping"}], + api_key="test-key", + litellm_logging_obj=logging_obj, + ) + response: Final = asyncio.create_task(capture_proxy_response(logging_obj, served)) + try: + await database.entered.wait() + assert not response.done() + assert database.committed == () + database.release.set() + await response + assert "bridged-chat:inbound" in database.committed + finally: + database.release.set() + await runtime.stop(1) + await response + + +def _served_message_start_id(events: Sequence[bytes]) -> object: + lines: Final = b"".join(events).decode().splitlines() + starts: Final = tuple( + json.loads(line.removeprefix("data:")) for line in lines if line.startswith("data:") and '"message_start"' in line + ) + assert len(starts) == 1 + return starts[0]["message"]["id"] + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"]) +@pytest.mark.asyncio +async def test_bridged_messages_call_is_logged_as_anthropic_messages_under_the_id_the_client_got( + monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "use_chat_completions_url_for_anthropic_messages", True) + recorder: Final = _LoggedCallTypes() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + response_body: Final = {**_searched_responses_body("openai/gpt-5.3-codex"), "output": [_PONG_MESSAGE]} + upstream: Final = ( + httpx.Response(200, text=_responses_sse(response_body), headers={"content-type": "text/event-stream"}) + if stream + else httpx.Response(200, json=response_body) + ) + + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith="https://api.openai.com/v1/responses").mock(return_value=upstream) + served: Final = await litellm.anthropic.messages.acreate( + model="openai/gpt-5.3-codex", + max_tokens=64, + messages=[{"role": "user", "content": "ping"}], + api_key="test-key", + stream=stream, + ) + served_id: Final = ( + _served_message_start_id([event async for event in cast(AsyncIterator[bytes], served)]) + if stream + else cast(Mapping[str, object], served)["id"] + ) + logged: Final = await asyncio.to_thread(recorder.seen_by_caller["acompletion"].get, timeout=10) + + assert (logged.call_type, logged.payload["call_type"]) == ("anthropic_messages", "anthropic_messages") + assert logged.payload["id"] == served_id + assert logged.payload["response"]["choices"][0]["message"]["content"] == "pong" + + +@pytest.mark.parametrize("caller", ["completion", "acompletion"]) +@pytest.mark.asyncio +async def test_bridged_chat_stream_keeps_the_upstream_headers_when_its_output_cannot_be_translated( + monkeypatch: pytest.MonkeyPatch, caller: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _LoggedCallTypes() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + response_body: Final = {**_searched_responses_body("openai/gpt-5.3-codex"), "output": []} + request: Final = _bridged_request("openai/gpt-5.3-codex", {}, stream=True) + + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith="https://api.openai.com/v1/responses").mock( + return_value=httpx.Response( + 200, + text=_responses_sse(response_body), + headers={"content-type": "text/event-stream", "x-request-id": "req_upstream"}, + ) + ) + await _served_chunk_ids(caller, request) + logged: Final = await asyncio.to_thread(recorder.seen_by_caller[caller].get, timeout=10) + + assert isinstance(logged.complete_streaming_response, ModelResponse) + assert logged.complete_streaming_response.id == response_body["id"] + assert logged.payload["hidden_params"]["additional_headers"]["llm_provider-x-request-id"] == "req_upstream" + + def test_convert_chat_completion_messages_to_responses_api_image_input(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, @@ -3064,6 +3479,38 @@ def test_map_optional_params_tool_choice_chat_nested_to_responses_api(): {"type": "custom", "name": "ApplyPatch"}, ), ({"type": "custom"}, {"type": "custom"}), + ( + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [ + {"type": "function", "function": {"name": "get_time"}}, + {"type": "custom", "custom": {"name": "apply_patch"}}, + {"type": "mcp", "server_label": "docs"}, + {"type": "function", "name": "already_flat"}, + ], + }, + }, + { + "type": "allowed_tools", + "mode": "required", + "tools": [ + {"type": "function", "name": "get_time"}, + {"type": "custom", "name": "apply_patch"}, + {"type": "mcp", "server_label": "docs"}, + {"type": "function", "name": "already_flat"}, + ], + }, + ), + ( + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}}, + ), + ( + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + ), ], ) def test_normalize_tool_choice_for_responses_api(tool_choice, expected): @@ -4105,6 +4552,16 @@ def test_output_item_done_with_stream_map_keeps_empty_delta(): ({"type": "required"}, "required"), ("auto", "auto"), ({"type": "function", "function": {"name": "get_weather"}}, {"type": "function", "name": "get_weather"}), + ( + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [{"type": "function", "function": {"name": "get_weather"}}], + }, + }, + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "get_weather"}]}, + ), ], ) async def test_acompletion_bridge_normalizes_tool_choice_on_the_wire( @@ -4177,6 +4634,7 @@ async def test_acompletion_bridge_normalizes_tool_choice_on_the_wire( post_kwargs = mock_post.call_args.kwargs request_body = post_kwargs["json"] if "json" in post_kwargs else json.loads(post_kwargs["data"]) assert request_body["tool_choice"] == expected_wire_tool_choice + assert [tool["name"] for tool in request_body["tools"]] == ["get_weather"] def _make_incomplete_responses_api_response( @@ -5439,6 +5897,38 @@ def test_transform_request_strips_the_breakpoint_when_the_served_providers_row_o ] +@pytest.mark.parametrize( + ("sibling_flag", "model_info", "keeps_the_breakpoint"), + [ + (True, {"id": "flagless"}, False), + (True, {"id": "opted-out", "supports_prompt_cache_breakpoint": False}, False), + (False, {"id": "opted-in", "supports_prompt_cache_breakpoint": True}, True), + ], + ids=["flagless", "opted-out", "opted-in"], +) +def test_transform_request_reads_the_routed_deployments_own_flag_not_a_siblings_on_the_shared_key( + monkeypatch: pytest.MonkeyPatch, + sibling_flag: bool, + model_info: Mapping[str, object], + keeps_the_breakpoint: bool, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "bedrock_mantle/qa-mantle-shared", + {**_MANTLE_GPT_ROW, "supports_prompt_cache_breakpoint": sibling_flag}, + ) + + system_content: Final = _bridged_input_for_served_provider( + "qa-mantle-shared", "bedrock_mantle", model_info=model_info + )[0]["content"] + + assert system_content == [ + {**_MARKED_SYSTEM_PART, "type": "input_text"} + if keeps_the_breakpoint + else {"type": "input_text", "text": "Stable prefix"} + ] + + def test_transform_request_strips_the_breakpoint_for_a_served_provider_without_a_row_of_its_own( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -5884,9 +6374,9 @@ _HAND_WRITTEN_PROMPT_CACHE_BREAKPOINT_BLOCKS: Final = ( ) -def test_convert_chat_completion_messages_to_responses_api_drops_malformed_prompt_cache_breakpoint_under_drop_params(): - """OpenAI's Responses API answered "Supported values are: '30m'" for a 1h breakpoint ttl on 2026-10-07, - so an unsupported ttl drops the marker as a unit while an unknown key is dropped from a valid one.""" +def test_convert_chat_completion_messages_to_responses_api_drops_only_an_unreadable_prompt_cache_breakpoint_under_drop_params(): + """The target's Responses config decides which ttl and keys its API takes, so the bridge keeps every marker + whose mode some target reads.""" handler = LiteLLMResponsesTransformationHandler() messages = [{"role": "user", "content": list(_HAND_WRITTEN_PROMPT_CACHE_BREAKPOINT_BLOCKS)}] @@ -5899,12 +6389,12 @@ def test_convert_chat_completion_messages_to_responses_api_drops_malformed_promp None, None, None, - None, - {"mode": "explicit"}, + {"mode": "explicit", "ttl": "1h"}, + {"mode": "explicit", "scope": "all"}, {"mode": "explicit", "ttl": "30m"}, {"mode": "explicit"}, ] - assert all("prompt_cache_breakpoint" not in block for block in content[:4]) + assert all("prompt_cache_breakpoint" not in block for block in content[:3]) def test_convert_chat_completion_messages_to_responses_api_keeps_malformed_prompt_cache_breakpoint_by_default(): @@ -5937,6 +6427,41 @@ def test_transform_request_drop_params_in_litellm_params_gates_the_prompt_cache_ assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0] +_TTL_MARKED_CONTENT: Final = ( + {"type": "text", "text": "unsupported ttl", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}}, + {"type": "text", "text": "well formed", "prompt_cache_breakpoint": {"mode": "explicit"}}, +) + + +@pytest.mark.parametrize( + "global_drop_params,deployment_drop_params,expected", + [ + ("true", None, [None, {"mode": "explicit"}]), + (False, True, [{"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}]), + (True, None, [{"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}]), + ], + ids=("string-global-target-keeps-raw-markers", "deployment-drop-params", "global-drop-params"), +) +def test_transform_request_hands_a_marker_to_the_target_whole_only_when_the_target_trims_it( + monkeypatch: pytest.MonkeyPatch, + global_drop_params: object, + deployment_drop_params: bool | None, + expected: list[object], +) -> None: + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + request: Final = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-6.1-sol", + messages=[{"role": "user", "content": list(_TTL_MARKED_CONTENT)}], + optional_params={}, + litellm_params={"custom_llm_provider": "openai", "drop_params": deployment_drop_params}, + headers={}, + litellm_logging_obj=Mock(), + ) + (item,) = cast(list[dict[str, object]], request["input"]) # cast-ok: the bridge emits message item mappings + parts: Final = cast(list[dict[str, object]], item["content"]) # cast-ok: message items carry content parts + assert [part.get("prompt_cache_breakpoint") for part in parts] == expected + + def test_convert_response_output_folds_two_messages_and_function_call_into_one_choice() -> None: first: Final = ResponseOutputMessage.model_validate( { diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index ed7dd5ec146..914eb1f2508 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -216,6 +216,7 @@ def _reset_aws_auth_caches() -> None: def _flush_client_caches() -> None: litellm.in_memory_llm_clients_cache.flush_cache() image_handling_module.in_memory_cache.flush_cache() + image_handling_module.pdf_probe_verdicts.flush_cache() _reset_aws_auth_caches() diff --git a/tests/unit/cookbook/__init__.py b/tests/unit/cookbook/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/cookbook/test_gpt_realtime_translate.py b/tests/unit/cookbook/test_gpt_realtime_translate.py new file mode 100644 index 00000000000..a40a3265600 --- /dev/null +++ b/tests/unit/cookbook/test_gpt_realtime_translate.py @@ -0,0 +1,141 @@ +import asyncio +import base64 +import json +import wave +from pathlib import Path +from types import SimpleNamespace +from typing import Final, cast + +import pytest +from websockets.asyncio.client import ClientConnection + +from cookbook import gpt_realtime_translate as translate + + +@pytest.mark.asyncio +async def test_short_upload_waits_for_first_translated_audio(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(translate, "RESPONSE_TIMEOUT_SECONDS", 0.1) + audio: Final = bytes(480) + events: Final = iter( + ( + {"type": "session.output_audio.delta", "delta": base64.b64encode(audio).decode()}, + {"type": "error", "error": {"message": "session closed"}}, + ) + ) + + async def recv() -> str: + event: Final = next(events) + if event["type"] == "session.output_audio.delta": + await asyncio.sleep(0.03) + return json.dumps(event) + + sender_finished: Final = asyncio.Event() + sender_finished.set() + output: Final = tmp_path / "translation.wav" + + result: Final = await translate.receive_translation( + cast(ClientConnection, SimpleNamespace(recv=recv)), output, sender_finished + ) + + assert result == 'Realtime API error: {"message": "session closed"}' + with wave.open(str(output), "rb") as rendered: + assert rendered.readframes(240) == audio + + +@pytest.mark.asyncio +async def test_long_upload_keeps_waiting_for_first_translated_audio_while_still_sending( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(translate, "RESPONSE_TIMEOUT_SECONDS", 0.05) + audio: Final = bytes(480) + pending: Final[asyncio.Queue[str]] = asyncio.Queue() + sender_finished: Final = asyncio.Event() + + async def speak_after_a_long_pause() -> None: + await asyncio.sleep(0.2) + await pending.put(json.dumps({"type": "session.output_audio.delta", "delta": base64.b64encode(audio).decode()})) + await asyncio.sleep(0.02) + sender_finished.set() + await pending.put(json.dumps({"type": "session.closed"})) + + output: Final = tmp_path / "translation.wav" + speaker: Final = asyncio.create_task(speak_after_a_long_pause()) + + result: Final = await translate.receive_translation( + cast(ClientConnection, SimpleNamespace(recv=pending.get)), output, sender_finished + ) + await speaker + + assert result is None + with wave.open(str(output), "rb") as rendered: + assert rendered.readframes(240) == audio + + +@pytest.mark.asyncio +async def test_a_quiet_gap_after_the_sender_stops_keeps_the_rest_of_the_translation( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(translate, "RESPONSE_TIMEOUT_SECONDS", 1.0) + first: Final = bytes(480) + second: Final = bytes(range(240)) * 2 + events: Final = iter( + ( + {"type": "session.output_audio.delta", "delta": base64.b64encode(first).decode()}, + {"type": "session.output_audio.delta", "delta": base64.b64encode(second).decode()}, + {"type": "session.closed"}, + ) + ) + + async def recv() -> str: + event: Final = next(events) + if event.get("delta") == base64.b64encode(second).decode(): + await asyncio.sleep(0.3) + return json.dumps(event) + + sender_finished: Final = asyncio.Event() + sender_finished.set() + output: Final = tmp_path / "translation.wav" + + result: Final = await translate.receive_translation( + cast(ClientConnection, SimpleNamespace(recv=recv)), output, sender_finished + ) + + assert result is None + with wave.open(str(output), "rb") as rendered: + assert rendered.readframes(480) == first + second + + +@pytest.mark.asyncio +async def test_a_stream_that_never_closes_the_session_reports_an_error( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(translate, "RESPONSE_TIMEOUT_SECONDS", 0.05) + audio: Final = bytes(480) + pending: Final[asyncio.Queue[str]] = asyncio.Queue() + await pending.put(json.dumps({"type": "session.output_audio.delta", "delta": base64.b64encode(audio).decode()})) + sender_finished: Final = asyncio.Event() + sender_finished.set() + + result: Final = await translate.receive_translation( + cast(ClientConnection, SimpleNamespace(recv=pending.get)), tmp_path / "translation.wav", sender_finished + ) + + assert result == "The translation stream went quiet for 0.05 s without closing the session" + + +@pytest.mark.asyncio +async def test_the_sender_closes_the_session_after_the_last_audio_chunk() -> None: + sent: Final[asyncio.Queue[str]] = asyncio.Queue() + finished: Final = asyncio.Event() + pcm: Final = bytes(translate.CHUNK_BYTES + 2) + + await translate.send_audio(cast(ClientConnection, SimpleNamespace(send=sent.put)), pcm, 0.0, finished) + + messages: Final = tuple(json.loads(sent.get_nowait()) for _ in range(sent.qsize())) + assert [message["type"] for message in messages] == [ + "session.input_audio_buffer.append", + "session.input_audio_buffer.append", + "session.close", + ] + assert b"".join(base64.b64decode(message["audio"]) for message in messages[:-1]) == pcm + assert finished.is_set() diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index 89fff365c3f..942433b39fb 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -9,18 +9,24 @@ import asyncio import base64 import json import logging +from collections.abc import Callable, Mapping from types import MappingProxyType import pytest -from typing import Final, Optional +from typing import TYPE_CHECKING, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch +from litellm._internal_context import internal_parent_request_id, is_internal_call from litellm.llms.base_llm.files.transformation import FileListModelScope from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM +from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.types.llms.openai import FileListPage, OpenAIFileObject from litellm.types.utils import LiteLLMBatch +if TYPE_CHECKING: + from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles + def _make_file_object(file_id: str = "file-output-abc") -> OpenAIFileObject: return OpenAIFileObject( @@ -2268,3 +2274,324 @@ async def test_batch_create_hook_passes_the_router_deployment_as_creating_model_ first, second = managed_files.store_unified_object_id.await_args_list assert first.kwargs["creating_model_id"] == "deploy-team-b" assert second.kwargs["creating_model_id"] is None + + +class _FakeFilesRouter: + def __init__( + self, + deployments_by_model: Mapping[str, tuple[str, ...]], + delete_error: Exception | None = None, + delete_gate: asyncio.Event | None = None, + provider: str = "openai", + gone_file_ids: frozenset[str] = frozenset(), + ) -> None: + self.deployments_by_model = deployments_by_model + self.delete_error = delete_error + self.gone_file_ids = gone_file_ids + self.delete_gate = delete_gate + self.provider = provider + self.upload_count = 0 + self.deletes: tuple[tuple[str, str, Mapping[str, object], Mapping[str, str]], ...] = () + self.deletes_marked_internal: tuple[bool, ...] = () + self.delete_parent_request_ids: tuple[str | None, ...] = () + + async def acreate_file( + self, + model: str, + file: tuple[str, bytes, str], + purpose: str, + litellm_metadata: Mapping[str, object], + ) -> OpenAIFileObject: + deployment_ids: Final = self.deployments_by_model[model] + provider_file_ids: Final = tuple( + f"file-upload-{self.upload_count + offset + 1}" for offset in range(len(deployment_ids)) + ) + self.upload_count += len(deployment_ids) + response: Final = _make_file_object(provider_file_ids[0]).model_copy(update={"purpose": purpose}) + response._hidden_params = { + "model_id": deployment_ids[0], + "custom_llm_provider": self.provider, + "model_file_id_mapping": dict(zip(deployment_ids, provider_file_ids)), + "litellm_call_id": f"call-of-{provider_file_ids[0]}", + } + return response + + def registry(self) -> "_FakeFilesRouter": + return self + + def credentials(self, model_id: str, form: str) -> dict[str, str]: + assert form == "with_provider" + return {"custom_llm_provider": self.provider, "api_key": f"key-of-{model_id}"} + + async def afile_delete( + self, + model: str, + file_id: str, + litellm_metadata: Mapping[str, object], + _litellm_internal_model_credentials: Mapping[str, str], + litellm_call_id: str | None = None, + ) -> None: + if self.delete_gate is not None: + await self.delete_gate.wait() + if isinstance(litellm_metadata, dict): + litellm_metadata.update({"model_info": {"id": model}}) + self.deletes = (*self.deletes, (model, file_id, litellm_metadata, _litellm_internal_model_credentials)) + self.deletes_marked_internal = (*self.deletes_marked_internal, is_internal_call.get()) + self.delete_parent_request_ids = (*self.delete_parent_request_ids, internal_parent_request_id.get()) + if file_id in self.gone_file_ids: + raise RuntimeError(f"No such File object: {file_id}") + if self.delete_error is not None: + raise self.delete_error + + +_REQUEST_METADATA: Final = MappingProxyType({"user_api_key_team_id": "test-team"}) + + +async def _eventually(condition: Callable[[], bool]) -> None: + async def _poll() -> None: + while not condition(): + await asyncio.sleep(0) + + await asyncio.wait_for(_poll(), timeout=1) + + +async def _upload( + managed_files: "PROXY_LiteLLMManagedFiles", + router: _FakeFilesRouter, + target_model_names: list[str], + litellm_metadata: Mapping[str, object] = _REQUEST_METADATA, + content_type: str = "text/plain", +) -> OpenAIFileObject: + return await managed_files.acreate_file( + create_file_request={ + "file": ("notes.txt", b"the launch code is ORCHID-42\n", content_type), + "purpose": "user_data", + "litellm_metadata": litellm_metadata, + }, + llm_router=router, + target_model_names_list=target_model_names, + litellm_parent_otel_span=None, + user_api_key_dict=_make_team_member_api_key_dict(), + ) + + +async def _create_managed_file( + router: _FakeFilesRouter, target_model_names: list[str] +) -> tuple[dict[str, str], list[str], tuple[str, str]]: + managed_files, mock_prisma = _make_real_managed_files_instance() + response: Final = await _upload(managed_files, router, target_model_names) + stored: Final = mock_prisma.db.litellm_managedfiletable.upsert.await_args.kwargs["data"]["create"] + decoded: Final = base64.urlsafe_b64decode(response.id + "=" * (-len(response.id) % 4)).decode() + carried: Final = ( + managed_files.get_model_id_from_unified_file_id(decoded), + managed_files.get_output_file_id_from_unified_file_id(decoded), + ) + return json.loads(stored["model_mappings"]), stored["flat_model_file_ids"], carried + + +@pytest.mark.asyncio +async def test_acreate_file_keeps_the_upload_the_managed_id_carries_when_two_targets_share_a_deployment(): + router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}) + + model_mappings, flat_ids, carried = await _create_managed_file(router, ["gpt-a", "gpt-b"]) + await _eventually(lambda: len(router.deletes) == 1) + + assert carried == ("wildcard-dep", "file-upload-1") + assert model_mappings == {"wildcard-dep": "file-upload-1"} + assert flat_ids == ["file-upload-1"] + assert router.deletes == ( + ( + "wildcard-dep", + "file-upload-2", + _REQUEST_METADATA, + {"custom_llm_provider": "openai", "api_key": "key-of-wildcard-dep"}, + ), + ) + assert isinstance(router.deletes[0][3], MappingProxyType) + + +@pytest.mark.asyncio +async def test_acreate_file_deletes_only_the_repeat_upload_and_keeps_every_deployment_reached_once(): + router: Final = _FakeFilesRouter({"gpt-a": ("dep-1", "dep-2"), "gpt-b": ("dep-2",), "gpt-c": ("dep-3",)}) + + model_mappings, flat_ids, carried = await _create_managed_file(router, ["gpt-a", "gpt-b", "gpt-c"]) + await _eventually(lambda: len(router.deletes) == 1) + + assert carried == ("dep-1", "file-upload-1") + assert list(model_mappings.items()) == [ + ("dep-1", "file-upload-1"), + ("dep-2", "file-upload-2"), + ("dep-3", "file-upload-4"), + ] + assert flat_ids == ["file-upload-1", "file-upload-2", "file-upload-4"] + assert [delete[:2] for delete in router.deletes] == [("dep-2", "file-upload-3")] + + +@pytest.mark.asyncio +async def test_acreate_file_still_returns_the_managed_file_when_deleting_the_repeat_upload_fails(caplog: pytest.LogCaptureFixture): + router: Final = _FakeFilesRouter( + {"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}, + delete_error=RuntimeError("provider unavailable"), + ) + + with caplog.at_level(logging.WARNING): + model_mappings, flat_ids, carried = await _create_managed_file(router, ["gpt-a", "gpt-b"]) + await _eventually(lambda: "provider unavailable" in caplog.text) + + assert carried == ("wildcard-dep", "file-upload-1") + assert model_mappings == {"wildcard-dep": "file-upload-1"} + assert flat_ids == ["file-upload-1"] + assert [delete[:2] for delete in router.deletes] == [("wildcard-dep", "file-upload-2")] + assert "file-upload-2" in caplog.text and "wildcard-dep" in caplog.text + assert "provider unavailable" in caplog.text + + +@pytest.mark.asyncio +async def test_acreate_file_answers_before_deleting_the_repeat_upload_finishes(): + delete_gate: Final = asyncio.Event() + router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}, delete_gate=delete_gate) + + model_mappings, flat_ids, carried = await asyncio.wait_for( + _create_managed_file(router, ["gpt-a", "gpt-b"]), timeout=1 + ) + + assert carried == ("wildcard-dep", "file-upload-1") + assert model_mappings == {"wildcard-dep": "file-upload-1"} + assert flat_ids == ["file-upload-1"] + assert router.deletes == () + delete_gate.set() + await _eventually(lambda: len(router.deletes) == 1) + assert router.deletes[0][:2] == ("wildcard-dep", "file-upload-2") + + +@pytest.mark.asyncio +async def test_shutdown_drain_waits_for_the_repeat_upload_delete(): + delete_gate: Final = asyncio.Event() + router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}, delete_gate=delete_gate) + + await _create_managed_file(router, ["gpt-a", "gpt-b"]) + asyncio.get_running_loop().call_later(0.05, delete_gate.set) + await drain_logging_tasks(timeout=1) + + assert [delete[:2] for delete in router.deletes] == [("wildcard-dep", "file-upload-2")] + + +@pytest.mark.asyncio +async def test_each_repeat_upload_delete_runs_as_an_internal_call_of_the_upload_that_made_the_copy(): + router: Final = _FakeFilesRouter( + {"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",), "gpt-c": ("wildcard-dep",)} + ) + + await _create_managed_file(router, ["gpt-a", "gpt-b", "gpt-c"]) + await drain_logging_tasks(timeout=1) + + assert router.deletes_marked_internal == (True, True) + assert sorted(zip((delete[1] for delete in router.deletes), router.delete_parent_request_ids)) == [ + ("file-upload-2", "call-of-file-upload-2"), + ("file-upload-3", "call-of-file-upload-3"), + ] + assert is_internal_call.get() is False + assert internal_parent_request_id.get() is None + + +@pytest.mark.asyncio +async def test_litellm_proxy_wildcard_keeps_its_last_per_target_upload_and_deletes_nothing(): + router: Final = _FakeFilesRouter( + {"gpt-a": ("proxy-wildcard-dep",), "gpt-b": ("proxy-wildcard-dep",)}, provider="litellm_proxy" + ) + + model_mappings, flat_ids, _ = await _create_managed_file(router, ["gpt-a", "gpt-b"]) + await drain_logging_tasks(timeout=1) + + assert model_mappings == {"proxy-wildcard-dep": "file-upload-2"} + assert flat_ids == ["file-upload-2"] + assert router.deletes == () + + +async def _upload_then_delete( + router: _FakeFilesRouter, target_model_names: list[str], content_type: str = "text/plain" +) -> tuple[str, ...]: + managed_files, mock_prisma = _make_real_managed_files_instance() + response: Final = await _upload(managed_files, router, target_model_names, content_type=content_type) + await drain_logging_tasks(timeout=1) + deletes_during_upload: Final = len(router.deletes) + managed_files.internal_usage_cache.async_get_cache = AsyncMock( + return_value=managed_files.internal_usage_cache.async_set_cache.await_args.kwargs["value"] + ) + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_managedfiletable.delete = AsyncMock() + + deleted: Final = await managed_files.afile_delete( + file_id=response.id, + litellm_parent_otel_span=None, + llm_router=router, + litellm_metadata=_REQUEST_METADATA, + litellm_call_id="delete-call", + ) + + assert deleted.deleted is True + mock_prisma.db.litellm_managedfiletable.delete.assert_awaited_once_with(where={"unified_file_id": response.id}) + return tuple(delete[1] for delete in router.deletes[deletes_during_upload:]) + + +@pytest.mark.asyncio +async def test_deleting_a_litellm_proxy_wildcard_file_removes_the_upload_its_id_carries_too(): + router: Final = _FakeFilesRouter( + {"gpt-a": ("proxy-wildcard-dep",), "gpt-b": ("proxy-wildcard-dep",)}, provider="litellm_proxy" + ) + + assert await _upload_then_delete(router, ["gpt-a", "gpt-b"]) == ("file-upload-2", "file-upload-1") + + +@pytest.mark.asyncio +async def test_deleting_a_wildcard_file_removes_its_kept_upload_once(): + router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}) + + assert await _upload_then_delete(router, ["gpt-a", "gpt-b"]) == ("file-upload-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content_type", + [ + "text/plain;llm_output_file_id,file-of-another-team;llm_output_file_model_id,wildcard-dep", + "text/plain;llm_output_file_id,file-of-another-team", + ], +) +async def test_deleting_a_file_removes_only_the_uploads_the_gateway_made_for_it(content_type: str): + router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}) + + assert await _upload_then_delete(router, ["gpt-a", "gpt-b"], content_type) == ("file-upload-1",) + + +@pytest.mark.asyncio +async def test_deleting_a_file_still_succeeds_when_the_upload_its_id_carries_is_already_gone( + caplog: pytest.LogCaptureFixture, +): + router: Final = _FakeFilesRouter( + {"gpt-a": ("proxy-wildcard-dep",), "gpt-b": ("proxy-wildcard-dep",)}, + provider="litellm_proxy", + gone_file_ids=frozenset({"file-upload-1"}), + ) + + with caplog.at_level(logging.WARNING): + assert await _upload_then_delete(router, ["gpt-a", "gpt-b"]) == ("file-upload-2", "file-upload-1") + + assert "file-upload-1" in caplog.text and "No such File object" in caplog.text + assert router.deletes_marked_internal == (False, True) + assert router.delete_parent_request_ids == (None, "delete-call") + + +@pytest.mark.asyncio +async def test_repeat_upload_delete_leaves_the_upload_request_metadata_unchanged(): + router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}) + upload_metadata: Final = {"user_api_key_team_id": "test-team"} + managed_files, _ = _make_real_managed_files_instance() + + await _upload(managed_files, router, ["gpt-a", "gpt-b"], upload_metadata) + await drain_logging_tasks(timeout=1) + + assert upload_metadata == {"user_api_key_team_id": "test-team"} + assert [delete[2] for delete in router.deletes] == [ + {"user_api_key_team_id": "test-team", "model_info": {"id": "wildcard-dep"}} + ] diff --git a/tests/unit/integrations/azure_storage/test_azure_storage.py b/tests/unit/integrations/azure_storage/test_azure_storage.py index 0f0ce24ff84..ee12291e7a0 100644 --- a/tests/unit/integrations/azure_storage/test_azure_storage.py +++ b/tests/unit/integrations/azure_storage/test_azure_storage.py @@ -5,19 +5,28 @@ import re import sys import threading from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS +from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.azure_storage import azure_storage as azure_storage_module from litellm.integrations.azure_storage.azure_storage import ( AzureBlobStorageLogger, _cached_credential_chain_token_provider, adls_safe_file_name, ) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType -from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload +from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + StandardAuditLogPayload, + StandardLoggingPayload, +) @pytest.fixture @@ -648,3 +657,171 @@ async def test_audit_log_is_uploaded_to_a_dated_audit_logs_path_with_an_account_ assert json.loads(uploaded) == audit_log assert file_client.append_data.await_args.kwargs == {"data": uploaded, "offset": 0, "length": len(uploaded)} file_client.flush_data.assert_awaited_once_with(position=len(uploaded), offset=0) + + +BACKGROUND_RESPONSE_ID: Final = "resp_YWJjZA==" +BACKGROUND_COST_POLL_KWARGS: Final = { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}, + }, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} + + +def _background_create_cost_and_other_events( + cost_kwargs: dict[str, object], +) -> tuple[tuple[dict[str, object], StandardLoggingPayload], ...]: + create: Final = StandardLoggingPayload( + id=BACKGROUND_RESPONSE_ID, + call_type="aresponses", + messages=[{"role": "user", "content": "background create input"}], + ) + cost: Final = StandardLoggingPayload(id=BACKGROUND_RESPONSE_ID, call_type="aget_responses", response_cost=0.42) + ordinary: Final = StandardLoggingPayload(id="chatcmpl-ordinary", call_type="acompletion") + batch_retrieve: Final = StandardLoggingPayload(id="batch_abc123", call_type="aretrieve_batch", response_cost=0.07) + return ( + ({"call_type": "aresponses", "litellm_params": {"metadata": {}}, "standard_logging_object": create}, create), + ({**cost_kwargs, "standard_logging_object": cost}, cost), + ( + {"call_type": "acompletion", "litellm_params": {"metadata": {}}, "standard_logging_object": ordinary}, + ordinary, + ), + ( + { + "call_type": "aretrieve_batch", + "litellm_params": {"metadata": {}}, + "standard_logging_object": batch_retrieve, + }, + batch_retrieve, + ), + ) + + +def _expected_files( + events: tuple[tuple[dict[str, object], StandardLoggingPayload], ...], +) -> dict[str, StandardLoggingPayload]: + create, cost, ordinary, batch_retrieve = (payload for _, payload in events) + return { + "resp_YWJjZA.json": create, + "resp_YWJjZA_background_cost.json": cost, + "chatcmpl-ordinary.json": ordinary, + "batch_abc123.json": batch_retrieve, + } + + +async def _log_and_flush( + logger: AzureBlobStorageLogger, events: tuple[tuple[dict[str, object], StandardLoggingPayload], ...] +) -> None: + for kwargs, payload in events: + await logger.async_log_success_event(kwargs, {"id": payload["id"]}, None, None) + await logger.async_send_batch() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), +) +async def test_background_cost_event_keeps_its_own_file_beside_the_create_on_the_entra_path( + mock_env_vars, monkeypatch, cost_kwargs +): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + events: Final = _background_create_cost_and_other_events(cost_kwargs) + with patch( # test-quality-ok: REST client is created inside the method; assert emitted requests + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_http_client.put.return_value = MagicMock() + mock_http_client.patch.return_value = MagicMock() + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger() + logger.azure_auth_token = "mock-azure-ad-token" + logger.token_expiry = None + await _log_and_flush(logger, events) + + base_url: Final = "https://test-account.dfs.core.windows.net/test-container/" + appends: Final = [call for call in mock_http_client.patch.call_args_list if "action=append" in call.args[0]] + uploaded: Final = { + call.args[0].removeprefix(base_url).split("?")[0]: json.loads(call.kwargs["data"]) for call in appends + } + assert uploaded == _expected_files(events), ( + "the create, its background cost event and every other event must each keep their own file" + ) + + +@pytest.mark.asyncio +async def test_failure_event_is_uploaded_to_a_file_named_after_its_id(mock_env_vars, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + failure: Final = StandardLoggingPayload(id="chatcmpl-failed", call_type="acompletion", error_str="provider failure") + with patch( # test-quality-ok: REST client is created inside the method; assert emitted requests + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_http_client.put.return_value = MagicMock() + mock_http_client.patch.return_value = MagicMock() + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger() + logger.azure_auth_token = "mock-azure-ad-token" + logger.token_expiry = None + await logger.async_log_failure_event( + {"call_type": "acompletion", "litellm_params": {"metadata": {}}, "standard_logging_object": failure}, + None, + None, + None, + ) + await logger.async_send_batch() + + base_url: Final = "https://test-account.dfs.core.windows.net/test-container/" + appends: Final = [call for call in mock_http_client.patch.call_args_list if "action=append" in call.args[0]] + assert {call.args[0].removeprefix(base_url).split("?")[0]: json.loads(call.kwargs["data"]) for call in appends} == { + "chatcmpl-failed.json": failure + } + + +@pytest.mark.asyncio +async def test_background_cost_event_keeps_its_own_file_beside_the_create_with_an_account_key( + workload_identity_env_vars, monkeypatch +): + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr(azure_storage_module, "datetime", _FrozenDatetime) + events: Final = _background_create_cost_and_other_events(BACKGROUND_COST_POLL_KWARGS) + + uploaded: Final[dict[str, StandardLoggingPayload]] = {} + + def file_client_for(file_path: str) -> MagicMock: + file_client = MagicMock() + file_client.create_file = AsyncMock() + file_client.flush_data = AsyncMock() + + async def append_data(data: bytes, offset: int, length: int) -> None: + uploaded[file_path] = json.loads(data) + + file_client.append_data = append_data + return file_client + + file_system_client = MagicMock() + file_system_client.get_file_client = MagicMock(side_effect=file_client_for) + service_client = MagicMock() + service_client.get_file_system_client = MagicMock(return_value=file_system_client) + fake_aio_module = MagicMock() + fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client) + + with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}): + logger = AzureBlobStorageLogger() + await _log_and_flush(logger, events) + + today: Final = _FrozenDatetime.now().strftime("%Y-%m-%d") + assert uploaded == {f"{today}/{name}": payload for name, payload in _expected_files(events).items()}, ( + "the create, its background cost event and every other event must each keep their own file" + ) diff --git a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py index f94d7bbee37..42e69f95bb4 100644 --- a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -20,10 +20,16 @@ from litellm.integrations.clickhouse.clickhouse_spend_logger import ( strip_cache_hit_suffix, ) from litellm.integrations.clickhouse.context import lens_analysis +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.secret_redaction import REDACTED from litellm.tracing.types import SpendLogRecord -from litellm.types.utils import StandardLoggingPayload +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, StandardLoggingPayload _JSON_OBJECT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue]) @@ -349,6 +355,41 @@ def test_cache_hit_id_is_stripped_for_response_id(): assert strip_cache_hit_suffix("chatcmpl-xyz") == "chatcmpl-xyz" +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + }, + id="cost-poll", + ), + pytest.param( + { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, + }, + id="passthrough-settlement", + ), + ), +) +def test_background_cost_row_keeps_its_own_request_id_beside_its_create(cost_kwargs: Mapping[str, object]) -> None: + payload: Final = cast(StandardLoggingPayload, _payload(id="resp_YWJjZA", call_type="aresponses")) + create_kwargs: Final[Mapping[str, object]] = {"call_type": "aresponses", "litellm_params": {"metadata": {}}} + + rows: Final = [spend_log_row_from_payload(payload, kwargs) for kwargs in (create_kwargs, cost_kwargs)] + + assert [row["request_id"] for row in rows] == [ + "resp_YWJjZA", + f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}", + ] + assert [row["response_id"] for row in rows] == ["resp_YWJjZA", "resp_YWJjZA"] + + def test_parse_traceparent_valid_missing_malformed(): assert parse_traceparent(TRACEPARENT) == (TRACE_ID, SPAN_ID) assert parse_traceparent(None) == ("", "") diff --git a/tests/unit/integrations/conftest.py b/tests/unit/integrations/conftest.py index 48a01ed8d48..9278a280fe7 100644 --- a/tests/unit/integrations/conftest.py +++ b/tests/unit/integrations/conftest.py @@ -1,14 +1,18 @@ import functools import http.server +import importlib.abc import ipaddress import os import queue import ssl +import sys import threading -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from importlib.machinery import ModuleSpec from pathlib import Path +from types import ModuleType from typing import Final import pytest @@ -75,6 +79,36 @@ def write_self_signed_cert(directory: Path, stem: str) -> tuple[Path, Path]: return certificate_path, key_path +PROXY_ONLY_PACKAGES: Final = frozenset({"fastapi", "starlette"}) + + +class _ProxyOnlyPackageBlocker(importlib.abc.MetaPathFinder): + def find_spec( + self, + fullname: str, + path: Sequence[str] | None = None, + target: ModuleType | None = None, + ) -> ModuleSpec | None: + if fullname.partition(".")[0] in PROXY_ONLY_PACKAGES: + raise ModuleNotFoundError(f"No module named {fullname!r}") + return None + + +def _needs_proxy_extras(module_name: str) -> bool: + return ( + module_name.partition(".")[0] in PROXY_ONLY_PACKAGES + or module_name == "litellm.proxy" + or module_name.startswith("litellm.proxy.") + ) + + +@pytest.fixture +def sdk_only_install(monkeypatch: pytest.MonkeyPatch) -> None: + for module_name in tuple(filter(_needs_proxy_extras, sys.modules)): + monkeypatch.delitem(sys.modules, module_name) + monkeypatch.setattr(sys, "meta_path", [_ProxyOnlyPackageBlocker(), *sys.meta_path]) + + @pytest.fixture(autouse=True) def restore_process_environment() -> Iterator[None]: original: Final = dict(os.environ) diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py index 6f143a244f6..c29985b1e41 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py @@ -13,9 +13,11 @@ import litellm import pytest import respx +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN mock_response_data: Final = { @@ -172,9 +174,7 @@ async def test_get_payload_not_found(monkeypatch): assert payload is None -def _gcs_logger_with_upload_status( - monkeypatch: pytest.MonkeyPatch, status_code: int -) -> tuple[GCSBucketLogger, Mock]: +def _gcs_logger_with_upload_status(monkeypatch: pytest.MonkeyPatch, status_code: int) -> tuple[GCSBucketLogger, Mock]: monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) monkeypatch.setattr(google.auth, "default", _google_default_credentials) monkeypatch.setenv("GCS_FLUSH_INTERVAL", "3600") @@ -394,3 +394,98 @@ async def test_batch_line_items_are_serialized_off_the_event_loop(monkeypatch: p (thread,) = threads assert thread is not threading.main_thread() + + +@pytest.mark.asyncio +async def test_unbatched_background_cost_event_keeps_its_own_object_readable_by_its_spend_log_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "false") + logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200) + object_prefix: Final = "/storage/v1/b/test-bucket/o/" + stored: Final[dict[str, object]] = {} + + def bucket(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + stored[request.url.params["name"]] = json.loads(request.content) + return httpx.Response(200, json={"name": request.url.params["name"]}) + name: Final = request.url.path.removeprefix(object_prefix) + if name not in stored: + return httpx.Response(404, text="No such object") + return httpx.Response(200, content=json.dumps(stored[name]).encode("utf-8")) + + logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(bucket)) + response_id: Final = "resp_YmFja2dyb3VuZA" + create: Final = { + "id": response_id, + "messages": [{"role": "user", "content": "background create input"}], + "error_str": None, + } + cost: Final = {"id": response_id, "response_cost": 0.42, "error_str": None} + ordinary: Final = {"id": "chatcmpl-ordinary", "error_str": None} + batch_retrieve: Final = {"id": "batch_abc123", "response_cost": 0.07, "error_str": None} + events: Final = ( + ({"call_type": "aresponses", "litellm_params": {"metadata": {}}}, create), + ( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + }, + cost, + ), + ({"call_type": "acompletion", "litellm_params": {"metadata": {}}}, ordinary), + ({"call_type": "aretrieve_batch", "litellm_params": {"metadata": {}}}, batch_retrieve), + ) + + for kwargs, payload in events: + await logger.async_log_success_event( + {**kwargs, "standard_logging_object": payload}, {"id": payload["id"]}, None, None + ) + await logger.flush_queue() + + assert {name.split("/", maxsplit=1)[1]: payload for name, payload in stored.items()} == { + response_id: create, + f"{response_id}_background_cost": cost, + "chatcmpl-ordinary": ordinary, + "batch_abc123": batch_retrieve, + }, "the create, its background cost event and every other event must each keep their own object" + now: Final = datetime.now(timezone.utc) + assert await logger.get_request_response_payload(f"{response_id}_background_cost", now, None) == cost + assert await logger.get_request_response_payload(response_id, now, None) == create + + +@pytest.mark.asyncio +async def test_unbatched_gcs_log_id_still_names_a_background_cost_event_object(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "false") + logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200) + stored: Final[dict[str, object]] = {} + + def bucket(request: httpx.Request) -> httpx.Response: + stored[request.url.params["name"]] = json.loads(request.content) + return httpx.Response(200, json={"name": request.url.params["name"]}) + + logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(bucket)) + cost: Final = {"id": "resp_YmFja2dyb3VuZA", "response_cost": 0.42, "error_str": None} + await logger.async_log_success_event( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": { + INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + "gcs_log_id": "qa-run-7", + } + }, + "standard_logging_object": cost, + }, + {"id": cost["id"]}, + None, + None, + ) + await logger.flush_queue() + + ((name, payload),) = stored.items() + object_name: Final = name.split("/", maxsplit=1)[1] + assert object_name.startswith("custom-") and object_name.endswith("-qa-run-7"), object_name + assert payload == cost diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 778c8dfe2f5..766038d5c4f 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -8,14 +8,18 @@ import subprocess import sys import textwrap from collections.abc import Mapping +from types import MappingProxyType from typing import Final, List, Optional, Tuple, cast from unittest.mock import MagicMock, patch import pytest +import respx +from httpx import Response from openai.types.chat import ChatCompletionMessage from pydantic import BaseModel, ConfigDict import litellm +from litellm import Router from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import ( CARRY_UNMATCHED_MESSAGE_POINTS, @@ -1397,6 +1401,109 @@ def _marks_the_route_sends(model: str, custom_llm_provider: str | None, messages _SECOND_SEARCH_CALL: Final = {**_SEARCH_CALL, "id": "srvtoolu_02Search"} +_CLIENT_MARKED_CONVERSATION: Final = ( + {"role": "system", "content": [{"type": "text", "text": "Handbook", "prompt_cache_breakpoint": {"mode": "explicit"}}]}, + {"role": "user", "content": [{"type": "text", "text": "First", "prompt_cache_breakpoint": {"mode": "explicit"}}]}, + {"role": "assistant", "content": "First answer"}, + {"role": "user", "content": "Second"}, +) + + +def _hook_result( + model: str, points: tuple[Mapping[str, object], ...] = ({"location": "message", "index": -1},), **params: object +) -> tuple[list[AllMessageValues], Mapping[str, object]]: + _, processed, returned_params = AnthropicCacheControlHook().get_chat_completion_prompt( + model=model, + messages=cast(list[AllMessageValues], copy.deepcopy(list(_CLIENT_MARKED_CONVERSATION))), + non_default_params={"cache_control_injection_points": list(points), **params}, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + return processed, returned_params + + +def _hook_output( + model: str, points: tuple[Mapping[str, object], ...] = ({"location": "message", "index": -1},), **params: object +) -> list[AllMessageValues]: + return _hook_result(model, points, **params)[0] + + +def test_cache_control_injection_points_replace_client_prompt_cache_breakpoints_on_claude(): + request: Final = AnthropicConfig().transform_request( + model="claude-sonnet-5-5", + messages=_hook_output("anthropic/claude-sonnet-5-5"), + optional_params={"max_tokens": 64}, + litellm_params={}, + headers={}, + ) + + assert _marks_in(request["system"]) == 0 + assert [_marks_in(message) for message in request["messages"]] == [0, 0, 1] + assert "prompt_cache_breakpoint" not in json.dumps(request) + + +@pytest.mark.parametrize( + "sibling_model_infos", + [ + pytest.param((), id="one-deployment"), + pytest.param( + ({"supports_prompt_cache_breakpoint": True, "base_model": "anthropic.claude-unmapped-v1:0"},), + id="beside-a-deployment-of-the-same-profile-with-a-base-model", + ), + ], +) +@pytest.mark.parametrize("route", ["bedrock", "bedrock/converse"]) +def test_cache_control_injection_points_replace_client_prompt_cache_breakpoints_on_a_flagged_bedrock_profile( + local_model_cost_map, sibling_model_infos: tuple[Mapping[str, object], ...], route: str +): + from litellm import Router + + model: Final = f"{route}/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude-flagged" + Router( + model_list=[ + { + "model_name": f"claude-profile-{position}", + "litellm_params": {"model": model, "aws_region_name": "us-east-1"}, + "model_info": model_info, + } + for position, model_info in enumerate(({"supports_prompt_cache_breakpoint": True}, *sibling_model_infos)) + ] + ) + processed, params = _hook_result(model) + + assert "prompt_cache_breakpoint" not in json.dumps(processed) + assert "prompt_cache_options" not in params + assert [_marks_in(message) for message in processed] == [0, 0, 0, 1] + + +def test_tool_config_injection_point_leaves_client_prompt_cache_breakpoints_on_claude(): + processed: Final = _hook_output("anthropic/claude-sonnet-5-5", points=({"location": "tool_config"},)) + + assert [message["content"] for message in processed[:2]] == [ + message["content"] for message in _CLIENT_MARKED_CONVERSATION[:2] + ] + + +@pytest.mark.parametrize( + ("model", "params"), + [ + pytest.param("openai/gpt-5.6", {"api_base": "https://gateway.example.com/v1"}, id="reads-them"), + pytest.param("openai/gpt-4o", {}, id="openai-model-that-predates-them"), + pytest.param("hosted_vllm/qwen3", {"api_base": "http://vllm.example.com/v1"}, id="self-hosted"), + pytest.param("azure_ai/qa-unflagged-gpt", {"api_base": "https://qa.services.ai.azure.com/models"}, id="foundry"), + ], +) +def test_cache_control_injection_points_keep_client_prompt_cache_breakpoints_on_a_route_that_does_not_translate_them( + model: str, params: Mapping[str, object] +): + processed: Final = _hook_output(model, **params) + + assert [message["content"] for message in processed[:2]] == [ + message["content"] for message in _CLIENT_MARKED_CONVERSATION[:2] + ] + + @pytest.mark.parametrize( ("model", "custom_llm_provider"), [ @@ -3485,12 +3592,23 @@ class TestOpenAIPromptCacheBreakpoint: ("gpt-5.6-sol", True), ("gpt-5.6-luna", True), ("gpt-5.7", True), + ("gpt-5.6o", True), + ("gpt-5.10", True), + ("gpt-5.06", True), ("gpt-6", True), ("GPT-5.6", True), ("gpt-5.5", False), ("gpt-5", False), ("gpt-5-chat-latest", False), ("gpt-4.1", False), + ("gpt-5-6-sol", True), + ("gpt-6-1", True), + ("gpt-5-4-mini", False), + ("gpt-5-32k", False), + ("gpt-5-6b", False), + ("gpt-5-06", False), + ("gpt-5-2026-03-01", False), + ("gpt-4-1106-preview", False), ("o3", False), ("claude-sonnet-4-5", False), ], @@ -4950,6 +5068,33 @@ class TestMalformedInjectionPointsAreIgnored: assert kwargs == {"cache_control_injection_points": value} +def test_capping_thousands_of_marks_on_one_tool_result_counts_the_request_a_logarithmic_number_of_times(): + def marked_parts(prefix: str, count: int) -> list[dict]: + return [ + {"type": "text", "text": f"{prefix} {index}", "prompt_cache_breakpoint": {"mode": "explicit"}} + for index in range(count) + ] + + messages: Final = [ + {"role": "user", "content": marked_parts("part", 4)}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": marked_parts("result", 2000)}, + ] + sent_breakpoints: Final = MagicMock(side_effect=AnthropicCacheControlHook.count_converse_cache_points) + + capped: Final = AnthropicCacheControlHook.prompt_cache_breakpoints_as_capped_cache_control( + cast(List[AllMessageValues], messages), True, sent_breakpoints, lambda: 0 + ) + + assert [part.get("cache_control") for part in capped[0]["content"]] == [{"type": "ephemeral"}] * 3 + [None] + assert ["cache_control" in part for part in capped[2]["content"]] == [False] * 1997 + [True] * 3 + assert sent_breakpoints.call_count < 3 * 20 + + def test_cache_control_hook_leaves_the_duplicate_tool_result_warning_to_the_request_builder( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ): @@ -5044,6 +5189,635 @@ def test_cache_control_hook_injects_next_to_a_malformed_message_instead_of_raisi ] +_SHARED_FOUNDRY_MODEL: Final = "azure_ai/qa-shared-gpt-deployment" +_SIBLING_BREAKPOINT_FLAGS: Final = MappingProxyType( + { + "opted-in": MappingProxyType({"supports_prompt_cache_breakpoint": True}), + "opted-out": MappingProxyType({"supports_prompt_cache_breakpoint": False}), + } +) +_FOUNDRY_CHAT_ROUTE: Final = ( + "/models/chat/completions", + MappingProxyType( + { + "id": "chatcmpl-azure-ai", + "object": "chat.completion", + "created": 1700000000, + "model": "qa-shared-gpt-deployment", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + } + ), +) +_FOUNDRY_RESPONSES_ROUTE: Final = ( + "/openai/v1/responses", + MappingProxyType( + { + "id": "resp_azure_ai", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "qa-shared-gpt-deployment", + "output": [ + { + "type": "message", + "id": "msg_azure_ai", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + } + ), +) + + +def _foundry_siblings_sharing_one_model_string( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + registration_order: tuple[str, ...], + upstream: tuple[str, Mapping[str, object]], +) -> tuple[Router, Mapping[str, respx.Route]]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + path, reply = upstream + routes: Final = MappingProxyType( + { + host: respx_mock.post(host=f"{host}.services.ai.azure.com", path=path).mock( + return_value=Response(200, json=dict(reply)) + ) + for host in registration_order + } + ) + router: Final = Router( + model_list=[ + { + "model_name": host, + "litellm_params": { + "model": _SHARED_FOUNDRY_MODEL, + "api_base": f"https://{host}.services.ai.azure.com/models", + "api_key": "azure-ai-test-key", + "cache_control_injection_points": [{"location": "message", "role": "system"}], + }, + "model_info": {"id": f"{host}-id", **_SIBLING_BREAKPOINT_FLAGS.get(host, {})}, + } + for host in registration_order + ], + num_retries=0, + ) + return router, routes + + +async def _send_unmarked_system_prompt(router: Router, host: str, entrypoint: str) -> None: + chat_messages: Final = [{"role": "system", "content": "stable"}, {"role": "user", "content": "hi"}] + match entrypoint: + case "completion": + router.completion(model=host, messages=chat_messages) + case "acompletion": + await router.acompletion(model=host, messages=chat_messages) + case "aanthropic_messages": + await router.aanthropic_messages( + model=host, system="stable", messages=[{"role": "user", "content": "hi"}], max_tokens=16 + ) + case "aresponses": + await router.aresponses(model=host, input=chat_messages) + + +@pytest.mark.parametrize( + ("entrypoint", "upstream"), + [ + ("completion", _FOUNDRY_CHAT_ROUTE), + ("acompletion", _FOUNDRY_CHAT_ROUTE), + ("aanthropic_messages", _FOUNDRY_CHAT_ROUTE), + ("aresponses", _FOUNDRY_RESPONSES_ROUTE), + ], + ids=["completion", "acompletion", "aanthropic_messages", "aresponses"], +) +@pytest.mark.parametrize( + "registration_order", [("opted-in", "opted-out", "unflagged"), ("unflagged", "opted-out", "opted-in")] +) +async def test_foundry_deployment_reads_only_its_own_prompt_cache_flag_not_a_siblings_sharing_its_model_string( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + registration_order: tuple[str, ...], + entrypoint: str, + upstream: tuple[str, Mapping[str, object]], +) -> None: + router, routes = _foundry_siblings_sharing_one_model_string(respx_mock, monkeypatch, registration_order, upstream) + + for host in registration_order: + await _send_unmarked_system_prompt(router, host, entrypoint) + + def sent_prompt_caching(host: str) -> tuple[object, bool]: + body: Final = routes[host].calls.last.request.content + return json.loads(body).get("prompt_cache_options"), b"prompt_cache_breakpoint" in body + + assert {host: sent_prompt_caching(host) for host in registration_order} == { + "opted-in": ({"mode": "implicit"}, True), + "opted-out": (None, False), + "unflagged": (None, False), + } + + +@pytest.mark.parametrize( + ("model", "sibling_flag", "expected"), + [("qa-openai-shared", True, False), ("gpt-99.0-qa-shared", False, True)], + ids=["sibling-opt-in-does-not-turn-it-on", "sibling-opt-out-does-not-turn-it-off"], +) +def test_openai_deployment_reads_the_catalog_not_a_siblings_flag_on_the_shared_key( + monkeypatch: pytest.MonkeyPatch, model: str, sibling_flag: bool, expected: bool +) -> None: + sibling_entry: Final = { + "litellm_provider": "openai", + "mode": "chat", + "supports_prompt_cache_breakpoint": sibling_flag, + } + monkeypatch.setitem(litellm.model_cost, model, sibling_entry) + + assert ( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + model, "openai", "https://api.openai.com/v1", model_info={"id": "flagless-deployment"} + ) + is expected + ) + + +@pytest.mark.parametrize( + ("model", "declared", "api_base", "expected"), + [ + ("qa-openai-shared", True, "https://api.openai.com/v1", True), + ("qa-openai-shared", True, "https://compatible-host.example.com/v1", False), + ("gpt-99.0-qa-shared", False, "https://api.openai.com/v1", False), + ], + ids=["opt-in-on-openai", "opt-in-on-a-compatible-host-stays-off", "opt-out-beats-the-version-rule"], +) +def test_openai_deployment_declared_flag_answers_support_while_the_host_still_decides( + model: str, declared: bool, api_base: str, expected: bool +) -> None: + assert ( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + model, "openai", api_base, model_info={"supports_prompt_cache_breakpoint": declared} + ) + is expected + ) + + +@pytest.mark.parametrize( + "registration_order", [("opted-in", "opted-out", "unflagged"), ("unflagged", "opted-out", "opted-in")] +) +async def test_foundry_responses_drop_params_keeps_a_client_breakpoint_by_the_deployments_own_flag( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, registration_order: tuple[str, ...] +) -> None: + router, routes = _foundry_siblings_sharing_one_model_string( + respx_mock, monkeypatch, registration_order, _FOUNDRY_RESPONSES_ROUTE + ) + client_marked_system_prompt: Final = { + "role": "system", + "content": [{"type": "input_text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}], + } + + for host in registration_order: + await router.aresponses( + model=host, input=[client_marked_system_prompt, {"role": "user", "content": "hi"}], drop_params=True + ) + + assert { + host: b"prompt_cache_breakpoint" in routes[host].calls.last.request.content for host in registration_order + } == { + "opted-in": True, + "opted-out": False, + "unflagged": True, + } + + +@pytest.mark.parametrize( + ("entrypoint", "part_type", "reply"), + [("acompletion", "text", _FOUNDRY_CHAT_ROUTE[1]), ("aresponses", "input_text", _FOUNDRY_RESPONSES_ROUTE[1])], + ids=["chat", "responses"], +) +@pytest.mark.parametrize("registration_order", [("opted-in", "unflagged"), ("unflagged", "opted-in")]) +async def test_azure_drop_params_keeps_a_client_breakpoint_by_the_deployments_own_flag_not_a_siblings( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + registration_order: tuple[str, ...], + entrypoint: str, + part_type: str, + reply: Mapping[str, object], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + routes: Final = MappingProxyType( + { + host: respx_mock.post(host=f"{host}.openai.azure.com").mock(return_value=Response(200, json=dict(reply))) + for host in registration_order + } + ) + router: Final = Router( + model_list=[ + { + "model_name": host, + "litellm_params": { + "model": "azure/gpt-4.1-qa-shared", + "api_base": f"https://{host}.openai.azure.com", + "api_key": "azure-test-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": f"{host}-id", **_SIBLING_BREAKPOINT_FLAGS.get(host, {})}, + } + for host in registration_order + ], + num_retries=0, + ) + marked_system_prompt: Final = { + "role": "system", + "content": [{"type": part_type, "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}], + } + request: Final = [marked_system_prompt, {"role": "user", "content": "hi"}] + + for host in registration_order: + match entrypoint: + case "acompletion": + await router.acompletion(model=host, messages=request, drop_params=True) + case "aresponses": + await router.aresponses(model=host, input=request, drop_params=True) + + assert { + host: b"prompt_cache_breakpoint" in routes[host].calls.last.request.content for host in registration_order + } == {"opted-in": True, "unflagged": False} + + +@pytest.mark.parametrize("registration_order", [("opted-in", "unflagged"), ("unflagged", "opted-in")]) +async def test_openai_messages_bridge_sends_a_breakpoint_by_the_deployments_own_flag_not_a_siblings( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, registration_order: tuple[str, ...] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _, reply = _FOUNDRY_RESPONSES_ROUTE + responses_route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=Response(200, json=dict(reply)) + ) + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "openai/qa-shared-gpt-deployment", "api_key": "openai-test-key"}, + "model_info": {"id": f"{name}-id", **_SIBLING_BREAKPOINT_FLAGS.get(name, {})}, + } + for name in registration_order + ], + num_retries=0, + ) + + async def sends_a_breakpoint(name: str) -> bool: + await router.aanthropic_messages( + model=name, + system=[{"type": "text", "text": "stable", "cache_control": {"type": "ephemeral"}}], + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + ) + return b"prompt_cache_breakpoint" in responses_route.calls.last.request.content + + assert {name: await sends_a_breakpoint(name) for name in registration_order} == { + "opted-in": True, + "unflagged": False, + } + + +_OPENAI_ROUTES: Final = MappingProxyType( + { + "acompletion": ("https://api.openai.com/v1/chat/completions", _FOUNDRY_CHAT_ROUTE[1]), + "aresponses": ("https://api.openai.com/v1/responses", _FOUNDRY_RESPONSES_ROUTE[1]), + "aanthropic_messages": ("https://api.openai.com/v1/responses", _FOUNDRY_RESPONSES_ROUTE[1]), + } +) + + +def _openai_route(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, entrypoint: str) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + url, reply = _OPENAI_ROUTES[entrypoint] + return respx_mock.post(url).mock(return_value=Response(200, json=dict(reply))) + + +@pytest.mark.parametrize("declared", [True, False]) +async def test_a_wildcard_deployments_prompt_cache_flag_leaves_each_matched_model_to_the_catalog( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, declared: bool +) -> None: + chat: Final = _openai_route(respx_mock, monkeypatch, "acompletion") + router: Final = Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_key": "openai-test-key", + "cache_control_injection_points": [{"location": "message", "role": "system"}], + }, + "model_info": {"id": "wildcard-id", "supports_prompt_cache_breakpoint": declared}, + } + ], + num_retries=0, + ) + + async def sends_prompt_cache_options(model: str) -> bool: + await _send_unmarked_system_prompt(router, model, "acompletion") + return b"prompt_cache_options" in chat.calls.last.request.content + + assert {model: await sends_prompt_cache_options(model) for model in ("openai/gpt-4.1-qa", "openai/gpt-99.0-qa")} == { + "openai/gpt-4.1-qa": False, + "openai/gpt-99.0-qa": True, + } + + +async def test_an_azure_ai_wildcard_deployments_breakpoint_flag_still_speaks_for_each_matched_model( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + path, reply = _FOUNDRY_CHAT_ROUTE + chat: Final = respx_mock.post(host="qa-wildcard.services.ai.azure.com", path=path).mock( + return_value=Response(200, json=dict(reply)) + ) + router: Final = Router( + model_list=[ + { + "model_name": "azure_ai/*", + "litellm_params": { + "model": "azure_ai/*", + "api_base": "https://qa-wildcard.services.ai.azure.com/models", + "api_key": "azure-ai-test-key", + "cache_control_injection_points": [{"location": "message", "role": "system"}], + }, + "model_info": {"id": "wildcard-id", "supports_prompt_cache_breakpoint": True}, + } + ], + num_retries=0, + ) + + await _send_unmarked_system_prompt(router, "azure_ai/qa-unlisted-alias", "acompletion") + + body: Final = chat.calls.last.request.content + assert (json.loads(body).get("prompt_cache_options"), b"prompt_cache_breakpoint" in body) == ( + {"mode": "implicit"}, + True, + ) + + +@pytest.mark.parametrize("entrypoint", ["acompletion", "aresponses", "aanthropic_messages"]) +@pytest.mark.parametrize("declared_in", ["litellm_params", "model_info"]) +async def test_a_base_model_that_takes_breakpoints_leaves_an_openai_deployments_request_as_it_was( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, entrypoint: str, declared_in: str +) -> None: + upstream: Final = _openai_route(respx_mock, monkeypatch, entrypoint) + + def deployment(name: str, base_model: Mapping[str, str]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": { + "model": "openai/gpt-4.1-qa", + "api_key": "openai-test-key", + "cache_control_injection_points": [{"location": "message", "role": "system"}], + **(base_model if declared_in == "litellm_params" else {}), + }, + "model_info": {"id": f"{name}-id", **(base_model if declared_in == "model_info" else {})}, + } + + router: Final = Router( + model_list=[deployment("plain", {}), deployment("priced-as-breakpoint-model", {"base_model": "gpt-99.0-qa"})], + num_retries=0, + ) + + async def sent_body(name: str) -> object: + await _send_unmarked_system_prompt(router, name, entrypoint) + return json.loads(upstream.calls.last.request.content) + + assert await sent_body("priced-as-breakpoint-model") == await sent_body("plain") + + +def test_a_routed_call_honors_a_register_model_breakpoint_opt_out_made_under_the_provider_prefixed_name() -> None: + from litellm.utils import catalog_with_runtime_registrations + + model: Final = next( + key + for key, entry in catalog_with_runtime_registrations().items() + if entry.get("litellm_provider") == "openai" and entry.get("supports_prompt_cache_breakpoint") is True + ) + litellm.register_model({f"openai/{model}": {"supports_prompt_cache_breakpoint": False}}) + + assert supports_openai_prompt_cache_breakpoint(model, "openai", {"id": "plain-id"}) is False + + +_APPLICATION_INFERENCE_PROFILE: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/qa" + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected"), + [ + ("us.anthropic.claude-sonnet-5-5", "bedrock", False), + ("converse/us.anthropic.claude-sonnet-5-5", "bedrock", False), + ("invoke/us.anthropic.claude-sonnet-5-5", "bedrock", False), + ("converse_like/us.anthropic.claude-sonnet-5-5", "bedrock", False), + ("openai/gpt-oss-120b", "bedrock", False), + ("moonshotai.kimi-k3", "bedrock", False), + (_APPLICATION_INFERENCE_PROFILE, "bedrock", False), + ("claude-sonnet-5-5", "anthropic", False), + ("claude-sonnet-5-5", "azure_ai", False), + ("claude-sonnet-5-5", "vertex_ai", False), + ("gemini-3-pro", "gemini", False), + ("qa-custom-model", "qa-custom-handler", False), + ("qa-shared-gpt-deployment", "azure_ai", True), + ("my-custom-deployment", "azure", True), + ("qa-mantle-model", "bedrock_mantle", True), + ("anthropic/qa-gateway-claude", "openrouter", False), + ("qa-downstream-claude", "litellm_proxy", False), + ("qa-self-hosted-model", "hosted_vllm", False), + ], +) +def test_a_declared_breakpoint_flag_picks_openais_dialect_only_where_the_request_goes_out_openai_shaped( + model: str, custom_llm_provider: str, expected: bool +) -> None: + assert ( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + model, custom_llm_provider, model_info={"supports_prompt_cache_breakpoint": True} + ) + is expected + ) + + +def test_a_gateway_deployment_declaring_the_flag_picks_openais_dialect_for_a_model_the_cost_map_lists_there() -> None: + litellm.register_model({"openrouter/qa-listed-gpt": {"litellm_provider": "openrouter", "mode": "chat"}}) + + assert [ + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + model, "openrouter", model_info={"supports_prompt_cache_breakpoint": True} + ) + for model in ("qa-listed-gpt", "qa-unlisted-gpt") + ] == [True, False] + + +def test_a_gateway_deployment_declaring_the_flag_picks_openais_dialect_when_its_model_info_lists_that_provider() -> None: + assert [ + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + "qa-self-hosted-gpt", + "hosted_vllm", + model_info={"supports_prompt_cache_breakpoint": True, "litellm_provider": listed_provider}, + ) + for listed_provider in ("hosted_vllm", "openrouter") + ] == [True, False] + +def test_a_custom_handler_provider_declaring_the_flag_keeps_the_anthropic_mark() -> None: + litellm.register_model( + {"qa-custom-handler/qa-model": {"litellm_provider": "qa-custom-handler", "mode": "chat"}} + ) + + assert ( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + "qa-model", "qa-custom-handler", model_info={"supports_prompt_cache_breakpoint": True} + ) + is False + ) + + +def test_a_routed_deployment_without_its_own_flag_reads_a_register_model_opt_in_but_not_a_siblings() -> None: + litellm.register_model( + {"azure_ai/qa-registered-gpt": {"litellm_provider": "azure_ai", "supports_prompt_cache_breakpoint": True}} + ) + Router( + model_list=[ + { + "model_name": "qa-opted-in-sibling", + "litellm_params": {"model": "azure_ai/qa-sibling-gpt", "api_base": "https://qa.services.ai.azure.com"}, + "model_info": {"id": "qa-opted-in-sibling-id", "supports_prompt_cache_breakpoint": True}, + } + ] + ) + + assert [ + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + model, "azure_ai", model_info={"id": "qa-flagless-id"} + ) + for model in ("qa-registered-gpt", "qa-sibling-gpt") + ] == [True, False] + + +@pytest.mark.parametrize( + ("entrypoint", "injection_points"), + [ + ("completion", [{"location": "message", "role": "system"}]), + ("acompletion", [{"location": "message", "role": "system"}]), + ("aresponses", [{"location": "message", "role": "system"}]), + ("aresponses-client-marked", None), + ], +) +async def test_an_azure_custom_deployment_declaring_the_breakpoint_flag_gets_openais_breakpoint( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + entrypoint: str, + injection_points: list[dict[str, str]] | None, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + azure: Final = respx_mock.post(host="qa-azure.openai.azure.com").mock( + side_effect=lambda request: Response( + 200, + json=dict( + _FOUNDRY_RESPONSES_ROUTE[1] if request.url.path.endswith("/responses") else _FOUNDRY_CHAT_ROUTE[1] + ), + ) + ) + deployment_params: Final = { + "model": "azure/my-custom-deployment", + "api_base": "https://qa-azure.openai.azure.com", + "api_key": "azure-test-key", + "api_version": "2025-04-01-preview", + } + router: Final = Router( + model_list=[ + { + "model_name": "custom-azure", + "litellm_params": ( + deployment_params + if injection_points is None + else {**deployment_params, "cache_control_injection_points": injection_points} + ), + "model_info": {"id": "custom-azure-id", "supports_prompt_cache_breakpoint": True}, + } + ], + num_retries=0, + ) + client_marked_system_prompt: Final = { + "role": "system", + "content": [{"type": "input_text", "text": "stable", "cache_control": {"type": "ephemeral"}}], + } + + match entrypoint: + case "aresponses-client-marked": + await router.aresponses( + model="custom-azure", input=[client_marked_system_prompt, {"role": "user", "content": "hi"}] + ) + case _: + await _send_unmarked_system_prompt(router, "custom-azure", entrypoint) + + body: Final = azure.calls.last.request.content + assert (b"prompt_cache_breakpoint" in body, b"cache_control" in body) == (True, False) + + +_BEDROCK_CONVERSE_REPLY: Final = MappingProxyType( + { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}, + } +) +_BEDROCK_INVOKE_REPLY: Final = MappingProxyType( + { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + } +) + + +@pytest.mark.parametrize( + ("entrypoint", "expected_system"), + [ + ("completion", [{"text": "stable"}, {"cachePoint": {"type": "default"}}]), + ("acompletion", [{"text": "stable"}, {"cachePoint": {"type": "default"}}]), + ("aanthropic_messages", [{"type": "text", "text": "stable", "cache_control": {"type": "ephemeral"}}]), + ], +) +async def test_a_bedrock_claude_deployment_declaring_the_breakpoint_flag_keeps_bedrocks_own_cache_marks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, entrypoint: str, expected_system: list[object] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + bedrock: Final = respx_mock.post(host="bedrock-runtime.us-east-1.amazonaws.com").mock( + side_effect=lambda request: Response( + 200, + json=dict(_BEDROCK_INVOKE_REPLY if request.url.path.endswith("/invoke") else _BEDROCK_CONVERSE_REPLY), + ) + ) + router: Final = Router( + model_list=[ + { + "model_name": "claude-on-bedrock", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_access_key_id": "bedrock-test-key", + "aws_secret_access_key": "bedrock-test-secret", + "aws_region_name": "us-east-1", + "cache_control_injection_points": [{"location": "message", "role": "system"}], + }, + "model_info": {"id": "claude-on-bedrock-id", "supports_prompt_cache_breakpoint": True}, + } + ], + num_retries=0, + ) + + await _send_unmarked_system_prompt(router, "claude-on-bedrock", entrypoint) + + body: Final = bedrock.calls.last.request.content + assert (json.loads(body)["system"], b"prompt_cache" in body) == (expected_system, False) + + def _search_conversation_with(follow_up: dict[str, object]) -> list[AllMessageValues]: return [ cast(AllMessageValues, {"role": "user", "content": "Weather in Paris?"}), diff --git a/tests/unit/integrations/test_deepeval.py b/tests/unit/integrations/test_deepeval.py index 6dd3aef0a03..4523a9c61d2 100644 --- a/tests/unit/integrations/test_deepeval.py +++ b/tests/unit/integrations/test_deepeval.py @@ -2,11 +2,15 @@ import os import unittest from litellm._uuid import uuid from datetime import datetime, timezone +from typing import Final from unittest.mock import MagicMock, patch from litellm.integrations.deepeval.api import Endpoints, HttpMethods from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.deepeval.types import SpanApiType, TraceSpanApiStatus +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN class TestDeepEvalLogger(unittest.TestCase): @@ -94,6 +98,37 @@ class TestDeepEvalLogger(unittest.TestCase): self.assertEqual(llm_span["inputTokenCount"], 10) self.assertEqual(llm_span["outputTokenCount"], 20) + def test_background_cost_span_keeps_its_own_uuid_beside_its_create(self) -> None: + standard_logging_object: Final = { + "id": "resp_YWJjZA", + "trace_id": self.trace_id, + "model": self.model, + "response": {"usage": {}, "choices": [{"message": {"content": "OK"}}]}, + } + create_kwargs: Final = { + "input": self.input_str, + "call_type": "aresponses", + "litellm_params": {"metadata": {}}, + "standard_logging_object": standard_logging_object, + } + cost_kwargs: Final = { + "input": self.input_str, + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + "standard_logging_object": standard_logging_object, + } + + for kwargs in (create_kwargs, cost_kwargs): + self.logger.log_success_event(kwargs, self.mock_response_obj, self.start_time, self.end_time) + + span_uuids: Final = [ + call.kwargs["body"]["llmSpans"][0]["uuid"] + for call in self.mock_api_instance.send_request.call_args_list + ] + self.assertEqual(span_uuids, ["resp_YWJjZA", f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"]) + def test_log_failure_event(self): error_message = "This is an error." kwargs = { diff --git a/tests/unit/integrations/test_dynamodb.py b/tests/unit/integrations/test_dynamodb.py new file mode 100644 index 00000000000..d139c49a49e --- /dev/null +++ b/tests/unit/integrations/test_dynamodb.py @@ -0,0 +1,52 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Final +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.integrations.dynamodb import DyanmoDBLogger +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +BACKGROUND_COST_POLL_KWARGS: Final[Mapping[str, object]] = { + "call_type": "aget_responses", + "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final[Mapping[str, object]] = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} + + +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), +) +@pytest.mark.usefixtures("sdk_only_install") +def test_background_cost_event_keeps_its_own_item_beside_its_create( + monkeypatch: pytest.MonkeyPatch, cost_kwargs: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "dynamodb_table_name", "litellm-logs") + dynamodb: Final = MagicMock() + logger: Final = DyanmoDBLogger(dynamodb=dynamodb) + response_obj: Final = {"id": "resp_YWJjZA", "usage": {"total_tokens": 3}} + now: Final = datetime.now(timezone.utc) + create_kwargs: Final[Mapping[str, object]] = {"call_type": "aresponses", "litellm_params": {"metadata": {}}} + + for kwargs in (create_kwargs, cost_kwargs): + logger.log_event(kwargs, response_obj, now, now, print_verbose=print) + + dynamodb.Table.assert_called_with("litellm-logs") + written_ids: Final = [call.kwargs["Item"]["id"] for call in dynamodb.Table.return_value.put_item.call_args_list] + assert written_ids == ["resp_YWJjZA", f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"] diff --git a/tests/unit/integrations/test_openmeter.py b/tests/unit/integrations/test_openmeter.py index 2d09e1572db..fc4e2fcd7a3 100644 --- a/tests/unit/integrations/test_openmeter.py +++ b/tests/unit/integrations/test_openmeter.py @@ -1,11 +1,30 @@ import json import os +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.openmeter import OpenMeterLogger +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +BACKGROUND_COST_POLL_KWARGS: Final[Mapping[str, object]] = { + "call_type": "aget_responses", + "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final[Mapping[str, object]] = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} class TestOpenMeterIntegration: @@ -36,6 +55,31 @@ class TestOpenMeterIntegration: with pytest.raises(Exception, match=r"Missing keys.*OPENMETER_API_KEY"): OpenMeterLogger() + @pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), + ) + @pytest.mark.usefixtures("sdk_only_install") + def test_background_cost_event_gets_its_own_event_id_beside_its_create( + self, cost_kwargs: Mapping[str, object] + ) -> None: + logger: Final = OpenMeterLogger() + response_obj: Final = {"id": "resp_YWJjZA"} + create_kwargs: Final[Mapping[str, object]] = { + "user": "user-1", + "call_type": "aresponses", + "litellm_params": {"metadata": {}}, + } + + create_event: Final = logger._common_logic(create_kwargs, response_obj) + cost_event: Final = logger._common_logic({**cost_kwargs, "user": "user-1", "response_cost": 0.42}, response_obj) + + assert create_event["id"] == "resp_YWJjZA" + assert cost_event["id"] == f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}" + def test_common_logic_with_string_user(self): """Test that _common_logic correctly handles string user parameter""" logger = OpenMeterLogger() diff --git a/tests/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index 878f19c66f8..04eb7c3a5a4 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -5,6 +5,7 @@ Unit tests for prometheus metric labels configuration from typing import Final, Literal import pytest +from pydantic import JsonValue import litellm from litellm.types.integrations.prometheus import ( @@ -802,6 +803,47 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric() _clear_prometheus_registry() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata_id", "request_id", "logged_id", "expected"), + ( + (None, None, "logged-deployment", "logged-deployment"), + ("", None, "logged-deployment", "logged-deployment"), + ("metadata-deployment", "request-deployment", "logged-deployment", "metadata-deployment"), + (None, "request-deployment", "logged-deployment", "request-deployment"), + (None, None, None, "None"), + ), +) +async def test_failure_counters_keep_logged_deployment_attribution( + metadata_id: str | None, request_id: str | None, logged_id: str | None, expected: str +) -> None: + from fastapi import HTTPException + + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import UserAPIKeyAuth + + request_data: Final[dict[str, JsonValue]] = { + "model": "metric-test-alias", + "metadata": {"model_info": {"id": metadata_id}}, + "model_info": {"id": request_id}, + "standard_logging_object": {"model_id": logged_id, "custom_llm_provider": "openai"}, + } + _clear_prometheus_registry() + try: + await PrometheusLogger().async_post_call_failure_hook( + request_data=request_data, + original_exception=HTTPException(status_code=400, detail="guardrail rejected"), + user_api_key_dict=UserAPIKeyAuth(token="metric-test-token"), + ) + for metric_name in ("litellm_proxy_failed_requests_metric_total", "litellm_proxy_total_requests_metric_total"): + samples: Final = _collected_samples(metric_name) + assert [(sample.labels["model_id"], sample.value) for sample in samples] == [ + (expected, 1.0) + ] + finally: + _clear_prometheus_registry() + + @pytest.mark.asyncio async def test_failure_hook_labels_a_cli_session_with_the_per_user_alias_not_the_login_token(): from litellm.integrations.prometheus import PrometheusLogger diff --git a/tests/unit/litellm_core_utils/audio_utils/test_transcription_streaming.py b/tests/unit/litellm_core_utils/audio_utils/test_transcription_streaming.py new file mode 100644 index 00000000000..5296ea3b15a --- /dev/null +++ b/tests/unit/litellm_core_utils/audio_utils/test_transcription_streaming.py @@ -0,0 +1,319 @@ +import datetime +import io +import math +import wave +from collections.abc import AsyncIterator, Iterator +from typing import Final + +import anyio +import anyio.lowlevel +import httpx +import pytest +import respx +from openai import AsyncOpenAI, OpenAI + +import litellm +from litellm import Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( + litellm_exception_mapper, + wrap_transcription_stream, +) +from litellm.types.utils import TranscriptionResponse + +_STARTED_AT: Final = datetime.datetime(2026, 10, 28, 12, 0, 0) # noqa: DTZ001 # callback timestamps use the legacy naive contract + + +class _FirstDeltaThenStall(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b'data: {"type":"transcript.text.delta","delta":"Good"}\n\n' + await anyio.sleep_forever() + + +class _FirstDeltaThenFails(httpx.AsyncByteStream): + def __init__(self, failure: httpx.TransportError) -> None: + self._failure = failure + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b'data: {"type":"transcript.text.delta","delta":"Good"}\n\n' + raise self._failure + + +def _transcription_router(failing_api_base: str, healthy_api_base: str) -> Router: + return Router( + model_list=[ + { + "model_name": "transcriber", + "litellm_params": { + "model": "openai/gpt-transcribe", + "api_key": "sk-test", + "api_base": api_base, + "order": order, + }, + } + for order, api_base in ((1, failing_api_base), (2, healthy_api_base)) + ], + allowed_fails=0, + cooldown_time=60, + ) + + +class _RecordingLogging: + def __init__(self) -> None: + self.logged: tuple[TranscriptionResponse, ...] = () + + def success_handler( + self, result: TranscriptionResponse, start_time: datetime.datetime, end_time: datetime.datetime + ) -> None: + return None + + async def async_success_handler( + self, result: TranscriptionResponse, start_time: datetime.datetime, end_time: datetime.datetime + ) -> None: + await anyio.lowlevel.checkpoint() + self.logged = (*self.logged, result) + + def handle_sync_success_callbacks_for_async_calls( + self, result: TranscriptionResponse, start_time: datetime.datetime, end_time: datetime.datetime + ) -> None: + return None + + def failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime, + end_time: datetime.datetime, + *, + for_async_request: bool = False, + ) -> None: + return None + + async def async_failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + return None + + +@pytest.mark.asyncio +async def test_a_stream_the_client_drops_mid_transcript_is_still_logged_as_a_success() -> None: + upstream: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_FirstDeltaThenStall() + ) + ) + ), + ) + stream: Final = await upstream.audio.transcriptions.create( + model="gpt-transcribe", file=("sample.wav", b"audio"), stream=True + ) + logging_obj: Final = _RecordingLogging() + wrapped: Final = wrap_transcription_stream(stream, logging_obj, _STARTED_AT, litellm_exception_mapper({})) + + with anyio.CancelScope() as client_connection: + async for _ in wrapped: + client_connection.cancel() + + assert [result.text for result in logging_obj.logged] == ["Good"] + + +@pytest.mark.asyncio +async def test_a_deployment_whose_stream_times_out_mid_transcript_goes_into_cooldown( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _transcription_router("https://stalls.example/v1", "https://healthy.example/v1") + + with respx.mock(assert_all_called=False) as upstream: + upstream.post("https://stalls.example/v1/audio/transcriptions").mock( + side_effect=lambda request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_FirstDeltaThenFails(httpx.ReadTimeout("read timed out")), + ) + ) + upstream.post("https://healthy.example/v1/audio/transcriptions").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=b'data: {"type":"transcript.text.done","text":"Good morning"}\n\n', + ) + ) + stalled: Final = await router.atranscription(model="transcriber", file=("sample.wav", b"audio"), stream=True) + with pytest.raises(litellm.Timeout): + async for _ in stalled: + pass + await anyio.sleep(0.1) + retried: Final = await router.atranscription(model="transcriber", file=("sample.wav", b"audio"), stream=True) + events: Final = [event async for event in retried] + + assert [getattr(event, "text", None) for event in events] == ["Good morning"] + + +@pytest.mark.asyncio +async def test_a_stream_the_upstream_drops_mid_transcript_raises_a_litellm_error_to_the_caller() -> None: + upstream: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_FirstDeltaThenFails(httpx.RemoteProtocolError("peer closed connection")), + ) + ) + ), + ) + stream: Final = await litellm.atranscription( + model="openai/gpt-transcribe", file=("sample.wav", b"audio"), stream=True, client=upstream + ) + + with pytest.raises(litellm.APIConnectionError, match="peer closed connection"): + async for _ in stream: + pass + + +class _FailureHookCounter(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: tuple[str, ...] = () + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.failures = (*self.failures, "sync_failure") + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.failures = (*self.failures, "async_failure") + + +@pytest.mark.asyncio +async def test_an_async_stream_that_fails_mid_transcript_is_logged_as_one_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = _FailureHookCounter() + monkeypatch.setattr(litellm, "callbacks", [counter]) + upstream: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_FirstDeltaThenFails(httpx.ReadError("connection reset")), + ) + ) + ), + ) + stream: Final = await litellm.atranscription( + model="openai/gpt-transcribe", file=("sample.wav", b"audio"), stream=True, client=upstream + ) + + with pytest.raises(litellm.APIConnectionError): + async for _ in stream: + pass + + assert counter.failures == ("async_failure",) + + +class _FirstDeltaThenResetSync(httpx.SyncByteStream): + def __iter__(self) -> Iterator[bytes]: + yield b'data: {"type":"transcript.text.delta","delta":"Good"}\n\n' + raise httpx.ReadError("connection reset") + + +class _FailureCostRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.costs: tuple[object, ...] = () + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.costs = (*self.costs, kwargs.get("response_cost") if isinstance(kwargs, dict) else None) + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.costs = (*self.costs, kwargs.get("response_cost") if isinstance(kwargs, dict) else None) + + +def _silent_wav(seconds: int) -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(16_000) + wav.writeframes(b"\x00\x00" * 16_000 * seconds) + return buffer.getvalue() + + +def _reset_mid_transcript(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_FirstDeltaThenFails(httpx.ReadError("connection reset")), + ) + + +def _reset_mid_transcript_sync(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_FirstDeltaThenResetSync()) + + +def _per_second_rate(model: str) -> float: + rate: Final = litellm.model_cost[model]["input_cost_per_second"] + assert isinstance(rate, float) + return rate + + +@pytest.mark.asyncio +async def test_an_async_stream_that_fails_mid_transcript_bills_the_uploaded_seconds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder: Final = _FailureCostRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + upstream: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_reset_mid_transcript)), + ) + stream: Final = await litellm.atranscription( + model="openai/gpt-transcribe", file=("sample.wav", _silent_wav(4)), stream=True, client=upstream + ) + + with pytest.raises(litellm.APIConnectionError): + async for _ in stream: + pass + + assert len(recorder.costs) == 1 + assert isinstance(recorder.costs[0], float) + assert math.isclose(recorder.costs[0], 4 * _per_second_rate("gpt-transcribe")) + + +def test_a_sync_stream_that_fails_mid_transcript_bills_the_uploaded_seconds(monkeypatch: pytest.MonkeyPatch) -> None: + recorder: Final = _FailureCostRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + upstream: Final = OpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.Client(transport=httpx.MockTransport(_reset_mid_transcript_sync)), + ) + stream: Final = litellm.transcription( + model="openai/gpt-transcribe", file=("sample.wav", _silent_wav(4)), stream=True, client=upstream + ) + + with pytest.raises(litellm.APIConnectionError): + for _ in stream: + pass + + assert len(recorder.costs) == 1 + assert isinstance(recorder.costs[0], float) + assert math.isclose(recorder.costs[0], 4 * _per_second_rate("gpt-transcribe")) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 21b527250c4..0a08142c66f 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -29,6 +29,7 @@ from litellm.llms.gemini.image_generation.cost_calculator import ( from litellm.llms.vertex_ai.image_generation.cost_calculator import ( cost_calculator as vertex_image_generation_cost_calculator, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.utils import ( CacheCreationTokenDetails, CompletionTokensDetailsWrapper, @@ -42,6 +43,62 @@ from litellm.types.utils import ( ) +def test_realtime_cached_modality_breakdown_matches_prompt_cost(_local_model_cost_map): + model: Final = "gpt-realtime-2.1-mini" + rates: Final = litellm.model_cost[model] + usage: Final = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=400, + audio_tokens=400, + image_tokens=200, + cached_tokens=300, + cached_tokens_details=CachedTokensDetails(text_tokens=100, audio_tokens=150, image_tokens=50), + ), + ) + + prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown: Final = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + cached_cost: Final = ( + 100 * rates["cache_read_input_token_cost"] + + 150 * rates["cache_read_input_audio_token_cost"] + + 50 * rates["cache_read_input_image_token_cost"] + ) + uncached_cost: Final = ( + 300 * rates["input_cost_per_token"] + + 250 * rates["input_cost_per_audio_token"] + + 150 * rates["input_cost_per_image_token"] + ) + + assert breakdown.cache_read_cost == pytest.approx(cached_cost) + assert prompt_cost == pytest.approx(uncached_cost + cached_cost) + + +def test_image_model_bills_cached_image_tokens_at_its_cached_image_rate(_local_model_cost_map): + model: Final = "gpt-image-2" + rates: Final = litellm.model_cost[model] + usage: Final = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=200, + image_tokens=800, + cached_tokens=600, + cached_tokens_details=CachedTokensDetails(text_tokens=100, image_tokens=500), + ), + ) + + prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + cached_cost: Final = 100 * rates["cache_read_input_token_cost"] + 500 * rates["cache_read_input_image_token_cost"] + uncached_cost: Final = 100 * rates["input_cost_per_token"] + 300 * rates["input_cost_per_image_token"] + + assert rates["cache_read_input_image_token_cost"] != rates["cache_read_input_token_cost"] + assert prompt_cost == pytest.approx(uncached_cost + cached_cost) + + @pytest.fixture def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") @@ -2234,6 +2291,10 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): prompt_tokens_details: PromptTokensDetailsResult = { "cache_hit_tokens": 0, "cache_hit_audio_tokens": 0, + "cached_text_tokens": 0, + "cached_audio_tokens": 0, + "cached_image_tokens": 0, + "has_cached_tokens_details": False, "cache_creation_tokens": 0, "cache_creation_token_details": CacheCreationTokenDetails( ephemeral_5m_input_tokens=100, diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index d0f5bfc6b9f..3e197740e98 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -3089,6 +3089,73 @@ def test_convert_to_model_response_object_function_output(): assert result._hidden_params is not None +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": 151645}, {"stop_reason": None, "token_ids": None}], + ids=["stop-string-with-token-ids", "stop-token", "natural-stop"], +) +def test_vllm_choice_fields_stay_on_the_choice_through_a_cache_round_trip( + vllm_choice_fields: dict[str, object], +) -> None: + upstream: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + **vllm_choice_fields, + } + ], + } + response: Final = convert_to_model_response_object(response_object=upstream, model_response_object=ModelResponse()) + assert isinstance(response, ModelResponse) + cache_hit: Final = convert_to_model_response_object( + response_object=json.loads(response.model_dump_json()), model_response_object=ModelResponse() + ) + assert isinstance(cache_hit, ModelResponse) + + dumped_choices: Final = [json.loads(rebuilt.model_dump_json())["choices"][0] for rebuilt in (response, cache_hit)] + assert [{field: choice.get(field, "missing") for field in vllm_choice_fields} for choice in dumped_choices] == [ + vllm_choice_fields, + vllm_choice_fields, + ] + + +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": None, "token_ids": None}], + ids=["stop-string-with-token-ids", "natural-stop"], +) +def test_vllm_choice_fields_cached_only_under_provider_specific_fields_come_back_on_the_choice( + vllm_choice_fields: dict[str, object], +) -> None: + cached_before_upgrade: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + "provider_specific_fields": vllm_choice_fields, + } + ], + } + cache_hit: Final = convert_to_model_response_object( + response_object=cached_before_upgrade, model_response_object=ModelResponse() + ) + assert isinstance(cache_hit, ModelResponse) + + dumped_choice: Final = json.loads(cache_hit.model_dump_json())["choices"][0] + assert {field: dumped_choice.get(field, "missing") for field in vllm_choice_fields} == vllm_choice_fields + + @pytest.mark.parametrize( "reported_cost", [ diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py index 51249e01d4d..a8162860684 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py @@ -9,7 +9,10 @@ _BARE_FUNCTION: Final = {"name": "g", "parameters": {}} # Tool schemas as rendered by render_tools in encoding/encoding_dsv4.py (deepseek-ai/DeepSeek-V4-Flash) and by -# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17 +# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17, and by chat_template.jinja +# (deepseek-ai/DeepSeek-V4.1-Flash) read on 2026-10-29. Grok's template is closed, so its two renderings are the ones +# whose token counts left a constant preamble across seven tool shapes billed by Azure AI Foundry +# grok-4-1-fast-reasoning and xAI grok-4.3 on 2026-10-29 @pytest.mark.parametrize( ("model", "rendered"), [ @@ -17,6 +20,18 @@ _BARE_FUNCTION: Final = {"name": "g", "parameters": {}} "DeepSeek-V4-Flash-0731", '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', ), + ( + "DeepSeek-V4.1-Flash", + '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', + ), + ( + "grok-4-1-fast-reasoning", + '{"name":"f","description":"é","parameters":{}}\n{"name":"g","parameters":{}}', + ), + ( + "grok-4.3", + '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', + ), ( "databricks-qwen35-122b-a10b", '\n{"type": "function", "function": {"name": "f", "description": "é", "parameters": {}}}' diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1ef97370b2b..146bb08a4fd 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -2343,7 +2343,8 @@ _MALFORMED_PROMPT_CACHE_BREAKPOINTS: Final = ( "", {}, {"mode": "bogus"}, - {"mode": "explicit", "ttl": "1h"}, + {"mode": "bogus", "ttl": "30m"}, + {1: "explicit"}, ) _MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS: Final = ( "string", @@ -2352,7 +2353,8 @@ _MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS: Final = ( "empty-string", "empty-object", "bogus-mode", - "bad-ttl", + "bogus-mode-with-ttl", + "non-string-key", ) _VALID_PROMPT_CACHE_BREAKPOINTS: Final = ({"mode": "explicit"}, {"mode": "explicit", "ttl": "30m"}) @@ -2382,6 +2384,18 @@ def test_prompt_cache_breakpoint_for_wire_strips_an_unknown_key_under_drop_param assert prompt_cache_breakpoint_for_wire({"mode": "explicit", "note": "kept"}, True) == {"mode": "explicit"} +@pytest.mark.parametrize( + "marker", + ({"mode": "explicit", "ttl": "1h"}, {"mode": "explicit", "ttl": "1h", "note": "kept"}), + ids=("bad-ttl", "bad-ttl-and-unknown-key"), +) +def test_prompt_cache_breakpoint_for_wire_drops_an_unsupported_ttl_and_keeps_the_mode_under_drop_params( + marker: dict[str, str], +) -> None: + assert prompt_cache_breakpoint_for_wire(marker, True) == {"mode": "explicit"} + assert prompt_cache_breakpoint_for_wire(marker, False) is marker + + def test_cache_control_is_replaced_by_an_explicit_prompt_cache_breakpoint_on_the_same_block() -> None: marked: Final = {"type": "text", "text": "rules", "cache_control": {"type": "ephemeral", "ttl": "1h"}} diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 8a755d47817..d65f094971a 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1298,6 +1298,71 @@ def test_anthropic_messages_pt_uses_pdf_format_for_generic_mime_data_uri(llm_pro } +@pytest.mark.parametrize( + ("image_url", "expected_type"), + ( + ({"url": "https://example.com/download?id=7", "format": "application/pdf"}, "document"), + ({"url": "https://example.com/download?id=7", "format": "Application/PDF; charset=binary"}, "document"), + ({"url": "https://example.com/files/report.pdf", "format": "image/png"}, "image"), + ({"url": "https://example.com/pdf/photo.png"}, "image"), + ({"url": "https://example.com/pdf/photo.png", "format": 5}, "image"), + ({"url": "https://example.com/pdf/photo.png", "format": ["application/pdf"]}, "image"), + ), +) +def test_anthropic_messages_pt_picks_block_type_for_https_image_url( + image_url: dict[str, object], expected_type: str +) -> None: + result: Final = anthropic_messages_pt( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "image_url", "image_url": image_url}, + ], + } + ], + model="claude-opus-5-5", + llm_provider="anthropic", + ) + + assert result[0]["content"][1] == { + "type": expected_type, + "source": {"type": "url", "url": image_url["url"]}, + } + + +@pytest.mark.parametrize( + "image_url", + ( + {"url": "data:Application/PDF;base64,JVBERi0xLjQK"}, + {"url": "data:application/x-pdf;base64,JVBERi0xLjQK"}, + {"url": "data:application/octet-stream;base64,JVBERi0xLjQK", "format": "application/pdf; charset=binary"}, + ), +) +def test_anthropic_messages_pt_normalizes_a_pdf_media_type_with_case_or_parameters( + image_url: dict[str, str], +) -> None: + result: Final = anthropic_messages_pt( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "image_url", "image_url": image_url}, + ], + } + ], + model="claude-opus-5-5", + llm_provider="anthropic", + ) + + assert result[0]["content"][1] == { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + } + + def test_anthropic_messages_pt_keeps_png_image_url_as_image() -> None: image_data = "iVBORw0KGgo=" result: Final = anthropic_messages_pt( diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py new file mode 100644 index 00000000000..6ebbdcf73bf --- /dev/null +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -0,0 +1,100 @@ +import copy +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + ChatAllowedToolsToolChoice, + is_allowed_tools_tool_choice, + parse_chat_allowed_tools_tool_choice, + responses_allowed_tools_tool_choice, +) + +_CHAT_LAYOUT: Final[dict[str, object]] = { + "type": "allowed_tools", + "allowed_tools": { + "mode": "auto", + "tools": [ + {"type": "function", "function": {"name": "get_time"}}, + {"type": "custom", "custom": {"name": "apply_patch", "x-vendor": 1}, "x-trace": "a"}, + {"type": "mcp", "server_label": "docs", "server_url": "https://docs.example"}, + {"type": "function", "name": "already_flat"}, + ], + }, +} + + +def _parsed(tool_choice: Mapping[str, object]) -> ChatAllowedToolsToolChoice: + parsed: Final = parse_chat_allowed_tools_tool_choice(tool_choice) + assert parsed is not None + return parsed + + +def test_responses_layout_lifts_mode_and_tools_and_flattens_names_in_order(): + assert responses_allowed_tools_tool_choice(_parsed(_CHAT_LAYOUT)) == { + "type": "allowed_tools", + "mode": "auto", + "tools": [ + {"type": "function", "name": "get_time"}, + {"type": "custom", "name": "apply_patch", "x-vendor": 1, "x-trace": "a"}, + {"type": "mcp", "server_label": "docs", "server_url": "https://docs.example"}, + {"type": "function", "name": "already_flat"}, + ], + } + + +def test_responses_layout_leaves_the_caller_dict_untouched(): + original: Final = copy.deepcopy(_CHAT_LAYOUT) + flat: Final = responses_allowed_tools_tool_choice(_parsed(_CHAT_LAYOUT)) + assert flat is not None + next(iter(flat["tools"]))["name"] = "mutated" + assert _CHAT_LAYOUT == original + + +@pytest.mark.parametrize( + "tool", + [ + {"type": "function"}, + {"type": "function", "name": ""}, + {"type": "function", "function": {}}, + {"type": "function", "function": "get_time"}, + {"type": "custom", "custom": {"name": ""}}, + {"type": "function", "name": "get_time", "function": {"name": "get_weather"}}, + {"type": "function", "function": {"name": "get_time", "type": "custom"}}, + ], +) +def test_responses_layout_refuses_to_rewrite_a_malformed_or_conflicting_entry(tool: dict[str, object]): + tool_choice: Final[dict[str, object]] = { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [tool]}, + } + assert responses_allowed_tools_tool_choice(_parsed(tool_choice)) is None + + +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": []}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "any", "tools": [{"type": "function"}]}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}, "extra": 1}, + {"type": "function", "function": {"name": "get_time"}}, + ], +) +def test_parse_rejects_anything_but_the_chat_layout(tool_choice: dict[str, object]): + assert parse_chat_allowed_tools_tool_choice(tool_choice) is None + + +@pytest.mark.parametrize( + ("tool_choice", "expected"), + [ + (_CHAT_LAYOUT, True), + ({"type": "allowed_tools"}, True), + ({"type": "function"}, False), + ("auto", False), + (None, False), + ], +) +def test_is_allowed_tools_tool_choice(tool_choice: Mapping[str, object] | str | None, expected: bool): + assert is_allowed_tools_tool_choice(tool_choice) is expected diff --git a/tests/unit/litellm_core_utils/test_decisions_pricing.py b/tests/unit/litellm_core_utils/test_decisions_pricing.py index e9814b980d1..d9603230319 100644 --- a/tests/unit/litellm_core_utils/test_decisions_pricing.py +++ b/tests/unit/litellm_core_utils/test_decisions_pricing.py @@ -62,6 +62,66 @@ def test_apply_decisions_pricing_leaves_rows_without_decisions_fields_unchanged( assert apply_decisions_pricing(plain_row) == plain_row +@pytest.mark.parametrize("rate", [0.0, 0.002]) +def test_flat_decisions_rate_overrides_catalog_without_changing_chat( + monkeypatch: pytest.MonkeyPatch, rate: float +) -> None: + model: Final = "parallel_ai/decision" + monkeypatch.setitem( + litellm.model_cost, + model, + {**_ROW, "litellm_provider": "parallel_ai", "input_cost_per_query_decisions": rate}, + ) + assert cost_calculator.cost_per_token(model=model, prompt_tokens=100, call_type="decisions") == (rate, 0.0) + assert cost_calculator.cost_per_token(model=model, prompt_tokens=100, call_type="completion") == pytest.approx( + (100 * 1e-07, 0.0) + ) + + +@pytest.mark.parametrize("query_field", ["input_cost_per_query", "input_cost_per_query_decisions"]) +def test_query_price_on_a_provider_that_bills_tokens_stays_unbilled( + monkeypatch: pytest.MonkeyPatch, query_field: str +) -> None: + model: Final = "fireworks_ai/accounts/example/models/query-priced" + monkeypatch.setitem( + litellm.model_cost, + model, + {query_field: 0.5, "input_cost_per_token": 0.0, "output_cost_per_token": 0.0} + | {"litellm_provider": "fireworks_ai", "mode": "evaluation"}, + ) + assert cost_calculator.cost_per_token(model=model, prompt_tokens=100, call_type="decisions") == (0.0, 0.0) + + +@pytest.mark.parametrize("call_type", ["adecisions", "asystemone"]) +@pytest.mark.parametrize( + ("token_rates", "expected_costs"), + [ + ({"input_cost_per_token": 1e-05, "output_cost_per_token": 0.0}, (100 * 1e-05, 0.0)), + ({"input_cost_per_token": 0.0, "output_cost_per_token": 1e-05}, (0.0, 10 * 1e-05)), + ], +) +def test_plain_query_price_leaves_a_token_priced_row_billing_tokens( + monkeypatch: pytest.MonkeyPatch, + call_type: str, + token_rates: dict[str, float], + expected_costs: tuple[float, float], +) -> None: + model: Final = "fireworks_ai/accounts/example/models/priced-both" + monkeypatch.setitem( + litellm.model_cost, + model, + {"input_cost_per_query": 0.5, **token_rates, "litellm_provider": "fireworks_ai", "mode": "evaluation"}, + ) + + assert cost_calculator.cost_per_token( + model=model, + custom_llm_provider="fireworks_ai", + prompt_tokens=100, + completion_tokens=10, + call_type=cast(CallTypesLiteral, call_type), + ) == pytest.approx(expected_costs) + + _LUNA_USAGE: Final = Usage( prompt_tokens=100, completion_tokens=10, diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 2b1e4a5eed3..1b00a32d31d 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -603,6 +603,26 @@ def test_anthropic_context_window_error_mapping(error_message): assert excinfo.value.llm_provider == "anthropic" +@pytest.mark.parametrize("upstream_status_code", [400, 500]) +def test_vertex_ai_anthropic_prompt_too_long_maps_to_context_window_exceeded(upstream_status_code: int) -> None: + original_exception: Final = VertexAIError( + status_code=upstream_status_code, + message='{"type":"error","error":{"type":"invalid_request_error",' + '"message":"prompt is too long: 210123 tokens > 200000 maximum"}}', + ) + + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model="claude-opus-5-5", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 400 + assert excinfo.value.llm_provider == "vertex_ai" + assert excinfo.value.model == "claude-opus-5-5" + + # Test cases for Vertex AI RateLimitError mapping # As per https://github.com/BerriAI/litellm/issues/16189 vertex_rate_limit_test_cases = [ diff --git a/tests/unit/litellm_core_utils/test_image_handling.py b/tests/unit/litellm_core_utils/test_image_handling.py index d1bebde9917..2d4cd4b9a6d 100644 --- a/tests/unit/litellm_core_utils/test_image_handling.py +++ b/tests/unit/litellm_core_utils/test_image_handling.py @@ -1,23 +1,36 @@ import asyncio import base64 import copy +import threading import time import uuid +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final +from unittest.mock import MagicMock import pytest from httpx import Request, Response import litellm +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.image_handling import ( MAX_CONCURRENT_REMOTE_MEDIA_FETCHES, RemoteMedia, async_convert_url_to_base64, + async_format_pdf_links, async_inline_remote_media, + async_serves_pdf, convert_url_to_base64, + format_pdf_links, inline_remote_media, + is_unformatted_pdf_link_image, + serves_pdf, ) from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @pytest.fixture(autouse=True) @@ -695,3 +708,295 @@ async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}]) + + +PDF_BODY: Final = b"%PDF-1.7\n" + b"0" * 4096 +PNG_BODY: Final = b"\x89PNG\r\n\x1a\n" + b"0" * 64 + + +@dataclass(frozen=True, slots=True) +class _Served: + status: int = 200 + content_type: str | None = "application/octet-stream" + body: bytes = PDF_BODY + honors_range: bool = True + barrier: threading.Barrier | None = None + declared_length: int | None = None + hold_open_until: threading.Event | None = None + respond_after: threading.Event | None = None + drip_every: float | None = None + + +@dataclass(frozen=True, slots=True) +class _ProbeServer: + base_url: str + routes: dict[str, _Served] + seen: list[tuple[str, str | None]] + + def serve(self, path: str, served: _Served) -> str: + self.routes[path] = served + return f"{self.base_url}{path}" + + +@pytest.fixture +def probe_server(monkeypatch): + monkeypatch.setattr(litellm, "user_url_validation", False) + monkeypatch.setattr(image_handling, "safe_get", url_utils.safe_get) + monkeypatch.setattr(litellm, "module_level_client", HTTPHandler()) + monkeypatch.setattr(litellm, "module_level_aclient", AsyncHTTPHandler()) + routes: Final[dict[str, _Served]] = {} + seen: Final[list[tuple[str, str | None]]] = [] + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + served = routes[self.path] + requested_range = self.headers.get("Range") + seen.append((self.path, requested_range)) + if served.barrier is not None: + served.barrier.wait() + if served.respond_after is not None: + served.respond_after.wait(timeout=30) + ranged = served.honors_range and requested_range == "bytes=0-1023" + body = served.body[:1024] if ranged else served.body + self.send_response(206 if ranged and served.status == 200 else served.status) + if served.content_type is not None: + self.send_header("Content-Type", served.content_type) + self.send_header("Content-Length", str(served.declared_length or len(body))) + self.end_headers() + if served.drip_every is None: + self.wfile.write(body) + else: + for byte in body: + self.wfile.write(bytes([byte])) + self.wfile.flush() + time.sleep(served.drip_every) + if served.hold_open_until is not None: + served.hold_open_until.wait(timeout=30) + + def log_message(self, *args): + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield _ProbeServer(f"http://127.0.0.1:{server.server_port}", routes, seen) + server.shutdown() + server.server_close() + + +async def _serves_pdf(mode: str, url: str, **kwargs) -> bool: + return serves_pdf(url, **kwargs) if mode == "sync" else await async_serves_pdf(url, **kwargs) + + +async def _format_pdf_links(mode, messages): + if mode == "sync": + return format_pdf_links(messages, should_probe=is_unformatted_pdf_link_image) + return await async_format_pdf_links(messages, should_probe=is_unformatted_pdf_link_image) + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_asks_for_only_the_leading_kilobyte(probe_server, mode): + url = probe_server.serve("/files/report.pdf", _Served()) + + assert await _serves_pdf(mode, url) is True + assert probe_server.seen == [("/files/report.pdf", "bytes=0-1023")] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize( + ("served", "expected"), + ( + (_Served(content_type="application/pdf", body=PNG_BODY), True), + (_Served(content_type="Application/X-PDF; charset=binary", body=b""), True), + (_Served(content_type="binary/octet-stream", body=PDF_BODY), True), + (_Served(content_type=None, body=PDF_BODY), True), + (_Served(content_type="application/octet-stream", body=b"\n" * 1000 + PDF_BODY), True), + (_Served(content_type="application/octet-stream", body=b"\n" * 1024 + PDF_BODY), False), + (_Served(content_type="binary/octet-stream", body=PNG_BODY), False), + (_Served(content_type="image/png", body=PDF_BODY), False), + (_Served(content_type="text/html", body=b"%PDF-"), False), + (_Served(status=404, content_type="application/pdf"), False), + (_Served(honors_range=False), True), + ), + ids=( + "served-pdf", + "served-x-pdf", + "generic-type-pdf-bytes", + "no-type-pdf-bytes", + "signature-inside-the-first-kilobyte", + "signature-past-the-first-kilobyte", + "generic-type-png-bytes", + "served-png", + "served-html", + "not-found", + "range-ignored", + ), +) +async def test_serves_pdf_trusts_a_specific_served_type_and_sniffs_a_generic_one(probe_server, mode, served, expected): + url = probe_server.serve(f"/{uuid.uuid4()}/report.pdf", served) + + assert await _serves_pdf(mode, url) is expected + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_makes_no_request_when_url_downloads_are_disabled(probe_server, monkeypatch, mode): + monkeypatch.setattr(image_handling, "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 0) + url = probe_server.serve("/report.pdf", _Served(content_type="application/pdf")) + + assert await _serves_pdf(mode, url) is False + assert probe_server.seen == [] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_is_false_when_the_host_refuses_the_connection(probe_server, mode): + assert await _serves_pdf(mode, "http://127.0.0.1:1/report.pdf") is False + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_is_false_without_a_request_when_the_ssrf_check_blocks_the_link(probe_server, monkeypatch, mode): + monkeypatch.setattr(litellm, "user_url_validation", True) + url = probe_server.serve("/report.pdf", _Served(content_type="application/pdf")) + + assert await _serves_pdf(mode, url) is False + assert probe_server.seen == [] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_stops_reading_after_the_leading_kilobyte_when_the_host_ignores_range(probe_server, mode): + rest_of_body_sent = threading.Event() + url = probe_server.serve( + "/report.pdf", _Served(honors_range=False, declared_length=1024**3, hold_open_until=rest_of_body_sent) + ) + + try: + assert await _serves_pdf(mode, url) is True + finally: + rest_of_body_sent.set() + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_gives_up_within_its_timeout_when_the_host_sends_one_byte_at_a_time(probe_server, mode): + url = probe_server.serve("/report.pdf", _Served(body=b"0" * 1024, honors_range=False, drip_every=0.02)) + started = time.monotonic() + + assert await _serves_pdf(mode, url, timeout=0.5) is False + assert time.monotonic() - started < 3 + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_reuses_a_link_verdict_instead_of_probing_the_link_again(probe_server, mode): + pdf_link = probe_server.serve("/a/report.pdf", _Served()) + png_link = probe_server.serve("/a/page.pdf", _Served(content_type="image/png", body=PNG_BODY)) + + verdicts = [await _serves_pdf(mode, url) for url in (pdf_link, png_link, pdf_link, png_link)] + + assert verdicts == [True, False, True, False] + assert [path for path, _ in probe_server.seen] == ["/a/report.pdf", "/a/page.pdf"] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_remembers_a_timed_out_probe_so_the_next_turn_does_not_wait_again(probe_server, mode): + host_answers = threading.Event() + url = probe_server.serve("/report.pdf", _Served(respond_after=host_answers)) + + try: + first = await _serves_pdf(mode, url, timeout=0.3) + started = time.monotonic() + second = await _serves_pdf(mode, url, timeout=0.3) + second_took = time.monotonic() - started + finally: + host_answers.set() + + assert (first, second) == (False, False) + assert second_took < 0.1 + assert probe_server.seen == [("/report.pdf", "bytes=0-1023")] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("status", (404, 503)) +async def test_serves_pdf_probes_a_link_again_a_minute_after_its_host_answered_an_error( + probe_server, monkeypatch, mode, status +): + clock: Final = MagicMock(return_value=1_000.0) + monkeypatch.setattr(image_handling, "pdf_probe_verdicts", InMemoryCache(clock=clock)) + url = probe_server.serve("/report.pdf", _Served(status=status)) + + while_down = await _serves_pdf(mode, url) + probe_server.serve("/report.pdf", _Served()) + clock.return_value = 1_030.0 + within_the_minute = await _serves_pdf(mode, url) + clock.return_value = 1_061.0 + after_the_minute = await _serves_pdf(mode, url) + + assert (while_down, within_the_minute, after_the_minute) == (False, False, True) + assert len(probe_server.seen) == 2 + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_format_pdf_links_marks_only_links_that_serve_a_pdf(probe_server, mode): + pdf_link = probe_server.serve("/a/report.pdf", _Served()) + query_link = probe_server.serve("/a/download?file=summary.pdf", _Served(content_type="application/pdf")) + png_at_pdf_link = probe_server.serve("/a/page.pdf?fm=png", _Served(content_type="image/png", body=PNG_BODY)) + png_link = f"{probe_server.base_url}/a/photo.png" + messages = [ + {"role": "system", "content": "read the files"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "compare"}, + {"type": "image_url", "image_url": {"url": pdf_link, "detail": "high"}, "cache_control": {"type": "ephemeral"}}, + {"type": "image_url", "image_url": query_link}, + {"type": "image_url", "image_url": {"url": png_at_pdf_link}}, + {"type": "image_url", "image_url": {"url": png_link}}, + {"type": "image_url", "image_url": {"url": pdf_link, "format": "image/png"}}, + ], + }, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": pdf_link}}]}, + ] + original = copy.deepcopy(messages) + + formatted = await _format_pdf_links(mode, messages) + + assert messages == original + assert formatted == ( + messages[0], + { + "role": "user", + "content": [ + {"type": "text", "text": "compare"}, + { + "type": "image_url", + "image_url": {"url": pdf_link, "detail": "high", "format": "application/pdf"}, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "image_url", "image_url": {"url": query_link, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": png_at_pdf_link}}, + {"type": "image_url", "image_url": {"url": png_link}}, + {"type": "image_url", "image_url": {"url": pdf_link, "format": "image/png"}}, + ], + }, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": pdf_link, "format": "application/pdf"}}]}, + ) + assert sorted(path for path, _ in probe_server.seen) == ["/a/download?file=summary.pdf", "/a/page.pdf?fm=png", "/a/report.pdf"] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_format_pdf_links_probes_every_link_at_once(probe_server, mode): + both_in_flight = threading.Barrier(2, timeout=5) + links = [probe_server.serve(f"/{name}.pdf", _Served(barrier=both_in_flight)) for name in ("first", "second")] + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": link}} for link in links]}] + + formatted = await _format_pdf_links(mode, messages) + + assert [part["image_url"]["format"] for part in formatted[0]["content"]] == ["application/pdf", "application/pdf"] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_format_pdf_links_leaves_messages_without_pdf_links_alone(probe_server, mode): + messages = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": f"{probe_server.base_url}/a.png"}}]}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:application/pdf;base64,JVBERi0="}}]}, + ] + + assert await _format_pdf_links(mode, messages) == tuple(messages) + assert probe_server.seen == [] diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f1854563cd5..13e181b2b92 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -11090,6 +11090,40 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost() assert logging_obj.response_cost_calculator(result=assembled) == 0.0042 +@pytest.mark.parametrize("continued_response_id", ["resp_continuedafteracutstream", None]) +@pytest.mark.parametrize("client_facing_response_id", ["resp_encryptedidtheclientsaw", None]) +def test_sync_success_handler_logs_a_stream_under_the_id_the_client_saw( + client_facing_response_id, continued_response_id +): + class Recorder(CustomLogger): + def __init__(self): + super().__init__() + self.logged = [] + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged.append((response_obj.id, kwargs["standard_logging_object"]["id"])) + + recorder: Final = Recorder() + logging_obj: Final = _responses_stream_logging_obj() + logging_obj.dynamic_success_callbacks = [recorder] + logging_obj.client_facing_response_id = client_facing_response_id + if continued_response_id is not None: + logging_obj.record_continued_partial( + Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), 0.0, continued_response_id + ) + now: Final = datetime.datetime.now() + + logging_obj.success_handler( + result=_completed_responses_event(ResponseAPIUsage(input_tokens=3, output_tokens=2, total_tokens=5)), + start_time=now, + end_time=now, + cache_hit=False, + ) + + expected_id: Final = client_facing_response_id or continued_response_id or "resp-1" + assert recorder.logged == [(expected_id, expected_id)] + + def test_response_cost_calculator_prices_terminal_responses_event_from_its_response(): logging_obj: Final = _responses_stream_logging_obj() inner_response: Final = ResponsesAPIResponse( diff --git a/tests/unit/litellm_core_utils/test_model_param_helper.py b/tests/unit/litellm_core_utils/test_model_param_helper.py index 6293c748505..df37139bc7b 100644 --- a/tests/unit/litellm_core_utils/test_model_param_helper.py +++ b/tests/unit/litellm_core_utils/test_model_param_helper.py @@ -1,5 +1,9 @@ +from collections.abc import Sequence +from typing import Final +import pytest +from litellm.caching.caching import Cache from litellm.litellm_core_utils.model_param_helper import ModelParamHelper @@ -46,3 +50,18 @@ def test_langfuse_model_parameters_no_secret_leakage(): assert "secret_fields" not in sanitized assert "authorization" not in sanitized assert "headers" not in sanitized + + +@pytest.mark.parametrize( + "hint,first_value,second_value", + [("languages", ("en",), ("fr",)), ("keywords", ("LiteLLM",), ("Kubernetes",))], +) +def test_transcription_cache_key_separates_recognition_hints( + hint: str, first_value: Sequence[str], second_value: Sequence[str] +) -> None: + cache: Final = Cache() + request: Final = {"model": "gpt-4o-transcribe", "file": "meeting.wav", "response_format": "json"} + + assert cache.get_cache_key(**request, **{hint: first_value}) != cache.get_cache_key( + **request, **{hint: second_value} + ) diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index ba76138a268..eec1dfcf459 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -30,6 +30,8 @@ from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTransc from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.xai.realtime.transformation import XAIRealtimeNormalizer from litellm.llms.xai.realtime.usage_meter import XAIRealtimeUsageMeter + +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.realtime import RealtimeResponseTransformInput, RealtimeResponseTypedDict @@ -512,6 +514,76 @@ async def test_backend_to_client_beta_receives_normalized_events(): assert sent["output_index"] == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("event_type", ("session.created", "session.updated")) +async def test_beta_client_bridged_to_a_ga_backend_receives_the_beta_session_shape(event_type: str) -> None: + client_ws = MagicMock() + client_ws.scope = {"headers": [(b"openai-beta", b"realtime=v1")]} + client_ws.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps( + { + "type": event_type, + "event_id": "e1", + "session": { + "type": "realtime", + "object": "realtime.session", + "id": "sess_1", + "model": "gpt-realtime-2.1-mini", + "instructions": "Answer briefly", + "output_modalities": ["audio"], + "max_output_tokens": "inf", + "tools": [], + "audio": { + "input": { + "format": {"type": "audio/pcmu"}, + "transcription": {"model": "gpt-transcribe"}, + "noise_reduction": None, + "turn_detection": {"type": "server_vad", "create_response": True}, + }, + "output": {"format": {"type": "audio/pcm", "rate": 24000}, "voice": "marin", "speed": 1.1}, + }, + }, + } + ).encode(), + ConnectionClosed(None, None), + ] + ) + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, logging_obj, backend_uses_beta_protocol=False, bridges_beta_client_to_ga=True + ) + + await streaming.backend_to_client_send_messages() + + sent: Final = json.loads(client_ws.send_text.call_args_list[0].args[0]) + assert sent["type"] == event_type + assert sent["session"] == { + "object": "realtime.session", + "id": "sess_1", + "model": "gpt-realtime-2.1-mini", + "instructions": "Answer briefly", + "modalities": ["audio", "text"], + "max_response_output_tokens": "inf", + "tools": [], + "input_audio_format": "g711_ulaw", + "input_audio_transcription": {"model": "gpt-transcribe"}, + "input_audio_noise_reduction": None, + "turn_detection": {"type": "server_vad", "create_response": True}, + "output_audio_format": "pcm16", + "voice": "marin", + "speed": 1.1, + }, ( + "Beta and GA session fields, checked 2026-10-20: " + "https://github.com/openai/openai-python/blob/v2.33.0/src/openai/types/beta/realtime/session.py and " + "https://github.com/openai/openai-python/blob/v2.33.0/src/openai/types/realtime/realtime_session_create_response.py" + ) + + @pytest.mark.asyncio async def test_backend_to_client_stores_normalized_events_for_logging(): client_ws = MagicMock() @@ -648,6 +720,50 @@ async def test_client_ack_messages_keeps_beta_session_shape_for_beta_backend(): assert "audio" not in session +@pytest.mark.asyncio +@pytest.mark.parametrize("beta_client", (False, True)) +async def test_translation_session_update_omits_session_type(beta_client: bool) -> None: + client_ws = MagicMock() + client_ws.scope = {"headers": [(b"openai-beta", b"realtime=v1")] if beta_client else []} + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "session.update", + "session": { + "type": "translation", + "input_audio_format": "g711_ulaw", + "output_audio_format": "pcm16", + "audio": {"output": {"language": "fr"}}, + }, + } + ), + Exception("connection closed"), + ] + ) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + streaming = RealTimeStreaming( + client_ws, + backend_ws, + logging_obj, + translation_session=True, + ) + + await streaming.client_ack_messages() + + sent_to_backend = json.loads(backend_ws.send.call_args_list[0].args[0]) + assert "type" not in sent_to_backend["session"] + assert sent_to_backend["session"]["audio"]["output"]["language"] == "fr" + assert sent_to_backend["session"]["audio"]["input"]["format"]["type"] == "audio/pcmu", ( + "OpenAI Realtime audio formats, checked 2026-10-06: " + "https://github.com/openai/openai-python/blob/v2.33.0/src/openai/types/realtime/realtime_audio_formats.py" + ) + assert sent_to_backend["session"]["audio"]["output"]["format"]["type"] == "audio/pcm" + + def test_translate_event_to_beta_renames_delta_types(): ev = RealTimeStreaming._translate_event_to_beta( {"type": "response.output_audio.delta", "delta": "abc", "event_id": "e1"} @@ -1158,6 +1274,153 @@ async def test_transcription_session_update_enforces_authorized_nested_model(): assert streaming._is_transcription_session is True +@pytest.mark.asyncio +async def test_translation_session_update_rejects_disallowed_nested_transcription_model() -> None: + client_ws: Final = MagicMock(send_text=AsyncMock(), close=AsyncMock()) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + MagicMock(), + model="gpt-realtime-translate", + user_api_key_dict=UserAPIKeyAuth(models=["gpt-realtime-translate"]), + translation_session=True, + ) + + sent: Final = await streaming._send_to_backend( + json.dumps( + { + "type": "session.update", + "session": { + "type": "translation", + "audio": { + "input": { + "transcription": {"model": "gpt-live-transcribe"}, + } + }, + }, + } + ) + ) + assert sent is False + backend_ws.send.assert_not_awaited() + assert streaming._is_transcription_session is False + error: Final = json.loads(client_ws.send_text.await_args.args[0]) + assert error["type"] == "error" + assert error["error"]["type"] == "permission_error" + assert "gpt-live-transcribe" in error["error"]["message"] + assert client_ws.close.await_args.kwargs["code"] == 1008 + + +def _translation_transcription_update(transcription: Mapping[str, object]) -> str: + return json.dumps( + { + "type": "session.update", + "session": {"type": "translation", "audio": {"input": {"transcription": transcription}}}, + } + ) + + +@pytest.mark.asyncio +async def test_translation_session_update_authorizes_each_new_nested_transcription_model() -> None: + backend_ws: Final = MagicMock(send=AsyncMock()) + streaming: Final = RealTimeStreaming( + MagicMock(), + backend_ws, + MagicMock(), + model="gpt-realtime-translate", + user_api_key_dict=UserAPIKeyAuth( + models=["gpt-realtime-translate", "gpt-realtime-whisper", "gpt-live-transcribe"] + ), + translation_session=True, + ) + + for transcription in ( + {"model": "gpt-realtime-whisper", "language": "en"}, + {"model": "gpt-live-transcribe", "language": "fr"}, + {"language": "de"}, + ): + await streaming._send_to_backend(_translation_transcription_update(transcription)) + + assert [ + json.loads(call.args[0])["session"]["audio"]["input"]["transcription"] + for call in backend_ws.send.await_args_list + ] == [ + {"model": "gpt-realtime-whisper", "language": "en"}, + {"model": "gpt-live-transcribe", "language": "fr"}, + {"model": "gpt-live-transcribe", "language": "de"}, + ] + + +@pytest.mark.asyncio +async def test_translation_session_update_refuses_a_later_nested_transcription_model_the_key_cannot_call() -> None: + client_ws: Final = MagicMock(send_text=AsyncMock(), close=AsyncMock()) + backend_ws: Final = MagicMock(send=AsyncMock()) + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + MagicMock(), + model="gpt-realtime-translate", + user_api_key_dict=UserAPIKeyAuth(models=["gpt-realtime-translate", "gpt-realtime-whisper"]), + translation_session=True, + ) + + first_sent: Final = await streaming._send_to_backend( + _translation_transcription_update({"model": "gpt-realtime-whisper", "language": "en"}) + ) + second_sent: Final = await streaming._send_to_backend( + _translation_transcription_update({"model": "gpt-live-transcribe", "language": "fr"}) + ) + + assert (first_sent, second_sent) == (True, False) + assert backend_ws.send.await_count == 1 + error: Final = json.loads(client_ws.send_text.await_args.args[0]) + assert error["error"]["type"] == "permission_error" + assert "gpt-live-transcribe" in error["error"]["message"] + assert client_ws.close.await_args.kwargs["code"] == 1008 + + +@pytest.mark.asyncio +async def test_translation_session_update_sends_the_deployment_model_for_a_nested_transcription_alias( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + model_list: Final = [ + { + "model_name": "whisper-alias", + "litellm_params": { + "model": "azure/qa-whisper-deployment", + "api_base": "https://realtime-qa.openai.azure.com", + "api_key": "azure-test-key", + }, + } + ] + monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=model_list)) + monkeypatch.setattr(proxy_server, "llm_model_list", model_list) + backend_ws: Final = MagicMock(send=AsyncMock()) + streaming: Final = RealTimeStreaming( + MagicMock(), + backend_ws, + MagicMock(), + model="gpt-realtime-translate", + user_api_key_dict=UserAPIKeyAuth(models=["gpt-realtime-translate", "whisper-alias"]), + translation_session=True, + ) + + await streaming._send_to_backend(_translation_transcription_update({"model": "whisper-alias", "language": "en"})) + await streaming._send_to_backend(_translation_transcription_update({"language": "fr"})) + + assert [ + json.loads(call.args[0])["session"]["audio"]["input"]["transcription"] + for call in backend_ws.send.await_args_list + ] == [ + {"model": "qa-whisper-deployment", "language": "en"}, + {"model": "qa-whisper-deployment", "language": "fr"}, + ] + + @pytest.mark.asyncio async def test_normal_realtime_session_keeps_nested_transcription_model(): backend_ws = MagicMock() @@ -1847,8 +2110,7 @@ async def test_realtime_function_call_output_guardrail_allows_clean_output(monke @pytest.mark.asyncio async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pytest.MonkeyPatch): """ - Test that _has_realtime_guardrails returns True for a guardrail configured with - pre_call mode (not just realtime_input_transcription). + A pre_call guardrail gates typed user messages and tool output, never audio VAD. """ import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -1870,10 +2132,6 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pyt logging_obj: Final = _mock_logging() streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) - assert streaming._has_realtime_guardrails() is True, ( - "pre_call guardrail should be recognized as a realtime guardrail" - ) - # pre_call-only guardrails gate typed user messages / tool output, not audio VAD. assert streaming._has_audio_transcription_guardrails() is False, ( "pre_call-only guardrail must not disable server_vad auto-response" ) @@ -2024,7 +2282,6 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monke }, ) - assert streaming._has_realtime_guardrails() is True assert streaming._has_audio_transcription_guardrails() is False @@ -3561,6 +3818,444 @@ def test_store_message_skips_pydantic_for_unlogged_audio_delta(): assert streaming.messages == [] +@pytest.mark.parametrize( + "event_type", ["session.output_audio.delta", "response.output_audio.delta", "response.audio.delta"] +) +def test_translation_audio_duration_is_finalized_once(event_type: str): + import base64 + + streaming = RealTimeStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=_mock_logging(), + model="gpt-realtime-translate", + translation_session=True, + ) + payload = base64.b64encode(bytes(48000)).decode() + streaming._capture_translation_output_audio({"type": event_type, "delta": payload}) + streaming._finalize_translation_usage() + streaming._finalize_translation_usage() + + closed_events = [event for event in streaming.messages if event.get("type") == "session.closed"] + assert len(closed_events) == 1 + assert closed_events[0]["usage"] == {"type": "duration", "output_seconds": 1.0} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("output_bytes", (0, 48000)) +@pytest.mark.parametrize("format_event_type", ("session.created", "session.updated")) +@pytest.mark.parametrize( + "audio_format,bytes_per_second", + ( + ("pcm16", 48000), + ("g711_ulaw", 8000), + ("g711_alaw", 8000), + ({"type": "audio/pcm", "rate": 24000}, 48000), + ({"type": "audio/pcmu", "rate": 8000}, 8000), + ({"type": "audio/pcma", "rate": 8000}, 8000), + ({"type": "audio/pcmu"}, 8000), + ({"type": "audio/pcma"}, 8000), + ), +) +async def test_translation_disconnect_bills_sent_input_audio( + output_bytes: int, format_event_type: str, audio_format: str | Mapping[str, object], bytes_per_second: int +) -> None: + import base64 + + backend: Final = MagicMock() + backend.send = AsyncMock() + streaming: Final = RealTimeStreaming( + websocket=_ga_client_ws(), + backend_ws=backend, + logging_obj=_mock_logging(), + model="gpt-realtime-translate", + translation_session=True, + ) + format_event: Final = { + "type": format_event_type, + "session": {"audio": {"input": {"format": audio_format}}}, + } + streaming._capture_translation_output_audio(format_event) + await streaming._send_to_backend( + json.dumps( + { + "type": "session.input_audio_buffer.append", + "audio": base64.b64encode(bytes(bytes_per_second * 2)).decode(), + } + ) + ) + streaming._capture_translation_output_audio( + {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(output_bytes)).decode()} + ) + streaming._finalize_translation_usage() + streaming._finalize_translation_usage() + + assert streaming.messages == [ + { + "type": "session.closed", + "usage": {"type": "duration", "input_seconds": 2.0, "output_seconds": output_bytes / 48000}, + } + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rate", (24000, 1, 0, -1, float("nan"), float("inf"))) +async def test_translation_client_format_updates_do_not_change_billed_duration(rate: int | float) -> None: + backend: Final = MagicMock() + backend.send = AsyncMock() + streaming: Final = RealTimeStreaming( + websocket=_ga_client_ws(), backend_ws=backend, logging_obj=_mock_logging(), translation_session=True + ) + streaming._capture_translation_output_audio( + { + "type": "session.created", + "session": {"audio": {"input": {"format": "g711_ulaw"}, "output": {"format": "g711_ulaw"}}}, + } + ) + input_event: Final = json.dumps( + {"type": "session.input_audio_buffer.append", "audio": base64.b64encode(bytes(8000)).decode()} + ) + output_event: Final = {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(8000)).decode()} + await streaming._send_to_backend(input_event) + streaming._capture_translation_output_audio(output_event) + await streaming._send_to_backend( + json.dumps( + { + "type": "session.update", + "session": { + "audio": { + "input": {"format": {"type": "audio/pcm", "rate": rate}}, + "output": {"format": {"type": "audio/pcm", "rate": rate}}, + } + }, + } + ) + ) + await streaming._send_to_backend(input_event) + streaming._capture_translation_output_audio(output_event) + streaming._finalize_translation_usage() + + assert streaming.messages == [ + {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 2.0, "output_seconds": 2.0}} + ] + + +@pytest.mark.asyncio +async def test_translation_confirmed_format_changes_preserve_previous_audio_duration() -> None: + backend: Final = MagicMock() + backend.send = AsyncMock() + streaming: Final = RealTimeStreaming( + websocket=_ga_client_ws(), backend_ws=backend, logging_obj=_mock_logging(), translation_session=True + ) + await streaming._send_to_backend( + json.dumps({"type": "session.input_audio_buffer.append", "audio": base64.b64encode(bytes(48000)).decode()}) + ) + streaming._capture_translation_output_audio( + {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(48000)).decode()} + ) + streaming._capture_translation_output_audio( + { + "type": "session.updated", + "session": {"audio": {"input": {"format": "g711_alaw"}, "output": {"format": "g711_alaw"}}}, + } + ) + await streaming._send_to_backend( + json.dumps({"type": "session.input_audio_buffer.append", "audio": base64.b64encode(bytes(8000)).decode()}) + ) + streaming._capture_translation_output_audio( + {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(8000)).decode()} + ) + streaming._finalize_translation_usage() + + assert streaming.messages == [ + {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 2.0, "output_seconds": 2.0}} + ] + + +@pytest.mark.parametrize("rate", (1, 0, -1, float("nan"), float("inf"), 10**400)) +def test_translation_invalid_provider_format_preserves_audio_duration(rate: int | float) -> None: + streaming: Final = RealTimeStreaming( + websocket=_ga_client_ws(), backend_ws=MagicMock(), logging_obj=_mock_logging(), translation_session=True + ) + output_event: Final = {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(48000)).decode()} + streaming._capture_translation_output_audio(output_event) + streaming._capture_translation_output_audio( + {"type": "session.updated", "session": {"audio": {"output": {"format": {"type": "audio/pcm", "rate": rate}}}}} + ) + streaming._capture_translation_output_audio(output_event) + streaming._finalize_translation_usage() + + assert streaming.messages == [{"type": "session.closed", "usage": {"type": "duration", "output_seconds": 2.0}}] + + +@pytest.mark.asyncio +async def test_translation_failed_audio_send_is_not_billed() -> None: + backend: Final = MagicMock() + backend.send = AsyncMock(side_effect=RuntimeError("send failed")) + streaming: Final = RealTimeStreaming( + websocket=_ga_client_ws(), backend_ws=backend, logging_obj=_mock_logging(), translation_session=True + ) + + with pytest.raises(RuntimeError, match="send failed"): + await streaming._send_to_backend(json.dumps({"type": "session.input_audio_buffer.append", "audio": "AAAA"})) + streaming._finalize_translation_usage() + + assert streaming.messages == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("default_on,requested", ((False, False), (True, False), (False, True))) +@pytest.mark.parametrize( + "event_hook", + (GuardrailEventHooks.realtime_input_transcription, GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call), +) +async def test_translation_rejects_applicable_guardrails_before_forwarding_audio( + monkeypatch: pytest.MonkeyPatch, default_on: bool, requested: bool, event_hook: GuardrailEventHooks +) -> None: + guardrail: Final = CustomGuardrail(guardrail_name="speech-policy", event_hook=event_hook, default_on=default_on) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + input_event: Final = json.dumps( + {"type": "session.input_audio_buffer.append", "audio": base64.b64encode(bytes(4800)).decode()} + ) + output_event: Final = {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(4800)).decode()} + client: Final = _ga_client_ws() + client.receive_text = AsyncMock(side_effect=[input_event, ConnectionClosed(None, None)]) + client.send_text = AsyncMock() + client.close = AsyncMock() + backend: Final = MagicMock() + backend.send = AsyncMock() + backend.recv = AsyncMock(side_effect=[json.dumps(output_event), ConnectionClosed(None, None)]) + streaming: Final = RealTimeStreaming( + websocket=client, + backend_ws=backend, + logging_obj=_mock_logging(), + translation_session=True, + request_data={"litellm_metadata": {"guardrails": ["speech-policy"] if requested else []}}, + ) + + await streaming.bidirectional_forward() + + messages: Final = tuple(json.loads(call.args[0]) for call in client.send_text.await_args_list) + if default_on or requested: + assert messages[0]["error"]["type"] == "guardrail_error" + client.close.assert_awaited_once_with( + code=1008, reason="Translation sessions cannot enforce configured realtime guardrails" + ) + backend.send.assert_not_awaited() + else: + assert output_event in messages + backend.send.assert_awaited_once_with(input_event) + + +@pytest.mark.asyncio +async def test_translation_for_a_key_opted_out_of_a_legacy_callback_guardrail_gets_no_realtime_session_gate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="legacy-callback")]) + session_update: Final = json.dumps({"type": "session.update", "session": {"audio": {"output": {"language": "es"}}}}) + session_created: Final = {"type": "session.created", "session": {"type": "translation"}} + client: Final = _ga_client_ws() + client.receive_text = AsyncMock(side_effect=[session_update, ConnectionClosed(None, None)]) + client.close = AsyncMock() + backend: Final = MagicMock() + backend.send = AsyncMock() + backend.recv = AsyncMock(side_effect=[json.dumps(session_created), ConnectionClosed(None, None)]) + streaming: Final = RealTimeStreaming( + websocket=client, + backend_ws=backend, + logging_obj=_mock_logging(), + translation_session=True, + request_data={"litellm_metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}}, + ) + + await streaming.bidirectional_forward() + + assert tuple(call.args[0] for call in backend.send.await_args_list) == (session_update,) + assert session_created in tuple(json.loads(call.args[0]) for call in client.send_text.await_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retain_close", (False, True)) +@pytest.mark.parametrize("reported_input,expected_input", [(None, 2.0), (0.0, 0.0), (0.25, 0.25)]) +async def test_translation_terminal_usage_fills_only_missing_input_duration( + monkeypatch: pytest.MonkeyPatch, retain_close: bool, reported_input: float | None, expected_input: float +) -> None: + import base64 + + monkeypatch.setattr(litellm, "logged_real_time_event_types", "*" if retain_close else None) + backend: Final = MagicMock() + backend.send = AsyncMock() + streaming: Final = RealTimeStreaming( + websocket=_ga_client_ws(), backend_ws=backend, logging_obj=_mock_logging(), translation_session=True + ) + await streaming._send_to_backend( + json.dumps({"type": "session.input_audio_buffer.append", "audio": base64.b64encode(bytes(96000)).decode()}) + ) + close_event: Final = { + "type": "session.closed", + "usage": { + "type": "duration", + "output_seconds": 0.5, + **({"input_seconds": reported_input} if reported_input is not None else {}), + }, + } + streaming._capture_translation_output_audio(close_event) + streaming.store_message(close_event) + streaming._finalize_translation_usage() + + usage: Final = tuple(event["usage"] for event in streaming.messages if event.get("type") == "session.closed") + assert sum(item.get("input_seconds") or 0.0 for item in usage) == expected_input + assert sum(item.get("output_seconds") or 0.0 for item in usage) == 0.5 + + +def test_translation_audio_duration_uses_session_output_format(): + import base64 + + streaming = RealTimeStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=_mock_logging(), + model="gpt-realtime-translate", + translation_session=True, + ) + streaming._capture_translation_output_audio( + { + "type": "session.created", + "session": {"audio": {"output": {"format": {"type": "audio/pcmu", "rate": 8000}}}}, + } + ) + streaming._capture_translation_output_audio( + { + "type": "session.output_audio.delta", + "delta": base64.b64encode(bytes(8000)).decode(), + } + ) + streaming._finalize_translation_usage() + + assert streaming.messages[-1]["usage"] == {"type": "duration", "output_seconds": 1.0} + + +def test_translation_does_not_duplicate_provider_duration_usage(): + streaming = RealTimeStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=_mock_logging(), + model="gpt-realtime-translate", + translation_session=True, + ) + streaming._capture_translation_output_audio( + {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(48000)).decode()} + ) + streaming.messages.append({"type": "session.closed", "usage": {"type": "duration", "output_seconds": 0.5}}) + streaming._finalize_translation_usage() + + closed_events = [event for event in streaming.messages if event.get("type") == "session.closed"] + assert len(closed_events) == 1 + + +def test_translation_prefers_provider_duration_over_audio_byte_estimate(): + import base64 + + streaming = RealTimeStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=_mock_logging(), + model="gpt-realtime-translate", + translation_session=True, + ) + streaming._capture_translation_output_audio( + {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(48000)).decode()} + ) + streaming._capture_translation_output_audio( + {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 0.25, "output_seconds": 0.5}} + ) + streaming._finalize_translation_usage() + + closed_events = [event for event in streaming.messages if event.get("type") == "session.closed"] + assert len(closed_events) == 1 + assert closed_events[0]["usage"] == {"type": "duration", "input_seconds": 0.25, "output_seconds": 0.5} + + +@pytest.mark.parametrize( + ("output_audio_bytes", "expected_usage"), + [ + (0, {"type": "duration", "input_seconds": 0.25, "output_seconds": 0.0}), + (48000, {"type": "duration", "input_seconds": 0.25, "output_seconds": 1.0}), + ], +) +def test_translation_preserves_input_only_provider_usage( + output_audio_bytes: int, expected_usage: Mapping[str, str | float] +) -> None: + streaming = RealTimeStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=_mock_logging(), + model="gpt-realtime-translate", + translation_session=True, + ) + streaming._capture_translation_output_audio( + {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(output_audio_bytes)).decode()} + ) + streaming._capture_translation_output_audio( + {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 0.25}} + ) + streaming._finalize_translation_usage() + + closed_events = [event for event in streaming.messages if event.get("type") == "session.closed"] + assert len(closed_events) == 1 + assert closed_events[0]["usage"] == expected_usage + + +def test_translation_retained_input_only_close_event_bills_captured_output(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "logged_real_time_event_types", "*") + streaming = RealTimeStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=_mock_logging(), + model="gpt-realtime-translate", + translation_session=True, + ) + streaming._capture_translation_output_audio( + {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(48000)).decode()} + ) + close_event: Final = {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 0.25}} + + streaming._capture_translation_output_audio(close_event) + streaming.store_message(close_event) + streaming._finalize_translation_usage() + + usage_events: Final = tuple(event["usage"] for event in streaming.messages if event.get("type") == "session.closed") + assert len(usage_events) == 2 + assert sum(usage.get("input_seconds", 0.0) for usage in usage_events) == 0.25 + assert sum(usage.get("output_seconds", 0.0) for usage in usage_events) == 1.0 + + +@pytest.mark.parametrize("malformed_output", [-1.0, float("nan"), float("inf"), 10**1000]) +def test_translation_malformed_provider_duration_uses_captured_audio(malformed_output: float | int) -> None: + streaming = RealTimeStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=_mock_logging(), + model="gpt-realtime-translate", + translation_session=True, + ) + streaming._capture_translation_output_audio( + {"type": "session.output_audio.delta", "delta": base64.b64encode(bytes(48000)).decode()} + ) + + streaming._capture_translation_output_audio( + { + "type": "session.closed", + "usage": {"type": "duration", "input_seconds": 0.25, "output_seconds": malformed_output}, + } + ) + streaming._finalize_translation_usage() + + close_events: Final = tuple(event for event in streaming.messages if event.get("type") == "session.closed") + assert len(close_events) == 1 + assert close_events[0]["usage"] == {"type": "duration", "input_seconds": 0.25, "output_seconds": 1.0} + + @pytest.mark.asyncio async def test_audio_delta_frame_parsed_at_most_once(): client_ws = _beta_client_ws() @@ -4200,6 +4895,48 @@ async def test_upstream_close_after_relayed_events_still_logs_the_session_as_suc client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) +def _model_not_found_error_event() -> bytes: + return json.dumps( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "model_not_found", + "message": "The model `gpt-realtime-translate-typo` does not exist", + }, + } + ).encode() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("close_code", (1000, 1008)) +async def test_upstream_error_before_any_session_logs_a_failure_not_a_success(close_code: int): + upstream_close: Final = ConnectionClosed(Close(close_code, ""), None) + session: Final = _relay_session( + _client_ws_that_never_sends(), _backend_ws_closing_with(_model_not_found_error_event(), upstream_close) + ) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert session.logging.logged_sessions == () + + +@pytest.mark.asyncio +async def test_upstream_error_after_session_created_still_logs_the_session_as_success(): + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + upstream_close: Final = ConnectionClosed(Close(1000, ""), None) + session: Final = _relay_session( + _client_ws_that_never_sends(), + _backend_ws_closing_with(session_created, _model_not_found_error_event(), upstream_close), + ) + + await session.run() + + assert len(session.logging.logged_sessions) == 1 + assert session.logging.logged_failures == () + + @pytest.mark.asyncio async def test_upstream_closing_while_a_client_message_is_forwarded_still_reaches_the_client(): upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) @@ -5029,3 +5766,264 @@ async def test_blocked_transcript_is_removed_from_the_conversation_before_the_re sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] assert [event["type"] for event in sent] == expected_frames assert all(event["item_id"] == "item_blocked" for event in sent if event["type"] == "conversation.item.delete") + + +@pytest.mark.asyncio +async def test_beta_client_bridged_to_a_ga_backend_sends_assistant_text_items_as_ga_output_text() -> None: + client_ws = MagicMock() + client_ws.scope = {"headers": [(b"openai-beta", b"realtime=v1")]} + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "The forecast is sunny"}], + }, + } + ), + json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Thanks"}]}, + } + ), + Exception("connection closed"), + ] + ) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, logging_obj, backend_uses_beta_protocol=False, bridges_beta_client_to_ga=True + ) + + await streaming.client_ack_messages() + + sent_items: Final = [json.loads(call.args[0])["item"] for call in backend_ws.send.call_args_list] + assert [(item["role"], item["content"]) for item in sent_items] == [ + ("assistant", [{"type": "output_text", "text": "The forecast is sunny"}]), + ("user", [{"type": "input_text", "text": "Thanks"}]), + ], ( + "Beta and GA item content types, checked 2026-10-28: " + "https://github.com/openai/openai-python/blob/v2.33.0/src/openai/types/beta/realtime/conversation_item_content.py and " + "https://github.com/openai/openai-python/blob/v2.33.0/src/openai/types/realtime/realtime_conversation_item_assistant_message.py" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_type", ("response.created", "response.done")) +async def test_beta_client_bridged_to_a_ga_backend_receives_the_beta_response_shape(event_type: str) -> None: + client_ws = MagicMock() + client_ws.scope = {"headers": [(b"openai-beta", b"realtime=v1")]} + client_ws.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps( + { + "type": event_type, + "event_id": "e1", + "response": { + "object": "realtime.response", + "id": "resp_1", + "status": "completed", + "conversation_id": "conv_1", + "output_modalities": ["audio"], + "max_output_tokens": "inf", + "audio": {"output": {"format": {"type": "audio/pcmu"}, "voice": "marin"}}, + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_audio", "transcript": "Sunny"}], + } + ], + }, + } + ).encode(), + ConnectionClosed(None, None), + ] + ) + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, logging_obj, backend_uses_beta_protocol=False, bridges_beta_client_to_ga=True + ) + + await streaming.backend_to_client_send_messages() + + sent: Final = json.loads(client_ws.send_text.call_args_list[0].args[0]) + assert sent["type"] == event_type + assert sent["response"] == { + "object": "realtime.response", + "id": "resp_1", + "status": "completed", + "conversation_id": "conv_1", + "modalities": ["audio", "text"], + "max_output_tokens": "inf", + "output_audio_format": "g711_ulaw", + "voice": "marin", + "output": [{"type": "message", "role": "assistant", "content": [{"type": "audio", "transcript": "Sunny"}]}], + }, ( + "Beta and GA response fields, checked 2026-10-28: " + "https://github.com/openai/openai-python/blob/v2.33.0/src/openai/types/beta/realtime/realtime_response.py and " + "https://github.com/openai/openai-python/blob/v2.33.0/src/openai/types/realtime/realtime_response.py" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", (False, True)) +async def test_beta_client_bridged_to_a_ga_backend_sends_ga_session_and_response_fields(drop_params: bool) -> None: + client_ws = MagicMock() + client_ws.scope = {"headers": [(b"openai-beta", b"realtime=v1")]} + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "session.update", + "session": { + "instructions": "Answer briefly", + "temperature": 0.7, + "max_response_output_tokens": 300, + "speed": 1.1, + "input_audio_noise_reduction": {"type": "near_field"}, + }, + } + ), + json.dumps( + { + "type": "response.create", + "response": { + "modalities": ["text", "audio"], + "voice": "marin", + "output_audio_format": "g711_ulaw", + "temperature": 0.7, + "max_response_output_tokens": 200, + "instructions": "Say hi", + }, + } + ), + Exception("connection closed"), + ] + ) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + logging_obj, + backend_uses_beta_protocol=False, + bridges_beta_client_to_ga=True, + drop_params=drop_params, + ) + + await streaming.client_ack_messages() + + sent: Final = { + event["type"]: event for event in (json.loads(call.args[0]) for call in backend_ws.send.call_args_list) + } + temperature: Final = {} if drop_params else {"temperature": 0.7} + assert sent["session.update"]["session"] == { + "type": "realtime", + "instructions": "Answer briefly", + "max_output_tokens": 300, + "audio": {"input": {"noise_reduction": {"type": "near_field"}}, "output": {"speed": 1.1}}, + **temperature, + } + assert sent["response.create"]["response"] == { + "output_modalities": ["audio"], + "audio": {"output": {"voice": "marin", "format": {"type": "audio/pcmu"}}}, + "max_output_tokens": 200, + "instructions": "Say hi", + **temperature, + } + + +@pytest.mark.asyncio +async def test_gemini_live_session_update_keeps_the_client_output_token_cap() -> None: + from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig + + client_ws: Final = _ga_client_ws() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "session.update", + "session": {"temperature": 0.4, "max_response_output_tokens": 77, "modalities": ["text"]}, + } + ), + _TEXT_ITEM_CREATE, + _RESPONSE_CREATE, + Exception("connection closed"), + ] + ) + backend_ws: Final = MagicMock(send=AsyncMock()) + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + _mock_logging(), + provider_config=GeminiRealtimeConfig(), + model="gemini-2.0-flash-live-001", + ) + + await streaming.client_ack_messages() + + setup_frame, content_frame = (json.loads(call.args[0]) for call in backend_ws.send.call_args_list) + assert "clientContent" in content_frame + assert setup_frame["setup"]["generationConfig"] == { + "temperature": 0.4, + "maxOutputTokens": 77, + "responseModalities": ["TEXT"], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event_normalizer", + (pytest.param(XAIRealtimeNormalizer(), id="xai"), pytest.param(None, id="openai-ga-client")), +) +async def test_unbridged_backend_receives_session_update_and_response_create_as_the_client_sent_them( + event_normalizer: XAIRealtimeNormalizer | None, +) -> None: + client_events: Final = ( + { + "type": "session.update", + "session": { + "type": "realtime", + "instructions": "Answer briefly", + "temperature": 0.7, + "max_response_output_tokens": 300, + "speed": 1.1, + "input_audio_noise_reduction": {"type": "near_field"}, + }, + }, + { + "type": "response.create", + "response": { + "modalities": ["text", "audio"], + "voice": "marin", + "output_audio_format": "g711_ulaw", + "temperature": 0.7, + "max_response_output_tokens": 200, + }, + }, + ) + client_ws: Final = _ga_client_ws() + client_ws.receive_text = AsyncMock( + side_effect=[*(json.dumps(event) for event in client_events), Exception("connection closed")] + ) + backend_ws: Final = MagicMock(send=AsyncMock()) + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, _mock_logging(), event_normalizer=event_normalizer, drop_params=True + ) + + await streaming.client_ack_messages() + + assert tuple(json.loads(call.args[0]) for call in backend_ws.send.call_args_list) == client_events diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 9e385c588ae..2fd3f16ab17 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -112,6 +112,57 @@ def test_redaction_reuses_message_rules_for_request_body_and_provider_fields(sha assert result == original_result and params == original_params +def _reasoning_items() -> list[dict[str, object]]: + return [{"type": "reasoning", "id": "rs_redaction", "summary": [{"type": "summary_text", "text": SECRET}]}] + + +def _message_response() -> ModelResponse: + return ModelResponse( + choices=[{"index": 0, "message": {"role": "assistant", "content": "answer", "reasoning_items": _reasoning_items()}}] + ) + + +def _delta_choices() -> dict[str, object]: + return {"choices": [{"index": 0, "delta": {"content": "answer", "reasoning_items": _reasoning_items()}}]} + + +@pytest.mark.parametrize( + ("details", "result"), + [ + pytest.param({}, _message_response(), id="message_object"), + pytest.param({}, _message_response().model_dump(), id="message_dict"), + pytest.param({}, _delta_choices(), id="delta_dict"), + pytest.param( + {"stream": True, "complete_streaming_response": ModelResponseStream(**_delta_choices())}, + None, + id="stream_chunk", + ), + pytest.param({"standard_logging_object": {"response": _message_response().model_dump()}}, None, id="logged"), + pytest.param({"standard_logging_object": {"response": _delta_choices()}}, None, id="logged_delta"), + ], +) +def test_redaction_hides_reasoning_item_summaries(details: dict[str, object], result: object) -> None: + assert SECRET in str(details) + str(result) + + redacted: Final = perform_redaction(details, result) + + assert SECRET not in str(redacted) and SECRET not in str(details), (redacted, details) + + +def test_an_opted_out_callback_gets_a_text_completion_with_its_answer_redacted() -> None: + logprobs: Final = {"tokens": [SECRET], "token_logprobs": [-0.1], "text_offset": [0], "top_logprobs": [{SECRET: -0.1}]} + response: Final = TextCompletionResponse( + choices=[{"index": 0, "text": SECRET, "logprobs": logprobs, "finish_reason": "stop"}] + ) + + redacted: Final = redacted_response_obj_for_custom_logger(response, CustomLogger(message_logging=False)) + + assert isinstance(redacted, TextCompletionResponse) + assert redacted.choices[0].text == "redacted-by-litellm" + assert SECRET not in redacted.model_dump_json() + assert response.choices[0].text == SECRET, "the caller's own response is not redacted" + + def test_redaction_redacts_reasoning_copied_into_hidden_params_in_place() -> None: hidden_params: Final = {"model_id": "deployment", "provider_specific_fields": _provider_fields()} original_hidden_params: Final = copy.deepcopy(hidden_params) @@ -561,6 +612,83 @@ class TestPerformRedaction: {"text": "redacted-by-litellm", "finish_reason": "length", "index": 1}, ] + def test_redacts_token_ids_that_decode_to_the_completion(self): + details = { + "stream": True, + "complete_streaming_response": litellm.ModelResponse( + choices=[ + litellm.Choices(message=litellm.Message(content="secret", role="assistant"), token_ids=[20, 11]) + ] + ), + "standard_logging_object": { + "response": { + "choices": [ + { + "token_ids": [20, 11], + "message": {"content": "secret", "provider_specific_fields": {"token_ids": [20, 11]}}, + }, + {"text": " Paris.", "token_ids": [20, 11], "index": 1}, + ] + } + }, + } + + perform_redaction(details, None) + + logged_choices = details["standard_logging_object"]["response"]["choices"] + assert [choice["token_ids"] for choice in logged_choices] == [None, None] + assert logged_choices[0]["message"]["provider_specific_fields"]["token_ids"] is None + assert details["complete_streaming_response"].choices[0].get("token_ids") is None + + def test_redacts_text_completion_stream_text_and_the_token_ids_that_decode_to_it(self): + details = { + "stream": True, + "async_complete_streaming_response": litellm.TextCompletionResponse( + choices=[ + litellm.utils.TextChoices( + index=0, text=" Paris.", reasoning_content="thinking", token_ids=[20, 11], finish_reason="stop" + ) + ] + ), + } + + perform_redaction(details, None) + + logged_choice = details["async_complete_streaming_response"].choices[0] + assert (logged_choice.text, logged_choice.reasoning_content) == ("redacted-by-litellm", "redacted-by-litellm") + assert logged_choice.get("token_ids") is None + assert logged_choice.finish_reason == "stop" + + def test_redacts_token_ids_a_cache_hit_keeps_on_the_choice_provider_specific_fields(self): + result = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret", role="assistant"), + provider_specific_fields={"token_ids": [20, 11], "stop_reason": "END"}, + ) + ] + ) + details = { + "standard_logging_object": { + "response": { + "choices": [ + { + "message": {"content": "secret"}, + "provider_specific_fields": {"token_ids": [20, 11], "stop_reason": "END"}, + } + ] + } + } + } + + redacted = perform_redaction(details, result) + + logged_fields = details["standard_logging_object"]["response"]["choices"][0]["provider_specific_fields"] + assert logged_fields == {"token_ids": None, "stop_reason": "END"} + assert redacted.choices[0].provider_specific_fields["token_ids"] is None + assert redacted.choices[0].provider_specific_fields["stop_reason"] == "END" + assert result.choices[0].provider_specific_fields["token_ids"] == [20, 11] + def test_redacts_object_choices_inside_model_response_dict(self): result = { "choices": [ diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index bbf1897b3a7..527c83d5760 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -2757,6 +2757,95 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015) +def _vllm_chat_chunk(index: int, content: str, finish_reason: str | None = None, **fields: object) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + created=1, + model="qwen3", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason, **fields)], + ) + + +@pytest.mark.parametrize("finish_content", ["", "4"], ids=["finish-only", "content-then-finish"]) +def test_stream_chunk_builder_keeps_the_vllm_stop_reason_and_every_token_id_of_each_choice( + finish_content: str, +) -> None: + built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(0, "1,", stop_reason=None, token_ids=[16, 11]), + _vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]), + _vllm_chat_chunk(0, "2,3", stop_reason=None, token_ids=[17, 11, 18]), + _vllm_chat_chunk(0, finish_content, "stop", stop_reason=",5", token_ids=[19]), + ] + ) + + assert built is not None + assert [ + {field: choice[field] for field in ("index", "stop_reason", "token_ids")} + for choice in built.model_dump()["choices"] + ] == [ + {"index": 0, "stop_reason": ",5", "token_ids": [16, 11, 17, 11, 18, 19]}, + {"index": 1, "stop_reason": None, "token_ids": [22, 11]}, + ] + + +def test_stream_chunk_builder_keeps_the_vllm_fields_of_a_sole_choice_streamed_on_a_nonzero_index() -> None: + built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]), + _vllm_chat_chunk(1, "8", "stop", stop_reason=",9", token_ids=[23]), + ] + ) + + assert built is not None + assert {field: built.model_dump()["choices"][0].get(field) for field in ("stop_reason", "token_ids")} == { + "stop_reason": ",9", + "token_ids": [22, 11, 23], + } + + +def test_stream_chunk_builder_keeps_null_vllm_fields_and_adds_none_for_other_providers() -> None: + vllm_built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(0, "hi", stop_reason=None, token_ids=None), + _vllm_chat_chunk(0, "", "stop", stop_reason=None, token_ids=None), + ] + ) + other_built: Final = stream_chunk_builder( + chunks=[_vllm_chat_chunk(0, "hi"), _vllm_chat_chunk(0, "", "stop")] + ) + + assert vllm_built is not None and other_built is not None + assert { + field: vllm_built.model_dump()["choices"][0][field] for field in ("stop_reason", "token_ids") + } == {"stop_reason": None, "token_ids": None} + assert {"stop_reason", "token_ids"}.isdisjoint(other_built.model_dump()["choices"][0]) + + +def test_stream_chunk_builder_text_completion_keeps_each_choices_vllm_stop_reason_and_token_ids() -> None: + def text_chunk(index: int, text: str, **fields: object) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", + object="text_completion", + created=1, + model="qwen3", + choices=[TextChoices(index=index, text=text, **fields)], + ) + + built: Final = stream_chunk_builder_text_completion( + chunks=[ + text_chunk(0, "1,", token_ids=[16, 11], stop_reason=None), + text_chunk(1, "7,", token_ids=[22, 11], stop_reason=None), + text_chunk(0, "2", token_ids=[17], stop_reason=", 3"), + text_chunk(1, "8", token_ids=[23], stop_reason=151645), + ] + ) + + assert [ + (choice["index"], choice["stop_reason"], choice["token_ids"]) for choice in built.model_dump()["choices"] + ] == [(0, ", 3", [16, 11, 17]), (1, 151645, [22, 11, 23])] + + @pytest.mark.parametrize( ("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")] ) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index da88f1358ad..cf8de9c40f5 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -10,6 +10,7 @@ import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain from pathlib import Path from types import CodeType, FrameType, MappingProxyType, TracebackType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast @@ -27,6 +28,10 @@ from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, +) from litellm.litellm_core_utils.streaming_handler import ( CustomStreamWrapper, _ProviderChunkEarlyReturn, @@ -34,6 +39,7 @@ from litellm.litellm_core_utils.streaming_handler import ( ) from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as GeminiModelResponseIterator, @@ -6971,6 +6977,327 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj: assert result.choices[0].delta.content == "hello" +_TOKEN_IDS: Final = TypeAdapter(list[int] | None) +_CHOICE_INDEX: Final = TypeAdapter(int) +_DELTA: Final = TypeAdapter(Mapping[str, object]) + + +def _joined_token_ids(choices: Sequence[Mapping[str, object]]) -> tuple[int, ...]: + return tuple(chain.from_iterable(_TOKEN_IDS.validate_python(choice.get("token_ids")) or () for choice in choices)) + + +def _hosted_vllm_stream( + logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] +) -> CustomStreamWrapper: + sse_lines: Final = ( + *( + "data: " + json.dumps({"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [c]}) + for c in upstream_choices + ), + "data: [DONE]", + ) + completion_stream: Final = HostedVLLMChatConfig().get_model_response_iterator( + streaming_response=iter(sse_lines), sync_stream=True, json_mode=False + ) + return CustomStreamWrapper( + completion_stream=completion_stream, model="m", custom_llm_provider="hosted_vllm", logging_obj=logging_obj + ) + + +def _hosted_vllm_stream_choices( + logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] +) -> tuple[dict[str, object], ...]: + return tuple(chunk.choices[0].model_dump() for chunk in _hosted_vllm_stream(logging_obj, upstream_choices)) + + +@pytest.mark.parametrize("stop_reason", [",5", 151645]) +@pytest.mark.parametrize("final_content", [",3", ""], ids=["content-then-finish", "finish-only"]) +def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id( + logging_obj: Logging, final_content: str, stop_reason: str | int +) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]}, + { + "index": 0, + "delta": {"content": final_content}, + "finish_reason": "stop", + "stop_reason": stop_reason, + "token_ids": [20], + }, + ), + ) + + assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [stop_reason] + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice["stop_reason"]) + for choice in emitted + if choice["finish_reason"] is None and choice.get("stop_reason") is not None + ] == ([(final_content, stop_reason)] if final_content else []) + assert _joined_token_ids(emitted) == (16, 11, 17, 20) + + +def test_hosted_vllm_stream_sends_empty_content_on_token_id_only_chunks(logging_obj: Logging) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [220]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [16]}, + ), + ) + + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids")) + for choice in emitted + if choice.get("token_ids") + ] == [("", [11]), ("", [220]), ("1", [16])] + + +def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(logging_obj: Logging) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 1, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [17]}, + {"index": 1, "delta": {"content": "2"}, "finish_reason": None, "token_ids": [18]}, + ), + ) + + assert [(choice["index"], choice.get("token_ids")) for choice in emitted if choice.get("token_ids")] == [ + (1, [16]), + (0, [17]), + (1, [18]), + ] + + +def _rejects_a_two(text: str) -> bool: + return "2" not in text + + +def _choices_sent_before_rule_failure(stream: CustomStreamWrapper) -> Iterator[dict[str, object]]: + try: + yield from (chunk.choices[0].model_dump() for chunk in stream) + except MidStreamFallbackError: + return + + +def test_hosted_vllm_stream_sends_no_token_ids_ahead_of_a_failing_post_call_rule( + logging_obj: Logging, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "post_call_rules", [_rejects_a_two]) + stream: Final = _hosted_vllm_stream( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [11, 17]}, + {"index": 0, "delta": {"content": ",2"}, "finish_reason": "stop", "stop_reason": None, "token_ids": [20]}, + ), + ) + + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids")) + for choice in _choices_sent_before_rule_failure(stream) + ] == [("", None), ("1", None)] + + +def _held_back_stop_string_stream(token_ids: Mapping[str, None]) -> tuple[Mapping[str, object], ...]: + return ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "stop_reason": ",3", **token_ids}, + ) + + +def test_hosted_vllm_stream_with_null_token_ids_emits_the_same_chunks_as_without_them(logging_obj: Logging) -> None: + with_null_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({"token_ids": None})) + without_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({})) + + assert [(c["delta"], c["finish_reason"]) for c in with_null_ids] == [ + (c["delta"], c["finish_reason"]) for c in without_ids + ] + + +_TWO_CHOICE_STOP_REASONS: Final[Mapping[int, str | int]] = MappingProxyType({0: ", 3", 1: 151645}) +_TWO_CHOICE_TOKEN_IDS: Final[Mapping[int, tuple[int, int]]] = MappingProxyType({0: (16, 17), 1: (26, 27)}) + + +def _two_choice_stream(first_finished: int, finish_contents: tuple[str, str]) -> tuple[Mapping[str, object], ...]: + return ( + *({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)), + *( + {"index": index, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [ids[0]]} + for index, ids in _TWO_CHOICE_TOKEN_IDS.items() + ), + *( + { + "index": index, + "delta": {"content": content}, + "finish_reason": "stop", + "stop_reason": _TWO_CHOICE_STOP_REASONS[index], + "token_ids": [_TWO_CHOICE_TOKEN_IDS[index][1]], + } + for index, content in zip((first_finished, 1 - first_finished), finish_contents) + ), + ) + + +def _carries_another_choices_fields(choice: Mapping[str, object]) -> bool: + own_index: Final = _CHOICE_INDEX.validate_python(choice["index"]) + token_ids: Final = _TOKEN_IDS.validate_python(choice.get("token_ids")) or () + return not set(token_ids) <= set(_TWO_CHOICE_TOKEN_IDS[own_index]) or choice.get("stop_reason") not in ( + None, + _TWO_CHOICE_STOP_REASONS[own_index], + ) + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", + [("", ""), ("2", ""), ("2", "2")], + ids=["both-finish-only", "content-then-finish-only", "both-content-then-finish"], +) +def test_hosted_vllm_stream_with_two_choices_keeps_ids_and_stop_reason_on_their_own_choice( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + + assert [ + (choice["index"], choice.get("token_ids"), choice.get("stop_reason")) + for choice in emitted + if _carries_another_choices_fields(choice) + ] == [] + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", [("2", ""), ("2", "2")], ids=["content-then-finish-only", "both-content-then-finish"] +) +def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final_content_chunk( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + content_chunk_stop_reasons: Final = frozenset( + (choice["index"], choice["stop_reason"]) + for choice in emitted + if _DELTA.validate_python(choice["delta"]).get("content") and choice.get("stop_reason") is not None + ) + finished_with_content: Final = frozenset( + (index, _TWO_CHOICE_STOP_REASONS[index]) + for index, content in zip((first_finished, 1 - first_finished), finish_contents) + if content + ) + + assert content_chunk_stop_reasons == finished_with_content + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", + [("", ""), ("2", ""), ("", "2"), ("2", "2")], + ids=["both-finish-only", "content-then-finish-only", "finish-only-then-content", "both-content-then-finish"], +) +def test_hosted_vllm_stream_with_two_choices_delivers_every_token_id_and_stop_reason_of_each_choice( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + + assert { + index: _joined_token_ids([choice for choice in emitted if choice["index"] == index]) + for index in _TWO_CHOICE_TOKEN_IDS + } == _TWO_CHOICE_TOKEN_IDS + assert frozenset( + (choice["index"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None + ) == frozenset(_TWO_CHOICE_STOP_REASONS.items()) + + +def test_hosted_vllm_stream_with_two_choices_finishes_the_first_choice_before_the_other_choice_streams_on( + logging_obj: Logging, +) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + *({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)), + {"index": 0, "delta": {"content": "1"}, "finish_reason": "stop", "stop_reason": ",2", "token_ids": [16]}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [26]}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]}, + {"index": 1, "delta": {"content": "1,"}, "finish_reason": None, "token_ids": [27]}, + {"index": 1, "delta": {"content": "3"}, "finish_reason": "stop", "stop_reason": 151645, "token_ids": [28]}, + ), + ) + first_choice_finish: Final = next( + position for position, choice in enumerate(emitted) if choice["index"] == 0 and choice["finish_reason"] + ) + + assert emitted[first_choice_finish]["stop_reason"] == ",2" + assert _joined_token_ids([choice for choice in emitted[first_choice_finish:] if choice["index"] == 1]) == ( + 26, + 11, + 27, + 28, + ) + + +_CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19) + + +def _cached_vllm_completion() -> dict[str, object]: + return { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + "stop_reason": ", 5", + "token_ids": list(_CACHED_VLLM_TOKEN_IDS), + } + ], + } + + +async def _replayed_cache_hit_choices(logging_obj: Logging, replay: str) -> tuple[dict[str, object], ...]: + if replay == "async": + async_wrapper: Final = CustomStreamWrapper( + completion_stream=convert_to_streaming_response_async(_cached_vllm_completion()), + model="m", + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return tuple([chunk.choices[0].model_dump() async for chunk in async_wrapper if chunk.choices]) + sync_wrapper: Final = CustomStreamWrapper( + completion_stream=convert_to_streaming_response(_cached_vllm_completion()), + model="m", + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return tuple(chunk.choices[0].model_dump() for chunk in sync_wrapper if chunk.choices) + + +@pytest.mark.parametrize("replay", ["sync", "async"]) +@pytest.mark.asyncio +async def test_a_streamed_cache_hit_keeps_the_vllm_stop_reason_and_every_token_id( + logging_obj: Logging, replay: str +) -> None: + emitted: Final = await _replayed_cache_hit_choices(logging_obj, replay) + + assert [ + (choice["finish_reason"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None + ] == [("stop", ", 5")] + assert _joined_token_ids(emitted) == _CACHED_VLLM_TOKEN_IDS + + def _later_chunk_outputs( custom_llm_provider: str, later_chunks: Sequence[ChatCompletionChunk], logging_obj: Logging ) -> tuple[ModelResponseStream, ...]: diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 697ac76b4e5..8ef118154d1 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -638,8 +638,14 @@ def test_token_counter_tools_match_billed_prompt_for_json_tool_chat_templates( ("deepseek/deepseek-chat:free", "DeepSeek-V4-Flash-0731"), ("deepseek-reasoner", "DeepSeek-V4-Flash-0731"), ("deepseek/deepseek-flash", "DeepSeek-V4-Flash-0731"), - ("DeepSeek-V4.1-Flash", "deepseek-v3.2"), - ("accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v3.2"), + ("DeepSeek-V4.1-Flash", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("databricks-deepseek-v4-1-flash", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("deepseek-ai/deepseek-v4.1-flash-maas", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("deepseek-v4p1-flash-1017", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("grok-4.1-fast-non-reasoning", "grok-4-1-fast-reasoning"), + ("xai.grok-4.1-fast", "grok-4-1-fast-reasoning"), + ("grok-4.3-latest", "grok-4.3"), + ("xai.grok-4-3", "grok-4.3"), ("Qwen/Qwen3.5-122B-A10B", "databricks-qwen35-122b-a10b"), ("qwen3-235b-a22b", "mistral-large-latest"), ], @@ -704,6 +710,48 @@ def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: obj ) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools) +# A request with a one-function tool billed 283 prompt tokens on xAI grok-4.3 without tool_choice and 297 with a +# named one, and 357 and 358 on Azure AI Foundry grok-4-1-fast-reasoning, on 2026-10-29 +@pytest.mark.parametrize("model", ["grok-4.3", "grok-4-1-fast-reasoning"]) +def test_token_counter_counts_a_named_tool_choice_for_grok(model: str): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + named: Final = {"type": "function", "function": {"name": "get_weather"}} + + assert token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=tools, tool_choice=named + ) > token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + + +def _counted_with_strict(model: str, **strict: bool) -> int: + return token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=[_openai_tool({**_WEATHER_FUNCTION, **strict})] + ) + + +# A one-function tool billed the same prompt tokens with strict unset and false on 2026-10-29: 357 through Azure AI +# Foundry grok-4-1-fast-reasoning chat and 283 through xAI grok-4.3 +@pytest.mark.parametrize("model", ["grok-4-1-fast-reasoning", "grok-4.3"]) +def test_token_counter_renders_a_false_strict_like_an_unset_one_for_grok(model: str): + assert ( + _counted_with_strict(model, strict=False) + == _counted_with_strict(model) + < _counted_with_strict(model, strict=True) + ) + + +# A one-function tool billed 5 prompt tokens more with strict false than unset, the same as with strict true, on +# 2026-10-30: 310 against 305 through Baseten DeepSeek-V4.1-Flash and 309 against 304 through Together +@pytest.mark.parametrize( + "model", ["accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v4-flash", "qwen3.5-plus"] +) +def test_token_counter_renders_a_false_strict_like_a_true_one_for_other_chat_templates(model: str): + assert ( + _counted_with_strict(model, strict=False) + == _counted_with_strict(model, strict=True) + > _counted_with_strict(model) + ) + + @pytest.mark.parametrize( ("tool_choice", "expected_tool_choice"), [ @@ -759,9 +807,17 @@ def test_token_counter_counts_anthropic_tools_as_sent_to_chat_template_families( # Qwen3.5's chat_template.jinja opens a system turn for the tool schemas when the request has none, 5 tokens that a # request with a system message already pays, while DeepSeek V4's encoding_dsv4.py appends them to the first turn -# either way (both read on Hugging Face on 2026-10-17) +# either way (both read on Hugging Face on 2026-10-17). A system message cost the same with and without a tool on +# Fireworks deepseek-v4p1-flash, Azure AI Foundry grok-4-1-fast-reasoning and xAI grok-4.3 on 2026-10-29 @pytest.mark.parametrize( - ("model", "implicit_system_turn_tokens"), [("DeepSeek-V4-Flash-0731", 0), ("databricks-qwen35-122b-a10b", 5)] + ("model", "implicit_system_turn_tokens"), + [ + ("DeepSeek-V4-Flash-0731", 0), + ("databricks-qwen35-122b-a10b", 5), + ("accounts/fireworks/models/deepseek-v4p1-flash", 0), + ("grok-4-1-fast-reasoning", 0), + ("grok-4.3", 0), + ], ) def test_token_counter_counts_the_tools_system_turn_only_without_a_system_message( model: str, implicit_system_turn_tokens: int diff --git a/tests/unit/litellm_core_utils/test_tool_fields.py b/tests/unit/litellm_core_utils/test_tool_fields.py index 23d7de98fbc..a99278d18dd 100644 --- a/tests/unit/litellm_core_utils/test_tool_fields.py +++ b/tests/unit/litellm_core_utils/test_tool_fields.py @@ -1,4 +1,5 @@ import datetime +import logging from collections.abc import Callable, Mapping from typing import Final @@ -6,7 +7,13 @@ import pytest import litellm from litellm import Router +from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.tool_fields import ( + tool_params_without_hosted_mcp, + tools_without_hosted_mcp, + tools_without_mcp_servers, +) from litellm.utils import async_pre_call_deployment_hook _GOOD_TOOL: Final = {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}} @@ -185,3 +192,92 @@ async def test_tools_core_accepts_reach_the_deployment(tools: object) -> None: sent: Final = await async_pre_call_deployment_hook(dict(_proxy_request("keeps", tools=tools)), "acompletion") assert sent["tools"] is tools + + +def test_a_hosted_mcp_error_names_every_server_and_marks_an_unlabeled_one(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + tools_without_hosted_mcp( + tools=[ + {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}, + {"type": "mcp", "server_url": "https://mcp.example.com/mcp"}, + ], + model="gemini-3.8-flash", + route="gemini", + drop_params=False, + ) + assert "the hosted MCP tool `deepwiki`, `unlabeled` " in str(raised.value) + + +_HOSTED_MCP_TOOL: Final = {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"} + + +def test_dropping_a_hosted_mcp_tool_keeps_the_other_tool_params_and_leaves_the_caller_dict_alone() -> None: + tool_choice: Final = {"type": "function", "function": {"name": "get_weather"}} + sent: Final = { + "tools": [_HOSTED_MCP_TOOL, _GOOD_TOOL], + "tool_choice": tool_choice, + "parallel_tool_calls": False, + "temperature": 0.2, + } + params: Final = tool_params_without_hosted_mcp(sent, model="deepseek-v4-flash", route="deepseek", drop_params=True) + + assert params == { + "tools": [_GOOD_TOOL], + "tool_choice": tool_choice, + "parallel_tool_calls": False, + "temperature": 0.2, + } + assert sent["tools"] == [_HOSTED_MCP_TOOL, _GOOD_TOOL] + + +def test_dropping_every_hosted_mcp_tool_clears_the_tool_selection_params() -> None: + params: Final = tool_params_without_hosted_mcp( + {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required", "parallel_tool_calls": True, "temperature": 0.2}, + model="deepseek-v4-flash", + route="deepseek", + drop_params=True, + ) + + assert params == {"temperature": 0.2} + + +class _Records(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.messages: Final[list[str]] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.mark.parametrize( + "drop", + [ + lambda model: tools_without_hosted_mcp( + [{"type": "mcp", "server_label": "wiki\r\nFORGED", "server_url": "https://mcp.example.com/mcp"}], + model=model, + route="deepseek", + drop_params=True, + ), + lambda model: tools_without_mcp_servers( + [{"type": "url", "name": "wiki\r\nFORGED", "url": "https://mcp.example.com/mcp"}], + tools=[], + model=model, + route="bedrock", + drop_params=True, + ), + ], + ids=["hosted_mcp_tool", "mcp_servers"], +) +def test_a_drop_warning_stays_on_one_log_line(drop: Callable[[str], object]) -> None: + records: Final = _Records() + verbose_logger.addHandler(records) + try: + drop("deepseek-v4-flash\nFORGED level=CRITICAL") + finally: + verbose_logger.removeHandler(records) + + assert len(records.messages) == 1 + assert "\n" not in records.messages[0] and "\r" not in records.messages[0] + assert "FORGED level=CRITICAL" in records.messages[0] diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 95068008d25..7e2786327dc 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2,6 +2,7 @@ import json import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import chain +from collections.abc import Callable from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -13,7 +14,8 @@ from openai.types.chat import ChatCompletionChunk import litellm from litellm._uuid import uuid -from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.constants import NON_LLM_CONNECTION_TIMEOUT, RESPONSE_FORMAT_TOOL_NAME +from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call, make_sync_call from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope @@ -253,6 +255,7 @@ def test_anthropic_tool_call_error_is_mapped_from_the_provider_response(): async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop(async_only_image_fetch): http_image_url = f"http://img.example/{uuid.uuid4()}.png" https_image_url = f"https://img.example/{uuid.uuid4()}.png" + https_pdf_link = f"https://img.example/{uuid.uuid4()}/page.pdf?w=600" captured = {} def handle(request): @@ -282,6 +285,7 @@ async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop {"type": "text", "text": "What colour is this?"}, {"type": "image_url", "image_url": {"url": http_image_url}}, {"type": "image_url", "image_url": {"url": https_image_url}}, + {"type": "image_url", "image_url": {"url": https_pdf_link}}, ], } ], @@ -290,11 +294,193 @@ async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop ) assert response.choices[0].message.content == "Green" - assert async_only_image_fetch.fetched == [http_image_url] + assert async_only_image_fetch.fetched == [http_image_url, https_pdf_link] sources = [part["source"] for part in captured["body"]["messages"][0]["content"] if part["type"] == "image"] assert sources == [ {"type": "base64", "media_type": "image/png", "data": async_only_image_fetch.base64_png}, {"type": "url", "url": https_image_url}, + {"type": "url", "url": https_pdf_link}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("download_limit_mb", (0, 50), ids=("downloads-disabled", "link-unreachable")) +async def test_anthropic_async_completion_sends_a_pdf_path_link_it_cannot_download_as_a_link( + monkeypatch: pytest.MonkeyPatch, download_limit_mb: int +) -> None: + pdf_link: Final = f"https://img.example/{uuid.uuid4()}/page.pdf?fm=png" + captured: Final[dict[str, object]] = {} + + async def unreachable(client: object, url: str, **kwargs: object) -> httpx.Response: + return httpx.Response(404, request=httpx.Request("GET", url)) + + def handle(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "A chart"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + monkeypatch.setattr(image_handling, "async_safe_get", unreachable) + monkeypatch.setattr(image_handling, "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", download_limit_mb) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response: Final = await litellm.acompletion( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": pdf_link}}]}], + client=client, + api_key="sk-ant-unit-test", + ) + + assert response.choices[0].message.content == "A chart" + assert captured["body"]["messages"][0]["content"] == [{"type": "image", "source": {"type": "url", "url": pdf_link}}] + + +@pytest.mark.asyncio +async def test_anthropic_async_completion_gives_up_on_a_pdf_link_within_the_non_llm_connection_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pdf_link: Final = f"https://unreachable.example/{uuid.uuid4()}/report.pdf" + fetch_timeouts: Final[list[object]] = [] + captured: Final[dict[str, object]] = {} + + async def unreachable(client: object, url: str, **kwargs: object) -> httpx.Response: + fetch_timeouts.append(kwargs.get("timeout")) + raise httpx.ConnectTimeout("no route to host") + + def handle(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "A chart"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + monkeypatch.setattr(image_handling, "async_safe_get", unreachable) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + await litellm.acompletion( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": pdf_link}}]}], + client=client, + api_key="sk-ant-unit-test", + ) + + assert captured["body"]["messages"][0]["content"] == [{"type": "image", "source": {"type": "url", "url": pdf_link}}] + assert fetch_timeouts == [NON_LLM_CONNECTION_TIMEOUT] + + +def _anthropic_reply(captured: dict[str, object]) -> Callable[[httpx.Request], httpx.Response]: + def handle(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "A report"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + return handle + + +def _served_pdf_slice(url: str) -> httpx.Response: + return httpx.Response( + 206, content=b"%PDF-1.7\n", headers={"content-type": "binary/octet-stream"}, request=httpx.Request("GET", url) + ) + + +@pytest.mark.asyncio +async def test_anthropic_async_completion_sends_a_link_that_serves_a_pdf_as_a_url_document( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pdf_link: Final = f"https://files.example/{uuid.uuid4()}/download?file=report.pdf" + probes: Final[list[object]] = [] + captured: Final[dict[str, object]] = {} + + async def serve_pdf_slice(client: object, url: str, **kwargs: object) -> httpx.Response: + probes.append(kwargs.get("headers")) + return _served_pdf_slice(url) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_pdf_slice) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_anthropic_reply(captured))) + + await litellm.acompletion( + model="anthropic/claude-haiku-5-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": pdf_link}, "cache_control": {"type": "ephemeral"}} + ], + } + ], + client=client, + api_key="sk-ant-unit-test", + ) + + assert captured["body"]["messages"][0]["content"] == [ + {"type": "document", "source": {"type": "url", "url": pdf_link}, "cache_control": {"type": "ephemeral"}} + ] + assert probes == [{"Range": "bytes=0-1023"}] + + +def test_anthropic_sync_completion_sends_a_link_that_serves_a_pdf_as_a_url_document( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pdf_link: Final = f"https://files.example/{uuid.uuid4()}/report.pdf" + png_at_pdf_link: Final = f"https://files.example/{uuid.uuid4()}/page.pdf?fm=png" + captured: Final[dict[str, object]] = {} + + def serve(client: object, url: str, **kwargs: object) -> httpx.Response: + if url == pdf_link: + return _served_pdf_slice(url) + return httpx.Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=httpx.Request("GET", url)) + + monkeypatch.setattr(image_handling, "safe_get", serve) + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_anthropic_reply(captured)))) + + litellm.completion( + model="anthropic/claude-haiku-5-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": pdf_link}}, + {"type": "image_url", "image_url": {"url": png_at_pdf_link}}, + ], + } + ], + client=client, + api_key="sk-ant-unit-test", + ) + + assert captured["body"]["messages"][0]["content"] == [ + {"type": "document", "source": {"type": "url", "url": pdf_link}}, + {"type": "image", "source": {"type": "url", "url": png_at_pdf_link}}, ] @@ -3705,6 +3891,40 @@ def test_streamed_end_turn_carries_no_stop_details(): assert "stop_reason" not in provider_fields +def test_message_start_surfaces_input_transformations_on_the_first_chunk() -> None: + input_transformations: Final = [{"type": "thinking_block_binding", "count": 2}] + chunk: Final = { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 1}, + "input_transformations": input_transformations, + }, + } + iterator: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + + first_chunk: Final = iterator.chunk_parser(chunk=chunk) + + assert first_chunk.choices[0].delta.provider_specific_fields == {"input_transformations": input_transformations} + assert first_chunk.model == "claude-sonnet-4-6" + + +def test_message_start_without_input_transformations_has_no_provider_specific_fields() -> None: + chunk: Final = { + "type": "message_start", + "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-6", "content": []}, + } + iterator: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + + assert iterator.chunk_parser(chunk=chunk).choices[0].delta.provider_specific_fields is None + + def _anthropic_chat_request_for_bound_thinking( extra_headers: dict[str, str] | None, drop_params: bool | None = None ) -> httpx.Request: @@ -3779,3 +3999,76 @@ def test_anthropic_chat_sends_block_binding_as_is_without_drop_params(monkeypatc monkeypatch.setattr(litellm, "drop_params", False) body: Final = _anthropic_chat_body_for_bound_thinking(None) assert body["thinking"] == {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}} + + +class _AnthropicChatUpstream: + def __init__(self) -> None: + self.anthropic_beta: tuple[str, ...] = () + + def respond(self, request: httpx.Request) -> httpx.Response: + self.anthropic_beta = tuple(request.headers.get_list("anthropic-beta")) + return httpx.Response( + 200, + json={ + "id": "msg_beta_row", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + +_BETAS_FOUNDRY_REJECTS_AND_ACCEPTS: Final = "advisor-tool-2026-03-01,thinking-binding-controls-2026-08-01" + + +@pytest.mark.parametrize( + ("api_base", "forwarded_betas"), + [ + pytest.param( + "https://res.services.ai.azure.com/anthropic", + frozenset({"thinking-binding-controls-2026-08-01"}), + id="foundry", + ), + pytest.param( + "https://res.cognitiveservices.azure.com/anthropic", + frozenset({"thinking-binding-controls-2026-08-01"}), + id="foundry-cognitiveservices", + ), + pytest.param( + "https://res.openai.azure.com/anthropic", + frozenset({"thinking-binding-controls-2026-08-01"}), + id="foundry-openai", + ), + pytest.param( + "https://api.anthropic.com", + frozenset(_BETAS_FOUNDRY_REJECTS_AND_ACCEPTS.split(",")), + id="anthropic", + ), + ], +) +def test_anthropic_chat_completion_picks_the_beta_row_by_api_base( + monkeypatch: pytest.MonkeyPatch, api_base: str, forwarded_betas: frozenset[str] +) -> None: + from litellm import anthropic_beta_headers_manager + + monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") + monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None) + upstream: Final = _AnthropicChatUpstream() + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(upstream.respond))) + try: + litellm.completion( + model="anthropic/claude-sonnet-4-6", + messages=[{"role": "user", "content": "Say ok"}], + api_base=api_base, + api_key="test-key", + extra_headers={"anthropic-beta": _BETAS_FOUNDRY_REJECTS_AND_ACCEPTS}, + client=client, + ) + finally: + client.close() + + assert frozenset(",".join(upstream.anthropic_beta).split(",")) == forwarded_betas diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 636e23168cf..fc873832956 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1,4 +1,5 @@ import copy +import itertools import json import logging from collections.abc import Iterator @@ -40,6 +41,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation im VertexAIAnthropicConfig, ) from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk from litellm.types.utils import ChatCompletionToolCallChunk, ModelResponse, ServerToolUse, Usage from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @@ -4421,6 +4423,184 @@ def test_transform_request_includes_cache_control(): assert result["cache_control"] == {"type": "ephemeral"} +def _anthropic_body(messages: list, litellm_params: dict | None = None, tools: list | None = None) -> dict: + return AnthropicConfig().transform_request( + model="claude-sonnet-5-5", + messages=messages, + optional_params={"max_tokens": 100, **({"tools": tools} if tools else {})}, + litellm_params=litellm_params or {}, + headers={}, + ) + + +def test_transform_request_carries_prompt_cache_breakpoint_to_cache_control(): + body: Final = _anthropic_body( + [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are terse.", "prompt_cache_breakpoint": {"mode": "explicit"}} + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit"}}, + { + "type": "text", + "text": "Bye", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "prompt_cache_breakpoint": {"mode": "explicit"}, + }, + ], + }, + ] + ) + + assert body["system"] == [{"type": "text", "text": "You are terse.", "cache_control": {"type": "ephemeral"}}] + assert body["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "Bye", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + ], + } + ] + + +@pytest.mark.parametrize( + "marker, expected_block", + [ + ({"mode": "explicit", "ttl": "30m"}, {"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}}), + ({"mode": "implicit"}, {"type": "text", "text": "Hello"}), + ], +) +def test_transform_request_prompt_cache_breakpoint_without_claude_equivalent(marker, expected_block): + messages: Final = [{"role": "user", "content": [{"type": "text", "text": "Hello", "prompt_cache_breakpoint": marker}]}] + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + _anthropic_body(copy.deepcopy(messages)) + + body: Final = _anthropic_body(messages, litellm_params={"drop_params": True}) + assert body["messages"] == [{"role": "user", "content": [expected_block]}] + + +def test_transform_request_names_the_marker_ttl_rule_for_a_ttl_only_claude_takes(): + messages: Final = [ + {"role": "user", "content": [{"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}}]} + ] + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="optional ttl of '30m' or '12h'"): + _anthropic_body(messages) + + +def test_transform_request_keeps_the_first_and_last_prompt_cache_breakpoints_within_the_cap(): + client_marked_part: Final = {"type": "text", "text": "part 0", "cache_control": {"type": "ephemeral"}} + marked_parts: Final = [ + {"type": "text", "text": f"part {index}", "prompt_cache_breakpoint": {"mode": "explicit"}} + for index in range(1, 5) + ] + messages: Final = [{"role": "user", "content": [client_marked_part, *marked_parts]}] + tools: Final = [{"name": "lookup", "input_schema": {"type": "object"}, "cache_control": {"type": "ephemeral"}}] + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="at most 4 cache breakpoints"): + _anthropic_body(copy.deepcopy(messages), tools=copy.deepcopy(tools)) + + body: Final = _anthropic_body(messages, litellm_params={"drop_params": True}, tools=tools) + assert [part.get("cache_control") for part in body["messages"][0]["content"]] == [ + {"type": "ephemeral"}, + {"type": "ephemeral"}, + None, + None, + {"type": "ephemeral"}, + ] + assert body["tools"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_transform_request_counts_a_cache_marked_tool_search_tool_toward_the_cap(): + config: Final = AnthropicConfig() + optional_params: Final = config.map_openai_params( + non_default_params={ + "max_tokens": 100, + "tools": [ + { + "type": "tool_search_tool_regex_20251119", + "name": "tool_search_tool_regex", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + optional_params={}, + model="claude-sonnet-5-5", + drop_params=False, + ) + marked_turns: Final = [ + { + "role": "user", + "content": [{"type": "text", "text": f"turn {index}", "prompt_cache_breakpoint": {"mode": "explicit"}}], + } + for index in range(4) + ] + + def body(litellm_params: dict) -> dict: + return config.transform_request( + model="claude-sonnet-5-5", + messages=copy.deepcopy(marked_turns), + optional_params=copy.deepcopy(optional_params), + litellm_params=litellm_params, + headers={}, + ) + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="this request has 5"): + body({}) + + capped: Final = body({"drop_params": True}) + assert capped["tools"][0]["cache_control"] == {"type": "ephemeral"} + sent_parts: Final = itertools.chain.from_iterable(message["content"] for message in capped["messages"]) + assert [part.get("cache_control") for part in sent_parts] == [ + {"type": "ephemeral"}, + {"type": "ephemeral"}, + None, + {"type": "ephemeral"}, + ] + + +@pytest.mark.parametrize( + "unsent_marks", + [ + pytest.param( + [ + { + "role": "assistant", + "content": [{"type": "text", "text": "Noted.", "prompt_cache_breakpoint": {"mode": "explicit"}}], + } + ], + id="assistant_part", + ), + pytest.param( + [{"role": "user", "content": "Also this.", "prompt_cache_breakpoint": {"mode": "explicit"}}], + id="message_level", + ), + ], +) +def test_transform_request_caps_only_the_marks_claude_receives(unsent_marks: list[dict]): + marked_parts: Final = [ + {"type": "text", "text": f"part {index}", "prompt_cache_breakpoint": {"mode": "explicit"}} for index in range(4) + ] + + body: Final = _anthropic_body( + [{"role": "user", "content": marked_parts}, *unsent_marks, {"role": "user", "content": "Go on."}] + ) + + sent_marks: Final = [ + part.get("cache_control") + for part in itertools.chain.from_iterable(message["content"] for message in body["messages"]) + if isinstance(part, dict) + ] + assert sent_marks.count({"type": "ephemeral"}) == 4 + + def test_compaction_block_empty_list_not_added(): """ Test that empty compaction_blocks list is not added to provider_specific_fields. @@ -6788,6 +6968,34 @@ def test_eager_input_streaming_not_set_on_computer_use_tool(): assert "eager_input_streaming" not in mapped_tool +def _computer_tool_with_defer_loading_and_allowed_callers(tool_type: str) -> dict[str, object]: + return { + "type": tool_type, + "function": {"name": "computer", "parameters": {"display_width_px": 1024, "display_height_px": 768}}, + "defer_loading": True, + "allowed_callers": ["code_execution_20250825"], + } + + +@pytest.mark.parametrize("tool_type", ["computer_20241022", "computer_20250124"]) +def test_legacy_computer_tools_drop_defer_loading_and_allowed_callers(tool_type: str): + mapped_tool, _ = AnthropicConfig().map_tool_helper(_computer_tool_with_defer_loading_and_allowed_callers(tool_type)) + + assert mapped_tool["type"] == tool_type + assert "defer_loading" not in mapped_tool + assert "allowed_callers" not in mapped_tool + + +def test_computer_20251124_keeps_defer_loading_and_allowed_callers(): + mapped_tool, _ = AnthropicConfig().map_tool_helper( + _computer_tool_with_defer_loading_and_allowed_callers("computer_20251124") + ) + + assert mapped_tool["type"] == "computer_20251124" + assert mapped_tool["defer_loading"] is True + assert mapped_tool["allowed_callers"] == ["code_execution_20250825"] + + def test_eager_input_streaming_reaches_anthropic_request_tools(): result = AnthropicConfig().map_openai_params( non_default_params={"tools": [_eager_chat_tool(eager_input_streaming=True)], "stream": True}, @@ -7520,10 +7728,10 @@ def test_process_anthropic_headers_with_no_matching_headers(): @pytest.mark.parametrize( "computer_tool_used, prompt_caching_set, expected_beta_header", [ - (True, False, True), - (False, True, False), - (True, True, True), - (False, False, False), + ("computer_20250124", False, True), + (None, True, False), + ("computer_20250124", True, True), + (None, False, False), ], ) def test_anthropic_beta_header(computer_tool_used, prompt_caching_set, expected_beta_header): @@ -9233,6 +9441,7 @@ async def test_litellm_anthropic_prompt_caching_system( } +@pytest.mark.usefixtures("local_beta_headers_config") def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(respx_mock: respx.MockRouter): route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( return_value=httpx.Response( @@ -9271,7 +9480,7 @@ def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(res sent_betas: Final = frozenset( beta.strip() for beta in route.calls.last.request.headers["anthropic-beta"].split(",") ) - assert sent_betas == frozenset({"computer-use-2025-01-24"}) + assert sent_betas == frozenset({"computer-use-2024-10-22", "computer-use-2025-01-24"}) assert json.loads(route.calls.last.request.content)["tools"] == [ { "type": "computer_20241022", @@ -9283,6 +9492,69 @@ def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(res ] +# Tool-to-beta pairs from Anthropic's computer use docs; api.anthropic.com answered 400 to +# computer_20251124 without computer-use-2025-11-24 and 200 with it on 2026-10-31 +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("tool_type", "expected_beta"), + [ + ("computer_20241022", "computer-use-2024-10-22"), + ("computer_20250124", "computer-use-2025-01-24"), + ("computer_20251124", "computer-use-2025-11-24"), + ], +) +def test_completion_adds_the_beta_matching_the_computer_tool_version( + respx_mock: respx.MockRouter, tool_type: str, expected_beta: str +): + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=_anthropic_message_body([{"type": "text", "text": "hi"}], "end_turn")) + ) + + litellm.completion( + model="anthropic/claude-opus-5", + api_key="mock_api_key", + messages=[{"role": "user", "content": "Say hi"}], + tools=[ + { + "type": tool_type, + "function": { + "name": "computer", + "parameters": {"display_height_px": 768, "display_width_px": 1024}, + }, + } + ], + ) + + sent_betas: Final = frozenset( + beta.strip() for beta in route.calls.last.request.headers["anthropic-beta"].split(",") + ) + assert sent_betas == frozenset({expected_beta}) + + +@pytest.mark.usefixtures("local_beta_headers_config") +def test_completion_sends_no_computer_use_beta_for_an_unknown_computer_tool_version(respx_mock: respx.MockRouter): + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=_anthropic_message_body([{"type": "text", "text": "hi"}], "end_turn")) + ) + + litellm.completion( + model="anthropic/claude-opus-5", + api_key="mock_api_key", + messages=[{"role": "user", "content": "Say hi"}], + tools=[ + { + "type": "computer_20991231", + "function": { + "name": "computer", + "parameters": {"display_height_px": 768, "display_width_px": 1024}, + }, + } + ], + ) + + assert "anthropic-beta" not in route.calls.last.request.headers + + def _anthropic_message_body(content: list[dict[str, object]], stop_reason: str) -> dict[str, object]: return { "id": "msg_01XrAv7gc5tQNDuoADra7vB4", @@ -9922,6 +10194,160 @@ def test_transform_request_tool_config_point_without_tools_drops_the_key(): assert litellm_params["metadata"] == {} +_INPUT_TRANSFORMATIONS: Final = [{"type": "thinking_block_binding", "count": 2}] + + +def _input_transformations_response() -> AnthropicMessagesResponse: + return AnthropicMessagesResponse( + id="msg_input_transformations", + type="message", + role="assistant", + model="claude-sonnet-4-6", + content=[{"type": "text", "text": "ok"}], + stop_reason="end_turn", + stop_sequence=None, + usage={"input_tokens": 12, "output_tokens": 3}, + input_transformations=_INPUT_TRANSFORMATIONS, + ) + + +def test_chat_response_surfaces_input_transformations_in_provider_specific_fields() -> None: + result: Final = AnthropicConfig().transform_response( + model="claude-sonnet-4-6", + raw_response=httpx.Response(200, json=_input_transformations_response()), + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].message.provider_specific_fields["input_transformations"] == _INPUT_TRANSFORMATIONS + + +def test_chat_response_without_input_transformations_adds_no_key() -> None: + completion_response: Final = { + key: value for key, value in _input_transformations_response().items() if key != "input_transformations" + } + result: Final = AnthropicConfig().transform_response( + model="claude-sonnet-4-6", + raw_response=httpx.Response(200, json=completion_response), + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert "input_transformations" not in result.choices[0].message.provider_specific_fields + + +def test_json_mode_chat_response_keeps_input_transformations() -> None: + completion_response: Final = { + **_input_transformations_response(), + "content": [{"type": "tool_use", "id": "toolu_json", "name": RESPONSE_FORMAT_TOOL_NAME, "input": {"answer": 42}}], + "stop_reason": "tool_use", + } + result: Final = AnthropicConfig().transform_response( + model="claude-sonnet-4-6", + raw_response=httpx.Response(200, json=completion_response), + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + json_mode=True, + ) + + assert json.loads(result.choices[0].message.content) == {"answer": 42} + assert result.choices[0].message.provider_specific_fields["input_transformations"] == _INPUT_TRANSFORMATIONS + + +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} +_ROUTES_WITHOUT_THE_MCP_CONNECTOR: Final = ( + pytest.param(AmazonAnthropicClaudeConfig, "bedrock", id="bedrock_invoke"), + pytest.param(VertexAIAnthropicConfig, "vertex_ai", id="vertex_ai"), +) + + +@pytest.mark.parametrize("config_class", [AnthropicConfig, AzureAnthropicConfig, BedrockClaudePlatformConfig]) +def test_routes_with_the_mcp_connector_send_a_hosted_mcp_tool_as_an_mcp_server(monkeypatch, config_class): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=False, + ) + assert [server["url"] for server in mapped["mcp_servers"]] == [_HOSTED_MCP_TOOL["server_url"]] + assert [tool["name"] for tool in mapped["tools"]] == ["get_weather"] + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_a_hosted_mcp_tool_answers_400_on_claude_routes_without_the_mcp_connector(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert f"claude-haiku-5-5 on {route} " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_claude_routes_without_the_mcp_connector_drop_a_hosted_mcp_tool_under_drop_params( + monkeypatch, caplog, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=True, + ) + assert [tool["name"] for tool in mapped["tools"]] == ["get_weather"] + assert "mcp_servers" not in mapped + assert any( + "`deepwiki`" in record.getMessage() and f" on {route} (" in record.getMessage() for record in caplog.records + ) + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_claude_routes_without_the_mcp_connector_send_no_tools_when_drop_params_drops_every_one( + monkeypatch, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=True, + ) + assert "tools" not in mapped + assert "mcp_servers" not in mapped + assert "tool_choice" not in mapped + + _KEPT_SEARCH_CALL: Final = { "role": "assistant", "content": None, diff --git a/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py b/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py index c34cb686928..12769e78329 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py @@ -27,11 +27,14 @@ Tests cover (consolidating PRs #23706 and #22727): import os import sys +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock, patch import pytest +import litellm + # Anchor sys.path to this file's location — not the working-directory-relative # pattern Greptile flagged on PR #23706. Resolves correctly regardless of # where pytest is invoked from. @@ -67,8 +70,8 @@ def _call_prepare(extra_kwargs, model="gpt-4o", output_format=None, **overrides) system=None, temperature=None, thinking=None, - tool_choice=None, - tools=None, + tool_choice=overrides.get("tool_choice"), + tools=overrides.get("tools"), top_k=None, top_p=None, output_format=output_format, @@ -263,3 +266,112 @@ async def test_messages_layer_converted_stream_flag_still_gets_a_message_back(co ) assert response["content"][0]["text"] == "hi" + + +_DEEPWIKI_SERVER: Final = {"type": "url", "url": "https://mcp.deepwiki.com/mcp", "name": "deepwiki"} + + +@pytest.mark.parametrize( + ("mcp_servers", "named"), + [ + ([_DEEPWIKI_SERVER], "`deepwiki`"), + ([_DEEPWIKI_SERVER, {"type": "url", "url": "https://mcp.example.com/mcp"}], "`deepwiki`, `unnamed`"), + ({"name": "deepwiki"}, "`unnamed`"), + ], + ids=["named", "one-unnamed", "not-a-list"], +) +def test_mcp_servers_a_chat_route_cannot_run_answer_400_naming_each_server( + monkeypatch: pytest.MonkeyPatch, mcp_servers: object, named: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + _call_prepare( + extra_kwargs={"mcp_servers": mcp_servers, "custom_llm_provider": "gemini"}, model="gemini-3.8-flash" + ) + + assert raised.value.status_code == 400 + assert f"to gemini-3.8-flash on gemini, so the MCP server {named} would not run" in str(raised.value) + + +@pytest.mark.parametrize( + ("route", "forwards"), + [("litellm_proxy", True), ("hosted_vllm", True), ("bedrock", False), ("vertex_ai", False)], +) +def test_mcp_servers_reach_routes_that_carry_them_and_answer_400_elsewhere( + monkeypatch: pytest.MonkeyPatch, route: str, forwards: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + extra_kwargs: Final = {"mcp_servers": [_DEEPWIKI_SERVER], "custom_llm_provider": route} + if not forwards: + with pytest.raises(litellm.UnsupportedParamsError, match=f"on {route}, so the MCP server `deepwiki`"): + _call_prepare(extra_kwargs=extra_kwargs, model="some-model") + return + + completion_kwargs, _ = _call_prepare(extra_kwargs=extra_kwargs, model="some-model") + + assert completion_kwargs["mcp_servers"] == [_DEEPWIKI_SERVER] + + +_DEEPWIKI_TOOLSET: Final = {"type": "mcp_toolset", "mcp_server_name": "deepwiki"} +_WEATHER_TOOL: Final = { + "name": "get_weather", + "description": "Get the weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, +} + + +@pytest.mark.parametrize( + ("tools", "kept_tool_names"), + [([_DEEPWIKI_TOOLSET, _WEATHER_TOOL], ["get_weather"]), ([_DEEPWIKI_TOOLSET], None)], + ids=["with-function", "toolset-only"], +) +def test_drop_params_drops_the_mcp_toolset_tools_of_dropped_mcp_servers( + monkeypatch: pytest.MonkeyPatch, tools: list[dict[str, object]], kept_tool_names: list[str] | None +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + completion_kwargs, _ = _call_prepare( + extra_kwargs={"mcp_servers": [_DEEPWIKI_SERVER], "custom_llm_provider": "gemini"}, + model="gemini-3.8-flash", + tools=tools, + tool_choice={"type": "auto"}, + ) + + sent_tools: Final = completion_kwargs.get("tools") + sent_names: Final = [tool["function"]["name"] for tool in sent_tools] if sent_tools else None + assert sent_names == kept_tool_names + assert ("tool_choice" in completion_kwargs) == bool(kept_tool_names) + + +_DROP_FORMS: Final[Mapping[str, Mapping[str, object]]] = { + "request": {"drop_params": True}, + "global": {}, + "additional_drop_params": {"additional_drop_params": ["mcp_servers"]}, +} + + +@pytest.mark.parametrize("drop_params", list(_DROP_FORMS)) +def test_drop_params_drops_mcp_servers_instead_of_forwarding_them( + monkeypatch: pytest.MonkeyPatch, drop_params: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_params == "global") + completion_kwargs, _ = _call_prepare( + extra_kwargs={ + "mcp_servers": [_DEEPWIKI_SERVER], + "custom_llm_provider": "gemini", + **_DROP_FORMS[drop_params], + }, + model="gemini-3.8-flash", + ) + + assert "mcp_servers" not in completion_kwargs + assert completion_kwargs["messages"] == MESSAGES + + +@pytest.mark.parametrize("mcp_servers", [None, []], ids=["absent", "empty"]) +def test_no_mcp_servers_needs_no_drop_params(monkeypatch: pytest.MonkeyPatch, mcp_servers: object) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + completion_kwargs, _ = _call_prepare( + extra_kwargs={"mcp_servers": mcp_servers, "custom_llm_provider": "gemini"}, model="gemini-3.8-flash" + ) + + assert "mcp_servers" not in completion_kwargs diff --git a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py index d4911427e73..a57011b6410 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py @@ -204,6 +204,20 @@ def test_splitter_splits_combined_chunk_into_content_then_finish(): assert finish_chunk.choices[0].delta.content is None +def test_a_finish_on_another_choice_index_does_not_end_the_message(): + chunks: Final = [ + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hello"))]), + ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content="Hi"))]), + ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content=""), finish_reason="length")]), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")]), + ] + + events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="qwen3")) + + stop_reasons: Final = [event["delta"]["stop_reason"] for event in events if event.get("type") == "message_delta"] + assert stop_reasons == ["end_turn"] + + def test_is_combined_false_when_choices_empty(): """A metadata-only chunk with no choices is never treated as combined.""" assert _CombinedChunkSplitter._is_combined(SimpleNamespace(choices=[])) is False diff --git a/tests/unit/llms/anthropic/messages/context_management/test_inference_view.py b/tests/unit/llms/anthropic/messages/context_management/test_inference_view.py new file mode 100644 index 00000000000..8c2747592cc --- /dev/null +++ b/tests/unit/llms/anthropic/messages/context_management/test_inference_view.py @@ -0,0 +1,239 @@ +import logging +from typing import Final +from unittest.mock import patch + +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.llms.anthropic.messages.context_management.inference_view import ( + CountContextEdits, + InferenceRequestView, + InferenceRequestViewError, + inference_request_view, +) + +_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None) + + +@pytest.mark.parametrize( + "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") +) +async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = await inference_request_view( + ( + {"role": "user", "content": "first question " * 50}, + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, + {"role": "user", "content": "latest question"}, + ), + system, + "anthropic", + _NO_EDITS, + ) + + assert isinstance(view, InferenceRequestView) + assert view.messages == [{"role": "user", "content": "latest question"}] + assert "Summary of the earlier turns." in str(view.system) + assert "Be brief." in str(view.system) + + +async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + view: Final = await inference_request_view( + ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + ), + None, + "anthropic", + _NO_EDITS, + ) + + assert isinstance(view, InferenceRequestView) + assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) + assert [block["type"] for block in assistant_blocks] == ["text"] + assert snippet in assistant_blocks[0]["text"] + + +async def test_openai_view_is_the_request_unchanged() -> None: + messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) + + view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS) + + assert isinstance(view, InferenceRequestView) + assert view.messages == messages + assert view.system == "Be brief." + + +_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 +_READ_FILE_TOOL: Final[dict[str, object]] = { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_EDIT: Final[JsonValue] = { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, +} +_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]} + + +def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]: + return ( + *( + turn + for index in range(first_round, first_round + rounds) + for turn in ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}], + }, + ) + ), # comprehension-ok: flattens each round into its three turns + {"role": "user", "content": "Summarize what you read."}, + ) + + +def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]: + return tuple( + _TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"]) + ) + + +def _edits(spec: JsonValue) -> CountContextEdits: + return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,)) + + +async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None: + view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (False, False, True) + + +async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None: + spec: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT] + } + turns: Final = ( + {"role": "user", "content": "first question " * 30000}, + {"role": "assistant", "content": "first answer"}, + *_agent_turns(3), + ) + + with ( + patch( + "litellm.llms.anthropic.messages.context_management.editors.compact._read_summary_model_setting", + return_value="gemini/gemini-2.5-flash", + ), + patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary, + ): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec)) + + assert summary.call_count == 0 + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (False, False, True) + + +async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + with ( + patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")), + caplog.at_level(logging.DEBUG, logger="LiteLLM"), + ): + view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True, True) + assert any("skipping edits" in record.getMessage() for record in caplog.records) + assert not any(record.levelno >= logging.ERROR for record in caplog.records) + + +async def test_anthropic_view_counts_the_turns_after_a_compaction_block_when_an_edit_fails() -> None: + turns: Final = ( + *_agent_turns(3)[:-1], + {"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]}, + *_agent_turns(2, first_round=3), + ) + + with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True) + assert "Read files 0 to 2." in str(view.system) + + +async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None: + clear_before_compact: Final[JsonValue] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 3}, + "keep": {"type": "tool_uses", "value": 1}, + }, + {"type": "compact_20260112"}, + ] + } + turns: Final = ( + *_agent_turns(3)[:-1], + {"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]}, + *_agent_turns(2, first_round=3), + ) + + with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True) + assert "Read files 0 to 2." in str(view.system) + + +async def test_anthropic_view_leaves_a_system_that_is_not_text_blocks_and_its_turns_to_the_provider_check() -> None: + image_system: Final[JsonValue] = [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}] + + view: Final = await inference_request_view( + _agent_turns(3), image_system, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS) + ) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True, True) + assert view.system == image_system + + +async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None: + low_trigger: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] + } + + view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger)) + + assert isinstance(view, InferenceRequestViewError) + assert view.status_code == 400 diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index af3108e5072..8a924dd32d6 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -3,10 +3,13 @@ import copy import json import os import uuid +from collections.abc import Mapping from dataclasses import dataclass +from pathlib import Path from typing import Any, Dict, Final, List import httpx +import openai import responses import respx import pytest @@ -1136,6 +1139,260 @@ async def test_messages_strips_provider_prefix_exactly_once(requested_model, exp assert captured["url"] == expected_url +_GET_WEATHER_TOOL: Final = { + "name": "get_weather", + "description": "Get the current weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tools, expected_path", + [ + ([_GET_WEATHER_TOOL], "/openai/v1/responses"), + (None, "/openai/deployments/gpt-6-luna/chat/completions"), + ], + ids=["function_tool_bridged_to_responses", "no_tools_on_chat"], +) +@pytest.mark.parametrize("credentials_from_env", [False, True], ids=["credentials_in_params", "credentials_from_env"]) +async def test_azure_ai_on_azure_openai_host_streams_natively_with_the_bare_deployment( + tools: list[dict[str, object]] | None, + expected_path: str, + credentials_from_env: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.delenv("AZURE_API_BASE", raising=False) + monkeypatch.delenv("AZURE_API_KEY", raising=False) + monkeypatch.setenv("AZURE_AI_API_BASE", "https://qa-resource.cognitiveservices.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "env-api-key") + credentials: Final = ( + {} + if credentials_from_env + else {"api_base": "https://qa-resource.cognitiveservices.azure.com", "api_key": "params-api-key"} + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route(host="qa-resource.cognitiveservices.azure.com").mock( + side_effect=httpx.ConnectError("cut at the wire") + ) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "What is the weather in Berlin?"}], + model="azure_ai/gpt-6-luna", + stream=True, + tools=tools, + **credentials, + ) + + bodies: Final = [json.loads(call.request.content) for call in upstream.calls] + assert {call.request.url.path for call in upstream.calls} == {expected_path} + assert {call.request.headers["api-key"] for call in upstream.calls} == {credentials.get("api_key", "env-api-key")} + assert {body["model"] for body in bodies} == {"gpt-6-luna"} + assert all(body["stream"] is True for body in bodies) + + +@pytest.mark.asyncio +async def test_bridged_messages_keep_the_region_of_a_regional_bedrock_mantle_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(litellm.exceptions.ServiceUnavailableError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="bedrock_mantle/eu-west-1/openai.gpt-oss-120b", + api_key="test-api-key", + ) + + assert {call.request.url.host for call in upstream.calls} == {"bedrock-mantle.eu-west-1.api.aws"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["anthropic/claude-sonnet-5-5", "ragflow/gpt-4o"]) +async def test_bridged_messages_keep_the_callers_explicit_provider_over_the_model_prefix( + model: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model=model, + custom_llm_provider="openai", + api_base="https://gateway.example.com/v1", + api_key="test-api-key", + ) + + assert {call.request.url.host for call in upstream.calls} == {"gateway.example.com"} + assert {json.loads(call.request.content)["model"] for call in upstream.calls} == {model} + + +@pytest.mark.parametrize("model", ["chatgpt/gpt-5", "github_copilot/gpt-5"]) +def test_bridged_messages_reach_the_callers_provider_without_signing_in_to_the_prefix_provider( + model: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "server-openai-key") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + litellm.anthropic.messages.create( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model=model, + custom_llm_provider="openai", + ) + + assert {call.request.url.host for call in upstream.calls} == {"api.openai.com"} + assert {call.request.headers["authorization"] for call in upstream.calls} == {"Bearer server-openai-key"} + + +@pytest.mark.asyncio +async def test_bridged_messages_send_the_deployments_key_for_azure_models_served_as_openai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "server-openai-key") + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure/mistral-large-latest", + api_base="https://example-resource.services.ai.azure.com/models", + api_key="deployment-key", + ) + + assert {call.request.headers["authorization"] for call in upstream.calls} == {"Bearer deployment-key"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entra_credential", + [ + {"azure_ad_token": "entra-token"}, + {"azure_ad_token_provider": lambda: "entra-token"}, + {"tenant_id": "tenant", "client_id": "client", "client_secret": "secret"}, + {"azure_username": "user", "azure_password": "password", "client_id": "client"}, + ], +) +@pytest.mark.parametrize("api_base_from_env", [False, True], ids=["api_base_in_params", "api_base_from_env"]) +async def test_bridged_messages_keep_the_deployments_entra_credential_over_the_env_key( + monkeypatch: pytest.MonkeyPatch, entra_credential: Mapping[str, object], api_base_from_env: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key") + litellm.in_memory_llm_clients_cache.flush_cache() + api_base: Final = {} if api_base_from_env else {"api_base": "https://example-resource.cognitiveservices.azure.com"} + + with ( + patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=lambda: "entra-token"), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure + patch( + "litellm.llms.azure.common_utils.get_azure_ad_token_from_username_password", + return_value=lambda: "entra-token", + ), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure + respx.mock(assert_all_called=True) as respx_mock, + ): + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure_ai/gpt-6-luna", + tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + **api_base, + **entra_credential, + ) + + assert {call.request.headers.get("authorization") for call in upstream.calls} == {"Bearer entra-token"} + assert {call.request.headers.get("api-key") for call in upstream.calls} == {None} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "azure_env_credential, expected_header, expected_value", + [ + ({"AZURE_API_KEY": "azure-env-key"}, "api-key", "azure-env-key"), + ( + {"AZURE_TENANT_ID": "tenant", "AZURE_CLIENT_ID": "client", "AZURE_CLIENT_SECRET": "secret"}, + "authorization", + "Bearer entra-token", + ), + ], + ids=["azure_env_key", "azure_env_entra"], +) +async def test_bridged_messages_keep_the_azure_env_credential_when_the_deployment_sets_its_own_api_base( + azure_env_credential: Mapping[str, str], expected_header: str, expected_value: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + for name in ("AZURE_API_KEY", "AZURE_OPENAI_API_KEY", "AZURE_AD_TOKEN", "AZURE_TENANT_ID", "AZURE_CLIENT_ID"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AZURE_AI_API_KEY", "key-for-another-resource") + for name, value in azure_env_credential.items(): + monkeypatch.setenv(name, value) + litellm.in_memory_llm_clients_cache.flush_cache() + + with ( + patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=lambda: "entra-token"), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure + respx.mock(assert_all_called=True) as respx_mock, + ): + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure_ai/gpt-6-luna", + api_base="https://example-resource.cognitiveservices.azure.com", + tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + ) + + assert {call.request.headers.get(expected_header) for call in upstream.calls} == {expected_value} + + +@pytest.mark.asyncio +async def test_bridged_messages_still_route_through_the_litellm_proxy_when_asked( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://proxy.example.com") + monkeypatch.setenv("LITELLM_PROXY_API_KEY", "proxy-key") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key") + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure_ai/gpt-6-luna", + use_litellm_proxy=True, + ) + + assert {call.request.url.host for call in upstream.calls} == {"proxy.example.com"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "model_info, expected_url", diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py index 8fa9d534a8c..54b29ca6aa7 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py @@ -563,6 +563,36 @@ def test_drop_params_forwards_a_non_object_message_item_as_sent_on_the_native_re assert request["messages"][2]["content"][0]["content"] == [{"type": "text", "text": "sunny"}] +def test_native_response_keeps_input_transformations() -> None: + input_transformations: Final[list[JsonValue]] = [{"type": "thinking_block_binding", "count": 2}] + upstream: Final[dict[str, JsonValue]] = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 19, "output_tokens": 42}, + "input_transformations": input_transformations, + } + logging_obj: Final = Logging( + model="claude-sonnet-4-6", + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="test-call", + function_id="test-function", + ) + + response: Final = AnthropicMessagesConfig().transform_anthropic_messages_response( + model="claude-sonnet-4-6", raw_response=httpx.Response(200, json=upstream), logging_obj=logging_obj + ) + + assert response["input_transformations"] == input_transformations + + BINDING_AND_DISPLAY_BETAS: Final = "thinking-binding-controls-2026-08-01,thinking-display-updates-2026-08-18" diff --git a/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py b/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py index 73c2edbce19..c0e0ea9f803 100644 --- a/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py @@ -2,8 +2,12 @@ from importlib import import_module from typing import Final from unittest.mock import AsyncMock, patch +import httpx +import openai import pytest +import respx +import litellm from litellm.llms.anthropic.messages.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, ) @@ -56,6 +60,30 @@ def test_anthropic_messages_handler_routes_litellm_proxy_mcp_to_the_gateway(): assert result is not None +@pytest.mark.asyncio +async def test_anthropic_messages_with_mcp_reaches_an_env_only_deployment_with_its_env_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key") + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure_ai/gpt-6-luna", + tools=[{**MCP_REFERENCE, "server_url": "litellm_proxy"}], + ) + + assert {call.request.url.host for call in upstream.calls} == {"example-resource.cognitiveservices.azure.com"} + assert {call.request.headers.get("api-key") for call in upstream.calls} == {"azure-env-key"} + assert not any(b"litellm_proxy" in call.request.content for call in upstream.calls) + + def test_anthropic_messages_handler_skips_the_gateway_on_recursion(): """The gateway's own follow-up call must not re-enter the gateway.""" with patch( diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py index 382917fc756..c06c9d721ac 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py @@ -764,6 +764,57 @@ async def test_recomputed_cost_of_a_served_bridge_response_prices_the_reasoning_ assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) +_DEEPWIKI_SERVER: Final = {"type": "url", "url": "https://mcp.deepwiki.com/mcp", "name": "deepwiki"} + + +def test_mcp_servers_the_responses_bridge_cannot_run_answer_400(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as rejected: + _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER]}, + max_tokens=1024, + messages=MESSAGES, + ) + + assert rejected.value.status_code == 400 + assert "to gpt-5.6 on openai, so the MCP server `deepwiki` would not run" in rejected.value.message + + +@pytest.mark.parametrize( + "drop_form", + [{"drop_params": True}, {"additional_drop_params": ["mcp_servers"]}], + ids=["drop_params", "additional_drop_params"], +) +def test_drop_params_keeps_mcp_servers_off_the_responses_request( + monkeypatch: pytest.MonkeyPatch, drop_form: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + responses_kwargs: Final = _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER], **drop_form}, + max_tokens=1024, + messages=MESSAGES, + ) + + assert "mcp_servers" not in responses_kwargs + assert responses_kwargs["input"] + + +def test_drop_params_drops_the_mcp_toolset_tools_of_dropped_mcp_servers(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + weather_tool: Final = {"name": "get_weather", "input_schema": {"type": "object", "properties": {}}} + responses_kwargs: Final = _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER], "drop_params": True}, + max_tokens=1024, + messages=MESSAGES, + tools=[{"type": "mcp_toolset", "mcp_server_name": "deepwiki"}, weather_tool], + ) + + assert [tool["name"] for tool in responses_kwargs["tools"]] == ["get_weather"] + + TEXT_RESPONSE: Final = { "id": "resp_cleared", "object": "response", diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index cf0b8a39865..0ee5d52d50d 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -10,6 +10,7 @@ Verifies that: - ANTHROPIC_API_KEY / ANTHROPIC_API_BASE take precedence over their aliases. """ +import itertools import json import logging import os @@ -4728,3 +4729,142 @@ def test_tools_without_strict_drops_only_the_strict_key_and_counts_the_strict_to assert tools_without_strict(tools) == ToolsWithoutStrict( tools=({"name": "a"}, {"name": "b"}, {"name": "c"}), dropped_count=1 ) + + +@pytest.mark.parametrize("registration_order", [("caching-off", "unflagged"), ("unflagged", "caching-off")]) +def test_auto_injection_reads_a_deployments_own_prompt_caching_flag_not_a_siblings( + registration_order: tuple[str, ...], +) -> None: + from litellm import Router + from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control + + model: Final = "anthropic.claude-qa-cache-model" + litellm.register_model({model: {"litellm_provider": "bedrock", "mode": "chat", "supports_prompt_caching": True}}) + model_infos: Final = { + "caching-off": {"id": "caching-off-id", "supports_prompt_caching": False}, + "unflagged": {"id": "unflagged-id"}, + } + Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": f"bedrock/{model}", "aws_region_name": "us-east-1"}, + "model_info": model_infos[name], + } + for name in registration_order + ] + ) + + assert {name: supports_anthropic_cache_control(model, "bedrock", model_infos[name]) for name in registration_order} == { + "caching-off": False, + "unflagged": True, + } + + +def test_auto_injection_leaves_an_unmapped_claude_deployment_uncached_instead_of_raising() -> None: + from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control + + assert supports_anthropic_cache_control("claude-qa-unmapped-model", "anthropic", {"id": "unmapped-id"}) is False + + +def test_auto_injection_honors_a_register_model_opt_out_made_under_the_provider_prefixed_name() -> None: + from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control + from litellm.utils import catalog_with_runtime_registrations + + model: Final = next( + key + for key, entry in catalog_with_runtime_registrations().items() + if entry.get("litellm_provider") == "anthropic" and entry.get("supports_prompt_caching") is True + ) + litellm.register_model({f"anthropic/{model}": {"supports_prompt_caching": False}}) + + assert supports_anthropic_cache_control(model, "anthropic", {"id": "plain-id"}) is False + + +@pytest.mark.parametrize("route_prefix", ["bedrock/", "bedrock/converse/"]) +def test_auto_injection_honors_a_register_model_opt_out_made_under_a_bedrock_route_prefixed_name(route_prefix: str) -> None: + from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control + from litellm.utils import catalog_with_runtime_registrations + + model: Final = next( + key + for key, entry in catalog_with_runtime_registrations().items() + if entry.get("litellm_provider") == "bedrock_converse" + and entry.get("supports_prompt_caching") is True + and key.startswith("us.anthropic.claude") + ) + litellm.register_model({f"{route_prefix}{model}": {"supports_prompt_caching": False}}) + + assert ( + supports_anthropic_cache_control(model, "bedrock", {"id": "plain-id"}), + supports_anthropic_cache_control(model, "bedrock"), + ) == (False, False) + + +@pytest.mark.parametrize( + ("declared", "catalog_caching", "expected"), + [("true", False, True), (1, False, True), ("false", True, False), (0, True, False), ("maybe", True, True)], +) +def test_auto_injection_reads_a_string_or_numeric_deployment_prompt_caching_flag_as_the_bool_it_spells( + declared: object, catalog_caching: bool, expected: bool +) -> None: + from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control + + model: Final = f"claude-qa-lax-flag-model-{catalog_caching}" + litellm.register_model( + {model: {"litellm_provider": "anthropic", "mode": "chat", "supports_prompt_caching": catalog_caching}} + ) + + assert ( + supports_anthropic_cache_control(model, "anthropic", {"id": "lax-id", "supports_prompt_caching": declared}) + is expected + ) + + +@pytest.mark.parametrize("reload_cost_map", [False, True], ids=["before-reload", "after-reload"]) +@pytest.mark.parametrize( + ("registrations", "expected"), + [ + ((("", True), ("anthropic/", False)), False), + ((("anthropic/", False), ("", True)), True), + ((("", False), ("anthropic/", True), ("", False)), False), + ], + ids=["prefixed-opt-out-last", "bare-opt-in-last", "bare-opt-out-registered-again-last"], +) +def test_the_later_of_a_bare_and_a_prefixed_register_model_entry_wins_for_routed_and_sdk_calls( + registrations: tuple[tuple[str, bool], ...], expected: bool, reload_cost_map: bool +) -> None: + from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control + from litellm.utils import catalog_with_runtime_registrations, reapply_runtime_model_cost_registrations + + model: Final = next( + key + for key, entry in catalog_with_runtime_registrations().items() + if entry.get("litellm_provider") == "anthropic" and entry.get("supports_prompt_caching") is True + ) + for prefix, flag in registrations: + litellm.register_model({f"{prefix}{model}": {"supports_prompt_caching": flag}}) + if reload_cost_map: + reapply_runtime_model_cost_registrations() + + assert ( + supports_anthropic_cache_control(model, "anthropic", {"id": "plain-id"}), + supports_anthropic_cache_control(model, "anthropic"), + ) == (expected, expected) + + +def test_a_register_model_opt_out_for_one_claude_model_leaves_every_other_routed_models_catalog_flag() -> None: + from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control + from litellm.utils import catalog_with_runtime_registrations + + opted_out, untouched = itertools.islice( + ( + key + for key, entry in catalog_with_runtime_registrations().items() + if entry.get("litellm_provider") == "anthropic" and entry.get("supports_prompt_caching") is True + ), + 2, + ) + litellm.register_model({f"anthropic/{opted_out}": {"supports_prompt_caching": False}}) + + assert supports_anthropic_cache_control(untouched, "anthropic", {"id": "plain-id"}) is True diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index aa7fa2727b4..13c17ebd2bf 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -50,20 +50,30 @@ COMPACTED_HISTORY = [ "config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig) ) @pytest.mark.parametrize( - ("image_url", "source"), + ("image_url", "block"), ( - ("data:image/png;base64,aW1hZ2U=", {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}), - ({"url": "data:image/png;base64,aW1hZ2U="}, {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}), + ( + "data:image/png;base64,aW1hZ2U=", + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ), + ( + {"url": "data:image/png;base64,aW1hZ2U="}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ), ( {"url": "data:image/png;base64,aW1hZ2U=", "format": "image/jpeg", "detail": "high"}, - {"type": "base64", "media_type": "image/jpeg", "data": "aW1hZ2U="}, + {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "aW1hZ2U="}}, + ), + ( + "data:application/pdf;base64,JVBERi0=", + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}, ), ), ) def test_count_translates_openai_images_without_mutating_input( config_type: type[AnthropicCountTokensConfig], image_url: str | dict[str, JsonValue], - source: dict[str, JsonValue], + block: dict[str, JsonValue], ) -> None: cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"} messages: Final[list[dict[str, JsonValue]]] = [{ @@ -81,7 +91,7 @@ def test_count_translates_openai_images_without_mutating_input( "model": "claude-opus-5-5", "messages": [{ "role": "user", "content": [ {"type": "text", "text": "Count this image"}, - {"type": "image", "source": source, "cache_control": cache_control}, + {**block, "cache_control": cache_control}, ], }], } diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index f6dae2c304f..f93968eb2e6 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -75,6 +75,42 @@ def test_map_openai_params_with_preview_api_version(): ) +def _map_tool_choice_on_pre_tool_choice_api_version(tool_choice: dict[str, object]) -> dict[str, object]: + return AzureOpenAIConfig().map_openai_params( + {"tool_choice": tool_choice, "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + {}, + "gpt-4o", + drop_params=True, + api_version="2023-07-01-preview", + ) + + +def test_map_openai_params_drop_params_still_drops_a_named_tool_choice_on_old_api_version(): + mapped: Final = _map_tool_choice_on_pre_tool_choice_api_version( + {"type": "function", "function": {"name": "get_time"}} + ) + assert "tool_choice" not in mapped + + +def test_map_openai_params_drop_params_never_drops_allowed_tools_on_old_api_version(): + with pytest.raises(litellm.UnsupportedParamsError, match="allowed_tools") as exc: + _map_tool_choice_on_pre_tool_choice_api_version( + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}} + ) + assert exc.value.status_code == 400 + + +def test_map_openai_params_forwards_allowed_tools_on_current_api_version(): + tool_choice: Final[dict[str, object]] = { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + } + mapped: Final = AzureOpenAIConfig().map_openai_params( + {"tool_choice": tool_choice}, {}, "gpt-4o", drop_params=True, api_version="2025-04-01-preview" + ) + assert mapped["tool_choice"] == tool_choice + + def test_transform_request_hoists_tool_message_image(): """Azure builds its request via convert_to_azure_openai_messages without the OpenAIGPTConfig._transform_messages pipeline, so transform_request must hoist @@ -806,3 +842,89 @@ def test_azure_keeps_a_later_developer_message_in_place_as_system(config: AzureO ) assert [message["role"] for message in messages] == ["system", "user", "assistant", "system", "user"] assert messages[3]["content"] == "Answer with exactly one word." + + +_MARKED_AZURE_PARTS: Final = ( + {"type": "text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}, + {"type": "text", "text": "Implicit", "prompt_cache_breakpoint": {"mode": "implicit"}}, +) +_TRIMMED_AZURE_PARTS: Final = [ + {"type": "text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "text", "text": "Implicit"}, +] + + +@pytest.fixture +def registered_azure_breakpoint_reader(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem( + litellm.model_cost, + "azure/unit-breakpoint-reader", + {"litellm_provider": "azure", "mode": "chat", "supports_prompt_cache_breakpoint": True}, + ) + monkeypatch.setitem( + litellm.model_cost, + "azure/unit-breakpoint-ignorer", + {"litellm_provider": "azure", "mode": "chat", "supports_prompt_cache_breakpoint": False}, + ) + + +def _azure_request_parts( + config: AzureOpenAIConfig, model: str, litellm_params: dict[str, object] +) -> object: + request: Final = config.transform_request( + model=model, + messages=[{"role": "user", "content": list(_MARKED_AZURE_PARTS)}], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + return request["messages"][0]["content"] + + +@pytest.mark.parametrize("config", [AzureOpenAIConfig(), AzureOpenAIGPT5Config()]) +@pytest.mark.parametrize( + "model, litellm_params", + [ + ("unit-breakpoint-reader", {"drop_params": True}), + ("my-azure-alias", {"drop_params": True, "base_model": "azure/unit-breakpoint-reader"}), + ("my-azure-alias", {"drop_params": True, "model_info": {"base_model": "azure/unit-breakpoint-reader"}}), + ("my-azure-alias", {"drop_params": True}), + ], + ids=("deployment-name", "litellm-params-base-model", "model-info-base-model", "unknown-deployment"), +) +def test_azure_drop_params_cuts_part_breakpoints_down_to_what_azure_chat_accepts( + registered_azure_breakpoint_reader: None, + config: AzureOpenAIConfig, + model: str, + litellm_params: dict[str, object], +) -> None: + """Azure OpenAI chat answered 400 for a content part's breakpoint ttl and for implicit mode on 2026-10-29, + the same as OpenAI chat completions.""" + assert _azure_request_parts(config, model, litellm_params) == _TRIMMED_AZURE_PARTS + + +@pytest.mark.parametrize("litellm_params", [{}, {"drop_params": False}], ids=("no-drop-params", "drop-params-false")) +def test_azure_sends_part_breakpoints_as_written_without_drop_params( + registered_azure_breakpoint_reader: None, litellm_params: dict[str, object] +) -> None: + assert _azure_request_parts(AzureOpenAIConfig(), "unit-breakpoint-reader", litellm_params) == list( + _MARKED_AZURE_PARTS + ) + + +@pytest.mark.parametrize( + "model, litellm_params", + [ + ("unit-breakpoint-ignorer", {"drop_params": True}), + ("unit-breakpoint-reader", {"drop_params": True, "base_model": "azure/unit-breakpoint-ignorer"}), + ("my-azure-alias", {"drop_params": True, "model_info": {"base_model": "azure/gpt-4.1"}}), + ], + ids=("known-non-reader", "base-model-known-non-reader", "model-info-base-model-known-non-reader"), +) +def test_azure_drop_params_strips_part_breakpoints_from_a_known_non_reader( + registered_azure_breakpoint_reader: None, model: str, litellm_params: dict[str, object] +) -> None: + assert _azure_request_parts(AzureOpenAIConfig(), model, litellm_params) == [ + {"type": "text", "text": "Stable prefix"}, + {"type": "text", "text": "Implicit"}, + ] diff --git a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index f1db76c1e21..6745421b8f8 100644 --- a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -1,4 +1,7 @@ import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -259,6 +262,21 @@ async def test_construct_url_v1_protocol(): assert url.count("/realtime") == 1 +def test_construct_url_translation_protocol(): + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + + url = AzureOpenAIRealtime()._construct_url( + api_base="https://my-endpoint.openai.azure.com", + model="translate-deployment", + api_version=None, + realtime_protocol="GA", + query_params={"model": "translate-deployment"}, + realtime_mode="translation", + ) + + assert url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime/translations?model=translate-deployment" + + @pytest.mark.asyncio @pytest.mark.parametrize("protocol", ["ga", "Ga", "gA", "V1", "v1", "GA"]) async def test_construct_url_case_insensitive_protocol(protocol): @@ -752,3 +770,32 @@ def test_construct_url_with_a_trailing_slash_api_base_matches_the_url_without_it ) assert url_for("https://gw.example.com/azure/") == url_for("https://gw.example.com/azure") + + +@dataclass(frozen=True, slots=True) +class _ClientWebSocketWithHeaders: + scope: Mapping[str, object] + + +@pytest.mark.parametrize("pinned_protocol", ("GA", "ga", "v1")) +@pytest.mark.parametrize( + ("client_headers", "bridges_beta_client_to_ga"), + ( + pytest.param(((b"openai-beta", b"realtime=v1"),), True, id="beta-client"), + pytest.param((), False, id="ga-client"), + ), +) +def test_a_deployment_pinned_to_ga_bridges_only_a_beta_client( + pinned_protocol: str, client_headers: tuple[tuple[bytes, bytes], ...], bridges_beta_client_to_ga: bool +) -> None: + from litellm.llms.azure.realtime.handler import AzureRealtimeProtocol, azure_realtime_protocol_for_client + + protocol: Final = azure_realtime_protocol_for_client( + pinned_protocol, + model="my-realtime-deployment", + realtime_mode="realtime", + query_params=None, + websocket=_ClientWebSocketWithHeaders(scope={"headers": client_headers}), + ) + + assert protocol == AzureRealtimeProtocol(name=pinned_protocol, bridges_beta_client_to_ga=bridges_beta_client_to_ga) diff --git a/tests/unit/llms/azure/realtime/test_http_transformation.py b/tests/unit/llms/azure/realtime/test_http_transformation.py index 714ce9558cf..dae8ed10d1d 100644 --- a/tests/unit/llms/azure/realtime/test_http_transformation.py +++ b/tests/unit/llms/azure/realtime/test_http_transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import httpx @@ -7,7 +8,11 @@ import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler _DATED_PREVIEW_API_VERSION: Final = "2025-04-01-preview" -_API_BASES: Final = ("https://my.openai.azure.com", "https://my.openai.azure.com/", "https://my.openai.azure.com/openai") +_API_BASES: Final = ( + "https://my.openai.azure.com", + "https://my.openai.azure.com/", + "https://my.openai.azure.com/openai", +) class _AzureThatRecordsTheRequest: @@ -55,3 +60,62 @@ async def test_azure_sdp_exchange_targets_the_ga_v1_surface_with_the_ephemeral_k assert str(request.url) == "https://my.openai.azure.com/openai/v1/realtime/calls" assert request.headers.get("authorization") == "Bearer ek_upstream" assert "api-key" not in request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("api_base", "url"), + ( + ("https://gw.example.com/azure/v1", "https://gw.example.com/azure/v1/openai/realtime/transcription_sessions"), + ("https://my.openai.azure.com/openai", "https://my.openai.azure.com/openai/realtime/transcription_sessions"), + ("https://my.openai.azure.com/openai/v1", "https://my.openai.azure.com/openai/realtime/transcription_sessions"), + ), +) +async def test_azure_transcription_session_strips_only_a_trailing_openai_path(api_base: str, url: str) -> None: + azure: Final = _AzureThatRecordsTheRequest() + await litellm.acreate_realtime_transcription_session( + model="azure/gpt-realtime-whisper-deployment", + api_base=api_base, + api_key="azure-key", + api_version=_DATED_PREVIEW_API_VERSION, + client=azure.client(), + ) + (request,) = azure.requests + assert str(request.url) == f"{url}?api-version={_DATED_PREVIEW_API_VERSION}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller_headers", "sent_headers"), + ( + pytest.param( + {"Authorization": "Bearer caller-token"}, + {"authorization": "Bearer caller-token", "api-key": None}, + id="caller-authorization", + ), + pytest.param( + {"authorization": "Bearer caller-token"}, + {"authorization": "Bearer caller-token", "api-key": None}, + id="caller-lowercase-authorization", + ), + pytest.param({"api-key": "caller-key"}, {"authorization": None, "api-key": "caller-key"}, id="caller-api-key"), + pytest.param({}, {"authorization": "Bearer entra-token", "api-key": None}, id="no-caller-auth"), + ), +) +async def test_azure_entra_fallback_never_overrides_the_callers_own_auth_header( + caller_headers: Mapping[str, str], + sent_headers: Mapping[str, str | None], + monkeypatch: pytest.MonkeyPatch, + no_ambient_azure_credentials: None, +) -> None: + monkeypatch.delenv("AZURE_API_KEY", raising=False) + monkeypatch.setenv("AZURE_AD_TOKEN", "entra-token") + azure: Final = _AzureThatRecordsTheRequest() + await litellm.acreate_realtime_client_secret( + model="azure/gpt-realtime-deployment", + api_base="https://my.openai.azure.com", + extra_headers=caller_headers, + client=azure.client(), + ) + (request,) = azure.requests + assert {name: request.headers.get(name) for name in sent_headers} == sent_headers diff --git a/tests/unit/llms/azure/test_audio_transcriptions.py b/tests/unit/llms/azure/test_audio_transcriptions.py index c5d403b3c38..f86fae494ef 100644 --- a/tests/unit/llms/azure/test_audio_transcriptions.py +++ b/tests/unit/llms/azure/test_audio_transcriptions.py @@ -4,6 +4,7 @@ from typing import Final import httpx import pytest +import respx from openai import AsyncAzureOpenAI, AzureOpenAI import litellm @@ -160,3 +161,70 @@ async def test_azure_transcription_parses_response_formats( ) assert response.text == expected_text + + +@pytest.mark.parametrize("api_version", ["2024-06-01", "2025-03-01-preview", "v1", "latest", "preview"]) +async def test_azure_router_transcription_sends_the_default_api_version( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + api_version: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("AZURE_API_VERSION", raising=False) + route: Final = respx_mock.route(host="transcribe-route.openai.azure.com").mock( + return_value=httpx.Response(200, json={"text": "hello"}) + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "transcriber", + "litellm_params": { + "model": "azure/gpt-transcribe", + "api_key": "test-key", + "api_base": "https://transcribe-route.openai.azure.com", + "api_version": api_version, + }, + } + ] + ) + + response: Final = await router.atranscription(model="transcriber", file=_audio_file()) + + sent_url: Final = route.calls.last.request.url + assert sent_url.path == "/openai/deployments/gpt-transcribe/audio/transcriptions" + assert sent_url.params.get("api-version") == litellm.AZURE_DEFAULT_API_VERSION + assert response.text == "hello" + + +@pytest.mark.parametrize( + ("api_version", "expected_api_version"), + [ + ("v1", litellm.AZURE_DEFAULT_API_VERSION), + ("latest", litellm.AZURE_DEFAULT_API_VERSION), + ("preview", litellm.AZURE_DEFAULT_API_VERSION), + ("2025-03-01-preview", "2025-03-01-preview"), + ], +) +def test_azure_sync_transcription_uses_the_deployments_route( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + api_version: str, + expected_api_version: str, +) -> None: + monkeypatch.delenv("AZURE_API_VERSION", raising=False) + route: Final = respx_mock.route(host="transcribe-route.openai.azure.com").mock( + return_value=httpx.Response(200, json={"text": "hello"}) + ) + + response: Final = litellm.transcription( + model="azure/gpt-transcribe", + file=_audio_file(), + api_key="test-key", + api_base="https://transcribe-route.openai.azure.com", + api_version=api_version, + ) + + sent_url: Final = route.calls.last.request.url + assert sent_url.path == "/openai/deployments/gpt-transcribe/audio/transcriptions" + assert sent_url.params.get("api-version") == expected_api_version + assert response.text == "hello" diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 3f9da8be2d4..0a41861e088 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -491,6 +491,11 @@ def test_default_max_retries_env_var_reaches_azure_sdk_client(): "avector_store_search", "acreate_skill", "acreate_interaction", + "acreate_realtime_client_secret", + "acreate_realtime_transcription_session", + "acreate_realtime_translation_client_secret", + "arealtime_calls", + "arealtime_translation_calls", ] ], ) diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 7807975be24..faa33e559ca 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -1,5 +1,7 @@ import json import operator +import threading +from collections.abc import Mapping from typing import Final import httpx @@ -8,8 +10,9 @@ import respx from pydantic import ValidationError import litellm +from litellm import anthropic_beta_headers_manager from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountResponse API_BASE: Final = "https://my-resource.services.ai.azure.com" COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens" @@ -268,3 +271,92 @@ async def test_count_sends_the_tool_change_beta_when_the_history_adds_a_tool(mon assert result is not None assert result.total_tokens == 541 assert route.calls.last.request.headers["anthropic-beta"] == "mid-conversation-tool-changes-2026-07-01,token-counting-2024-11-01" + + +async def _betas_azure_inference_sends(context_management: Mapping[str, object]) -> frozenset[str]: + with respx.mock: + route: Final = respx.post(f"{API_BASE}/anthropic/v1/messages").mock( + return_value=httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 2}, + }, + ) + ) + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=[{"role": "user", "content": "Summarize the thread."}], + model="azure_ai/claude-sonnet-5-5", + api_key="azure-key", + api_base=API_BASE, + context_management=context_management, + ) + return frozenset(filter(None, route.calls.last.request.headers.get("anthropic-beta", "").split(","))) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize("edit_type", ["clear_tool_uses_20250919", "compact_20260112"]) +async def test_count_forwards_the_context_edits_with_the_betas_azure_inference_sends( + monkeypatch: pytest.MonkeyPatch, edit_type: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + spec: Final = {"edits": [{"type": edit_type}]} + inference_betas: Final = await _betas_azure_inference_sends(spec) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 88})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-5-5", + messages=[{"role": "user", "content": "Summarize the thread."}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-5-5", + context_management=CountContextManagement(spec=spec, drop_invalid_edits=False), + ) + assert result is not None + assert result.total_tokens == 88 + assert json.loads(route.calls.last.request.content)["context_management"] == spec + assert frozenset(route.calls.last.request.headers["anthropic-beta"].split(",")) == inference_betas | { + "token-counting-2024-11-01" + } + + +class _FetchThreads: + seen: tuple[int, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.seen = (*self.seen, threading.get_ident()) + return httpx.Response(200, json={}) + + +@pytest.mark.asyncio +async def test_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None) + fetch_threads: Final = _FetchThreads() + with respx.mock: + respx.get(litellm.anthropic_beta_headers_url).mock(side_effect=fetch_threads) + respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 88})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-5-5", + messages=[{"role": "user", "content": "Summarize the thread."}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-5-5", + context_management=CountContextManagement( + spec={"edits": [{"type": "clear_tool_uses_20250919"}]}, drop_invalid_edits=False + ), + ) + assert result is not None + assert len(fetch_threads.seen) == 1 + assert threading.get_ident() not in fetch_threads.seen diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 9f68a4fb73f..b1d52c9dd51 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -1271,6 +1271,12 @@ def test_azure_ai_falls_back_to_the_model_name_when_the_base_model_row_does_not_ {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}, ), ({"type": "ephemeral"}, {"drop_params": True}, {"type": "text", "text": "stable"}), + ({"mode": "implicit"}, {"drop_params": True}, {"type": "text", "text": "stable"}), + ( + {"mode": "explicit", "ttl": "1h"}, + {"drop_params": True}, + {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ), ( {"mode": "explicit", "ttl": "30m"}, {}, diff --git a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py index 908e727bc4c..d0c1c4a6fb2 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py @@ -5,9 +5,15 @@ Every azure_ai route must authenticate with an Entra ID token when no API key is instead of requiring an API key. """ +import json +import queue +from collections.abc import Mapping +from typing import Final from unittest.mock import patch +import httpx import pytest +import respx import litellm from litellm.llms.azure_ai.common_utils import ( @@ -16,6 +22,8 @@ from litellm.llms.azure_ai.common_utils import ( has_azure_entra_params, resolve_azure_ai_agent_auth_header, ) +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.custom_httpx.http_handler import HTTPHandler ENTRA_PARAMS = {"azure_ad_token": "entra-token"} @@ -290,3 +298,159 @@ async def test_agent_auth_header_is_the_entra_bearer(): headers = await resolve_azure_ai_agent_auth_header({"azure_ad_token": "entra-token"}) assert headers == {"Authorization": "Bearer entra-token"} + + +OPENAI_CHAT_COMPLETION: Final = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "gpt-6.1-sol", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +OPENAI_EMBEDDING: Final = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "cohere-embed-v3-english", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, +} +ANTHROPIC_MESSAGE: Final = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "OK"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} + + +class _PreCallOptionalParamsCapture(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.optional_params: tuple[dict[str, object], ...] = () + + def log_pre_api_call(self, model: str, messages: list[object], kwargs: dict[str, dict[str, object]]) -> None: + self.optional_params = (*self.optional_params, kwargs["optional_params"]) + + +@pytest.mark.parametrize( + ("model", "provider_response"), + [ + pytest.param("azure_ai/gpt-6.1-sol", OPENAI_CHAT_COMPLETION, id="openai-shaped"), + pytest.param("azure_ai/claude-sonnet-5-5", ANTHROPIC_MESSAGE, id="claude"), + ], +) +def test_completion_sends_the_entra_token_only_as_the_bearer_header( + model: str, provider_response: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +): + capture: Final = _PreCallOptionalParamsCapture() + monkeypatch.setattr(litellm, "callbacks", [capture]) + sent: Final[queue.Queue[httpx.Request]] = queue.Queue() + + def respond(request: httpx.Request) -> httpx.Response: + sent.put_nowait(request) + return httpx.Response(200, json=provider_response) + + response: Final = litellm.completion( + model=model, + api_base="https://example-resource.services.ai.azure.com", + azure_ad_token="entra-token", + messages=[{"role": "user", "content": "hi"}], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + ) + + request: Final = sent.get_nowait() + assert request.headers["Authorization"] == "Bearer entra-token" + assert "azure_ad_token" not in json.loads(request.content) + assert response.choices[0].message.content == "OK" + assert capture.optional_params + assert "entra-token" not in json.dumps(capture.optional_params, default=str) + + +@pytest.mark.respx(assert_all_called=True) +def test_embedding_sends_the_entra_token_only_as_the_bearer_header(respx_mock: respx.MockRouter): + route: Final = respx_mock.post(url__startswith="https://example-resource.services.ai.azure.com").mock( + return_value=httpx.Response(200, json=OPENAI_EMBEDDING) + ) + + litellm.embedding( + model="azure_ai/cohere-embed-v3-english", + api_base="https://example-resource.services.ai.azure.com", + azure_ad_token="entra-token", + input=["hello"], + ) + + request: Final = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-token" + assert "azure_ad_token" not in json.loads(request.content) + + +IMAGE_RESPONSE: Final = {"created": 1, "data": [{"b64_json": "aGk="}]} +PNG_BYTES: Final = b"\x89PNG\r\n\x1a\n" + b"\x00" * 32 + + +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + "extra_kwargs", + [{}, {"extra_body": {"azure_ad_token": "entra-token", "client_secret": "entra-token"}}], + ids=["kwarg", "extra_body"], +) +def test_image_edit_sends_the_entra_token_only_as_the_bearer_header( + respx_mock: respx.MockRouter, extra_kwargs: dict[str, dict[str, str]] +): + route: Final = respx_mock.post(url__startswith="https://example-resource.services.ai.azure.com").mock( + return_value=httpx.Response(200, json=IMAGE_RESPONSE) + ) + + litellm.image_edit( + model="azure_ai/FLUX.1-Kontext-pro", + api_base="https://example-resource.services.ai.azure.com", + api_version="2025-04-01-preview", + azure_ad_token="entra-token", + image=PNG_BYTES, + prompt="add a hat", + **extra_kwargs, + ) + + request: Final = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-token" + assert b"entra-token" not in request.content + + +@pytest.mark.respx(assert_all_called=True) +def test_image_generation_sends_the_entra_token_only_as_the_bearer_header(respx_mock: respx.MockRouter): + route: Final = respx_mock.post(url__startswith="https://example-resource.services.ai.azure.com").mock( + return_value=httpx.Response(200, json=IMAGE_RESPONSE) + ) + + litellm.image_generation( + model="azure_ai/FLUX-1.1-pro", + api_base="https://example-resource.services.ai.azure.com", + azure_ad_token="entra-token", + prompt="a red circle", + ) + + request: Final = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-token" + assert "azure_ad_token" not in json.loads(request.content) + + +@pytest.mark.respx(assert_all_called=True) +def test_transcription_sends_the_entra_token_only_as_the_bearer_header(respx_mock: respx.MockRouter): + route: Final = respx_mock.post(url__startswith="https://example-resource.services.ai.azure.com").mock( + return_value=httpx.Response(200, json={"text": "hi"}) + ) + + litellm.transcription( + model="azure_ai/whisper", + api_base="https://example-resource.services.ai.azure.com", + api_version="2024-06-01", + azure_ad_token="entra-token", + file=("hello.wav", b"RIFF" + b"\x00" * 32), + ) + + request: Final = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-token" + assert b"entra-token" not in request.content diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index 041c625eaef..2ef202762ff 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -12,6 +12,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( BaseDecisionsConfig, DecisionsInputLimits, + DecisionsResponseValidationError, ir_to_systemone_request, ir_to_systemone_response, parse_systemone_response, @@ -21,6 +22,7 @@ from litellm.llms.openai.decisions.transformation import openai_request_to_ir from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig from litellm.types.decisions import ( MAX_DECISION_QUESTIONS, + DecisionsIRChoiceAnswer, DecisionsIRRefusal, DecisionsIRRequest, DecisionsIRScoreAnswer, @@ -347,6 +349,49 @@ def test_a_box_answer_without_confidence_reaches_the_caller_unchanged() -> None: assert response.model_dump(mode="json") == payload +@pytest.mark.parametrize( + ("question", "answer"), + [ + ( + {"type": "choice", "criteria": {"yes": "Yes", "no": "No"}}, + {"type": "choice", "choice": "yes", "probabilities": {"yes": 0.9, "no": 0.1}}, + ), + ( + {"type": "score", "criteria": ["low", "high"]}, + {"type": "score", "score": 0.5, "legend": {"0": "low", "1": "high"}, "probabilities": {"0": 0.5, "1": 0.5}}, + ), + ], +) +@pytest.mark.parametrize("confidence", ({}, {"confidence": None})) +@pytest.mark.parametrize( + ("provider", "optional"), (("nace", False), ("fireworks_ai", False), ("parallel_ai", True)) +) +def test_only_a_provider_without_confidence_may_omit_it( + question: Mapping[str, object], + answer: Mapping[str, object], + confidence: Mapping[str, object], + provider: str, + optional: bool, +) -> None: + config: Final = ProviderConfigManager.get_provider_decisions_config("decision", LlmProviders(provider)) + assert config is not None + request: Final = systemone_request_to_ir( + DecisionsRequestBody.model_validate({"state": {"text": "Example"}, "questions": {"asked": question}}) + ) + upstream: Final = httpx.Response( + 200, json={"mode": "fast", "model": "decision", "answers": {"asked": {**answer, **confidence}}} + ) + + if optional: + parsed: Final = config.transform_decisions_response("decision", provider, upstream, request) + assert isinstance(parsed.answers[0], (DecisionsIRChoiceAnswer, DecisionsIRScoreAnswer)) + assert parsed.answers[0].confidence is None + return + with pytest.raises(DecisionsResponseValidationError) as raised: + config.transform_decisions_response("decision", provider, upstream, request) + assert raised.value.status_code == 500 + + @pytest.mark.parametrize("provider", ("nace", "openrouter", "fireworks_ai", "typesafe")) @pytest.mark.parametrize("legend", (["low", "high"], {"0": "low", "1": "high"})) def test_score_legends_preserve_labels_probabilities_and_metadata(provider: str, legend: object) -> None: diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index d40abb4860b..34a4cf5225b 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -3,7 +3,7 @@ import copy import json import uuid from types import SimpleNamespace -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import MagicMock, patch @@ -14,11 +14,13 @@ from pydantic import BaseModel, ConfigDict # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. import litellm +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +from tests.unit.conftest import NativeStructuredOutputKey from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE ONE_PIXEL_PNG = base64.b64decode( @@ -635,7 +637,6 @@ def test_shared_defs_reach_native_output_once_instead_of_inlined_per_use(set_nat assert "tools" not in result - class Relabel(BaseModel): model_config = ConfigDict(frozen=True) @@ -672,6 +673,7 @@ def test_nested_pydantic_response_format_reaches_native_output_with_refs_that_re assert item_ref.removeprefix("#/$defs/") in sent["$defs"] assert "tools" not in result + def test_output_config_format_forwarded_for_bedrock_chat_invoke_request(local_model_cost_map): """Bedrock Invoke chat path forwards ``output_config.format`` alongside effort for models with native structured-output support (Claude Sonnet 4.6).""" @@ -1050,6 +1052,64 @@ async def test_bedrock_invoke_claude_async_completion_inlines_document_url_sourc } in captured["body"]["messages"][0]["content"] +def test_bedrock_invoke_claude_sync_completion_downloads_a_pdf_link_once_and_sends_it_as_a_base64_document(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import image_handling + + pdf_url = f"https://docs.example/{uuid.uuid4()}/report.pdf" + pdf_bytes = b"%PDF-1.7\n" + b"0" * 64 + fetched = [] + captured = {} + + def serve_pdf(client, url, **kwargs): + fetched.append(url) + return httpx.Response( + 200, content=pdf_bytes, headers={"content-type": "application/pdf"}, request=httpx.Request("GET", url) + ) + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "A lease"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + monkeypatch.setattr(image_handling, "safe_get", serve_pdf) + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(handle)) + + response = litellm.completion( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this document?"}, + {"type": "image_url", "image_url": {"url": pdf_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "A lease" + assert fetched == [pdf_url] + assert captured["body"]["messages"][0]["content"][1] == { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": base64.b64encode(pdf_bytes).decode()}, + } + + @pytest.mark.parametrize( "model, expected_betas", [ @@ -1384,6 +1444,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="session") def event_loop(): try: @@ -1393,6 +1454,7 @@ def event_loop(): yield loop loop.close() + @pytest.fixture(scope="function") def setup_and_teardown(event_loop): import litellm @@ -1426,6 +1488,7 @@ def setup_and_teardown(event_loop): if pending: event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "set_verbose": getattr(litellm, "set_verbose", False), @@ -1497,6 +1560,7 @@ IN WITNESS WHEREOF, the parties have executed this Agreement. * 8 ) # Repeat to ensure we have enough tokens (need 1024+ for Claude models) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBedrockAnthropicPromptCachingRegression: """ @@ -1659,6 +1723,7 @@ class TestBedrockAnthropicPromptCachingRegression: if "anthropic_beta" in additional_fields: assert "prompt-caching-2024-07-31" not in additional_fields["anthropic_beta"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBedrockAnthropic1MContextRegression: """ @@ -1844,6 +1909,7 @@ class TestBedrockAnthropic1MContextRegression: assert "context-1m-2025-08-07" in beta_headers assert "computer-use-2024-10-22" in beta_headers + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBedrockAnthropicCombinedRegressions: """ @@ -2010,6 +2076,225 @@ def test_invoke_uses_opt_in_native_output_is_false_when_native_output_did_not_co ) +@pytest.mark.parametrize( + ("forced_tool_use", "map_flag", "model_info", "thinking_params", "native"), + [ + (False, None, None, {}, True), + (False, None, None, {"thinking": {"type": "adaptive"}}, False), + (False, True, None, {"thinking": {"type": "adaptive"}}, True), + (False, False, None, {}, False), + (False, None, {"supports_bedrock_native_structured_output": False}, {}, False), + (True, True, None, {}, False), + ], + ids=[ + "cannot-force-tool", + "cannot-force-tool-thinking", + "cannot-force-tool-thinking-key-true", + "cannot-force-tool-key-false", + "cannot-force-tool-deployment-opt-out", + "can-force-tool", + ], +) +def test_invoke_model_that_cannot_force_the_json_tool_gets_native_output_without_thinking_unless_opted_out( + set_native_structured_output: Callable[[str, NativeStructuredOutputKey], None], + monkeypatch: pytest.MonkeyPatch, + forced_tool_use: bool, + map_flag: bool | None, + model_info: Mapping[str, object] | None, + thinking_params: Mapping[str, object], + native: bool, +) -> None: + model: Final = "us.anthropic.claude-haiku-5-5" + monkeypatch.setitem(litellm.model_cost[model], "supports_forced_tool_use", forced_tool_use) + set_native_structured_output(model, "bedrock" if map_flag else None) + if map_flag is None: + monkeypatch.delitem(litellm.model_cost[model], "supports_bedrock_native_structured_output") + + config: Final = AmazonAnthropicClaudeConfig(model_info=model_info) + optional_params: Final = config.map_openai_params( + non_default_params={"response_format": _CITY_RESPONSE_FORMAT, "max_tokens": 2048, **thinking_params}, + optional_params={}, + model=model, + drop_params=False, + ) + body: Final = config.transform_request( + model=model, + messages=[{"role": "user", "content": "Capital of France?"}], + optional_params=optional_params, + litellm_params={} if model_info is None else {"model_info": model_info}, + headers={}, + ) + + assert (body.get("output_config", {}).get("format", {}).get("schema") == _CITY_SCHEMA) is native + assert ("tools" in body) is not native + + +@pytest.mark.parametrize( + "caller_schema_param", + [ + lambda schema: {"output_format": {"type": "json_schema", "schema": schema}}, + lambda schema: {"output_config": {"format": {"type": "json_schema", "schema": schema}}}, + ], + ids=["output_format", "output_config.format"], +) +@pytest.mark.parametrize( + ("model", "forced_tool_use", "native_key"), + [("us.anthropic.claude-opus-4-8", True, "legacy"), ("us.anthropic.claude-haiku-5-5", False, None)], + ids=["native-key", "cannot-force-tool"], +) +def test_invoke_caller_schema_drops_unsupported_constraints_before_going_native( + set_native_structured_output: Callable[[str, NativeStructuredOutputKey], None], + monkeypatch: pytest.MonkeyPatch, + model: str, + forced_tool_use: bool, + native_key: NativeStructuredOutputKey, + caller_schema_param: Callable[[Mapping[str, object]], dict[str, object]], +) -> None: + monkeypatch.setitem(litellm.model_cost[model], "supports_forced_tool_use", forced_tool_use) + set_native_structured_output(model, native_key) + if native_key is None: + monkeypatch.delitem(litellm.model_cost[model], "supports_bedrock_native_structured_output") + schema: Final = { + "type": "object", + "properties": {"age": {"type": "integer", "minimum": 0}}, + "required": ["age"], + "additionalProperties": False, + } + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=model, + messages=[{"role": "user", "content": "How old is the oldest tree?"}], + optional_params={**caller_schema_param(schema), "max_tokens": 2048}, + litellm_params={}, + headers={}, + ) + + sent_age: Final = body["output_config"]["format"]["schema"]["properties"]["age"] + assert sent_age["type"] == "integer" + assert "minimum" not in sent_age + + +BEDROCK_HAIKU_45_BASE_ID: Final = "anthropic.claude-haiku-4-5-20251001-v1:0" + + +@pytest.mark.parametrize( + "model", + [ + f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + f"invoke/{BEDROCK_HAIKU_45_BASE_ID}", + ], +) +def test_bedrock_invoke_claude_default_max_tokens_comes_from_base_model(local_model_cost_map, model: str): + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, ( + f"default max_tokens should be the base model's cap, got {body['max_tokens']}" + ) + + +def test_bedrock_invoke_claude_explicit_max_tokens_is_kept(local_model_cost_map): + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={"max_tokens": 123}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == 123, f"explicit max_tokens must not be overridden, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_takes_the_deployments_own_cap(local_model_cost_map): + deployment_cap: Final = 8192 + assert litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) != deployment_cap + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"model_info": {"max_output_tokens": deployment_cap}}, + headers={}, + ) + + assert body["max_tokens"] == deployment_cap, f"the deployment's own cap should win, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_ignores_a_sibling_deployments_cap(local_model_cost_map): + regional_id: Final = f"us.{BEDROCK_HAIKU_45_BASE_ID}" + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + litellm.register_model(model_cost={regional_id: {"max_output_tokens": 50}}, persist_across_reloads=False) + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/{regional_id}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"a sibling's cap on the shared id leaked, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_yields_to_anthropic_config(local_model_cost_map): + configured_max_tokens: Final = 1000 + assert litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) != configured_max_tokens + + body: Final = AmazonAnthropicClaudeConfig( + read_configured_max_tokens=lambda: configured_max_tokens + ).transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == configured_max_tokens, ( + f"AnthropicConfig(max_tokens) was ignored, got {body['max_tokens']}" + ) + + +@pytest.mark.parametrize(("configured_max_tokens", "deployment_cap"), [(100000, 50), (1000, 8192)]) +def test_bedrock_invoke_claude_default_max_tokens_is_the_lower_of_anthropic_config_and_the_deployments_cap( + local_model_cost_map, configured_max_tokens: int, deployment_cap: int +): + body: Final = AmazonAnthropicClaudeConfig( + read_configured_max_tokens=lambda: configured_max_tokens + ).transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"model_info": {"max_output_tokens": deployment_cap}}, + headers={}, + ) + + assert body["max_tokens"] == min(configured_max_tokens, deployment_cap), f"got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_reads_an_inference_profile_arns_base_model(local_model_cost_map): + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model="invoke/anthropic/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"base_model": BEDROCK_HAIKU_45_BASE_ID}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"an ARN should take its base_model's cap, got {body['max_tokens']}" + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("drop_params", [True, False]) def test_bedrock_chat_invoke_adds_the_binding_beta_so_block_binding_is_kept( diff --git a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py index 154be9f0370..154a102fe34 100644 --- a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py +++ b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -2,6 +2,8 @@ import base64 import json import uuid from types import SimpleNamespace +from typing import Final +from unittest.mock import MagicMock import httpx import pytest @@ -243,3 +245,54 @@ async def test_mantle_claude_async_completion_streams_a_wrapped_json_schema_answ chunks = tuple([chunk async for chunk in stream]) assert _streamed_answer(chunks) == BARE_STREAMED_ANSWER + + +def _json_tool_message(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-opus-5-5", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "json_tool_call", "input": RISK_ANSWER}], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 9}, + }, + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("model", ("bedrock_mantle/anthropic.claude-opus-5-5", "bedrock/mantle/anthropic.claude-opus-5-5")) +def test_mantle_claude_that_cannot_force_a_tool_keeps_json_schema_on_the_json_tool(model: str) -> None: + respond: Final = MagicMock(wraps=_json_tool_message) + + response: Final = litellm.completion( + **{**_json_schema_stream_kwargs(model), "stream": False}, + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + ) + + sent: Final = tuple(json.loads(call.args[0].content) for call in respond.call_args_list) + assert json.loads(response.choices[0].message.content) == RISK_ANSWER + assert [(body.get("output_config", {}).get("format"), [tool["name"] for tool in body["tools"]]) for body in sent] == [ + (None, ["json_tool_call"]) + ] + + +def test_mantle_claude_default_max_tokens_comes_from_the_mantle_row(local_model_cost_map): + from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig + + mapped_cap: Final = litellm.get_max_tokens("bedrock_mantle/anthropic.claude-haiku-4-5") + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonMantleConfig().transform_request( + model="mantle/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"Mantle should send its mapped cap, got {body['max_tokens']}" diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 9ab7060d766..7c19ef5b751 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -4,7 +4,7 @@ import json import logging import os import re -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping from typing import Final from unittest.mock import MagicMock, Mock, patch @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_tools_pt, ) from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import CONVERTED_SYSTEM_NOTE +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.anthropic.chat.transformation import ( NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES, NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS, @@ -31,6 +32,7 @@ from litellm.llms.bedrock.common_utils import BEDROCK_NATIVE_OUTPUT_MAX_PROPERTI from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.bedrock import ContentBlock, ConverseTokenUsageBlock, RequestObject from litellm.types.llms.openai import AllMessageValues +from tests.unit.conftest import NativeStructuredOutputKey def test_transform_usage(): @@ -4981,7 +4983,10 @@ ADAPTIVE_THINKING: Final = {"type": "adaptive"} def _converse_request_for_json_schema( - model: str, stream: bool, thinking_params: Mapping[str, object] | None = None + model: str, + stream: bool, + thinking_params: Mapping[str, object] | None = None, + model_info: Mapping[str, object] | None = None, ) -> tuple[Mapping[str, object], RequestObject]: config = AmazonConverseConfig() optional_params = config.map_openai_params( @@ -4989,12 +4994,13 @@ def _converse_request_for_json_schema( optional_params={}, model=model, drop_params=False, + model_info=model_info, ) request = config._transform_request( model=model, messages=[{"role": "user", "content": "Name the capital of France."}], optional_params=optional_params, - litellm_params={}, + litellm_params={} if model_info is None else {"model_info": model_info}, headers={}, ) return optional_params, request @@ -5061,6 +5067,115 @@ def test_bedrock_native_output_key_applies_only_when_thinking_is_requested( assert ("toolConfig" in request) is not native +_FORCED_JSON_TOOL_CHOICE: Final = {"tool": {"name": "json_tool_call"}} + + +@pytest.mark.parametrize( + ("forced_tool_use", "map_flag", "model_info", "thinking_params", "native", "tool_choice"), + [ + (False, None, None, None, True, None), + (False, None, None, {"thinking": {"type": "adaptive"}}, False, None), + (False, True, None, {"thinking": {"type": "adaptive"}}, True, None), + (False, False, None, None, False, None), + (False, None, {"supports_bedrock_native_structured_output": False}, None, False, None), + (True, True, None, None, False, _FORCED_JSON_TOOL_CHOICE), + ], + ids=[ + "cannot-force-tool", + "cannot-force-tool-thinking", + "cannot-force-tool-thinking-key-true", + "cannot-force-tool-key-false", + "cannot-force-tool-deployment-opt-out", + "can-force-tool", + ], +) +def test_a_model_that_cannot_force_the_json_tool_gets_native_output_without_thinking_unless_opted_out( + set_native_structured_output: Callable[[str, NativeStructuredOutputKey], None], + monkeypatch: pytest.MonkeyPatch, + forced_tool_use: bool, + map_flag: bool | None, + model_info: Mapping[str, object] | None, + thinking_params: Mapping[str, object] | None, + native: bool, + tool_choice: Mapping[str, object] | None, +) -> None: + model: Final = "us.anthropic.claude-haiku-5-5" + monkeypatch.setitem(litellm.model_cost[model], "supports_forced_tool_use", forced_tool_use) + set_native_structured_output(model, "bedrock" if map_flag else None) + if map_flag is None: + monkeypatch.delitem(litellm.model_cost[model], "supports_bedrock_native_structured_output") + + _, request = _converse_request_for_json_schema(model, False, thinking_params, model_info=model_info) + + assert ("outputConfig" in request) is native + assert ("toolConfig" in request) is not native + assert request.get("toolConfig", {}).get("toolChoice") == tool_choice + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("model", "thinking_params", "native"), + [ + ("us.anthropic.claude-opus-5-5", None, True), + ("us.anthropic.claude-opus-5-5", {"thinking": {"type": "adaptive"}}, False), + ("global.anthropic.claude-fable-5-1", None, True), + ("global.anthropic.claude-fable-5-1", {"reasoning_effort": "low"}, False), + ("us-gov.anthropic.claude-opus-5-5", None, False), + ("us-gov.anthropic.claude-fable-5-1", None, False), + ("us-gov-west-1/anthropic.claude-opus-5-5", None, False), + ("converse/us-gov-east-1/anthropic.claude-fable-5-1", None, False), + ("us-east-1/anthropic.claude-opus-5-5", None, True), + ], +) +def test_models_whose_thinking_stops_under_native_output_keep_the_json_tool_only_when_thinking( + model: str, thinking_params: Mapping[str, object] | None, native: bool +) -> None: + _, request = _converse_request_for_json_schema(model, False, thinking_params) + + assert ("outputConfig" in request) is native + assert ("toolConfig" in request) is not native + + +def _converse_capital_answer(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": '{"city":"Paris","country":"France"}'}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 12, "totalTokens": 13}, + }, + request=request, + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "model", ["us.anthropic.claude-opus-5-5", "global.anthropic.claude-sonnet-5-5", "global.anthropic.claude-fable-5-1"] +) +def test_streamed_json_schema_on_a_model_that_cannot_force_a_tool_still_calls_converse_without_streaming( + model: str, +) -> None: + respond: Final = MagicMock(wraps=_converse_capital_answer) + + stream: Final = litellm.completion( + model=f"bedrock/{model}", + messages=[{"role": "user", "content": "Name the capital of France."}], + response_format=CAPITAL_RESPONSE_FORMAT, + stream=True, + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + ) + + assert isinstance(stream, CustomStreamWrapper) + streamed: Final = "".join(chunk.choices[0].delta.content or "" for chunk in stream) + sent: Final = respond.call_args.args[0] + assert sent.url.path.endswith("/converse") + assert "textFormat" in json.loads(sent.content)["outputConfig"] + assert json.loads(streamed) == {"city": "Paris", "country": "France"} + + def test_adaptive_thinking_is_kept_next_to_json_schema_on_a_native_output_profile(set_native_structured_output): model: Final = "us.anthropic.claude-opus-5" set_native_structured_output(model, "bedrock") @@ -9958,6 +10073,461 @@ def test_bedrock_prompt_caching_message(messages, expected_cache_control): assert "cachePoint" not in json.dumps(transformed_messages) +def test_converse_carries_prompt_cache_breakpoint_to_cache_point(): + request: Final = litellm.AmazonConverseConfig()._transform_request( + model="anthropic.claude-sonnet-5-5", + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are terse.", "prompt_cache_breakpoint": {"mode": "explicit"}} + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit"}}], + }, + ], + optional_params={}, + litellm_params={}, + ) + + assert request["system"] == [{"text": "You are terse."}, {"cachePoint": {"type": "default"}}] + assert request["messages"] == [ + {"role": "user", "content": [{"text": "Hello"}, {"cachePoint": {"type": "default"}}]} + ] + + +def test_converse_keeps_the_first_and_last_prompt_cache_breakpoints_within_the_cap(): + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": f"part {index}", "prompt_cache_breakpoint": {"mode": "explicit"}} + for index in range(5) + ], + } + ] + + def transform(litellm_params: dict) -> dict: + return litellm.AmazonConverseConfig()._transform_request( + model="anthropic.claude-sonnet-5-5", + messages=copy.deepcopy(messages), + optional_params={}, + litellm_params=litellm_params, + ) + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="at most 4 cache breakpoints"): + transform({}) + + cache_point: Final = {"cachePoint": {"type": "default"}} + assert transform({"drop_params": True})["messages"] == [ + { + "role": "user", + "content": [ + {"text": "part 0"}, + cache_point, + {"text": "part 1"}, + cache_point, + {"text": "part 2"}, + {"text": "part 3"}, + cache_point, + {"text": "part 4"}, + cache_point, + ], + } + ] + + +_CONVERSE_CAP_TOOL: Final[dict[str, object]] = { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object", "properties": {}}}, +} + + +@pytest.mark.parametrize( + "tool_cache_params", + [ + {"tools": [{**_CONVERSE_CAP_TOOL, "cache_control": {"type": "ephemeral"}}]}, + { + "tools": [_CONVERSE_CAP_TOOL], + "cache_control_injection_points": [{"location": "tool_config", "control": {"type": "ephemeral"}}], + }, + ], + ids=["client_tool_cache_control", "tool_config_injection_point"], +) +def test_converse_reserves_the_tool_cache_point_within_the_cap(tool_cache_params: dict[str, object]): + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": f"part {index}", "prompt_cache_breakpoint": {"mode": "explicit"}} + for index in range(4) + ], + } + ] + + def transform(litellm_params: dict) -> dict: + return litellm.AmazonConverseConfig()._transform_request( + model="anthropic.claude-sonnet-5-5", + messages=copy.deepcopy(messages), + optional_params=copy.deepcopy(tool_cache_params), + litellm_params=litellm_params, + ) + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="this request has 5"): + transform({}) + + request: Final = transform({"drop_params": True}) + cache_point: Final = {"cachePoint": {"type": "default"}} + assert request["messages"][0]["content"] == [ + {"text": "part 0"}, + cache_point, + {"text": "part 1"}, + cache_point, + {"text": "part 2"}, + {"text": "part 3"}, + cache_point, + ] + assert request["toolConfig"]["tools"][-1] == cache_point + + + +@pytest.mark.parametrize( + "unsent_tool_params", + [ + { + "tools": [ + _CONVERSE_CAP_TOOL, + {"type": "tool_search_tool_regex_20251119", "name": "find", "cache_control": {"type": "ephemeral"}}, + ] + }, + {"tools": [_CONVERSE_CAP_TOOL, {"type": "web_search", "cache_control": {"type": "ephemeral"}}]}, + { + "tools": [{"type": "web_search"}], + "cache_control_injection_points": [{"location": "tool_config", "control": {"type": "ephemeral"}}], + }, + ], + ids=["tool_search_tool", "built_in_tool", "tool_config_point_with_no_tool_sent"], +) +@pytest.mark.parametrize("litellm_params", [{}, {"drop_params": True}], ids=["strict", "drop_params"]) +def test_converse_leaves_the_cap_to_markers_when_a_marked_tool_is_not_sent( + unsent_tool_params: dict[str, object], litellm_params: dict[str, object] +): + request: Final = litellm.AmazonConverseConfig()._transform_request( + model="anthropic.claude-sonnet-5-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": f"part {index}", "prompt_cache_breakpoint": {"mode": "explicit"}} + for index in range(4) + ], + } + ], + optional_params=copy.deepcopy(unsent_tool_params), + litellm_params=dict(litellm_params), + ) + + cache_point: Final = {"cachePoint": {"type": "default"}} + assert request["messages"][0]["content"] == [ + {"text": "part 0"}, + cache_point, + {"text": "part 1"}, + cache_point, + {"text": "part 2"}, + cache_point, + {"text": "part 3"}, + cache_point, + ] + assert cache_point not in request.get("toolConfig", {}).get("tools", []) + +def _converse_tool_turn(*result_parts: dict[str, object]) -> list[dict[str, object]]: + return [ + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": list(result_parts)}, + ] + + +def _marked_text_parts(count: int) -> list[dict[str, object]]: + return [ + {"type": "text", "text": f"part {index}", "prompt_cache_breakpoint": {"mode": "explicit"}} + for index in range(count) + ] + + +def _marked_tool_result_parts() -> list[dict[str, object]]: + return [ + {"type": "text", "text": f"result {index}", "prompt_cache_breakpoint": {"mode": "explicit"}} + for index in range(2) + ] + + +def test_converse_counts_one_cache_point_for_a_tool_result_with_several_marks(): + request: Final = litellm.AmazonConverseConfig()._transform_request( + model="anthropic.claude-sonnet-5-5", + messages=[ + {"role": "user", "content": _marked_text_parts(3)}, + *_converse_tool_turn(*_marked_tool_result_parts()), + ], + optional_params={"tools": [_CONVERSE_CAP_TOOL]}, + litellm_params={}, + ) + + assert len(_collect_cache_points(request["messages"])) == 4 + + +def test_converse_drops_marks_until_a_merged_tool_result_fits_the_cap(): + messages: Final = [ + {"role": "user", "content": _marked_text_parts(4)}, + *_converse_tool_turn(*_marked_tool_result_parts()), + ] + + def transform(litellm_params: dict) -> dict: + return litellm.AmazonConverseConfig()._transform_request( + model="anthropic.claude-sonnet-5-5", + messages=copy.deepcopy(messages), + optional_params={"tools": [_CONVERSE_CAP_TOOL]}, + litellm_params=litellm_params, + ) + + with pytest.raises(litellm.UnsupportedParamsError, match="this request has 5"): + transform({}) + + request: Final = transform({"drop_params": True}) + cache_point: Final = {"cachePoint": {"type": "default"}} + assert request["messages"][0]["content"] == [ + {"text": "part 0"}, + cache_point, + {"text": "part 1"}, + cache_point, + {"text": "part 2"}, + cache_point, + {"text": "part 3"}, + ] + assert len(_collect_cache_points(request["messages"][1:])) == 1 + + +def test_converse_drops_prompt_cache_breakpoint_from_a_nova_tool_result(): + messages: Final = [ + { + "role": "system", + "content": [{"type": "text", "text": "You are terse.", "prompt_cache_breakpoint": {"mode": "explicit"}}], + }, + {"role": "user", "content": "Look it up."}, + *_converse_tool_turn(*_marked_tool_result_parts()), + ] + + def transform(litellm_params: dict) -> dict: + return litellm.AmazonConverseConfig()._transform_request( + model="us.amazon.nova-lite-v1:0", + messages=copy.deepcopy(messages), + optional_params={"tools": [_CONVERSE_CAP_TOOL]}, + litellm_params=litellm_params, + ) + + with pytest.raises(litellm.UnsupportedParamsError, match="tool result"): + transform({}) + + request: Final = transform({"drop_params": True}) + assert request["system"] == [{"text": "You are terse."}, {"cachePoint": {"type": "default"}}] + assert _collect_cache_points(request["messages"]) == [] + + +def test_converse_translates_prompt_cache_breakpoint_on_an_unmapped_arn_only_when_flagged(local_model_cost_map): + from litellm import Router + + flagged_arn: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude-flagged" + unflagged_arn: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/mistral" + Router( + model_list=[ + { + "model_name": "claude-profile", + "litellm_params": {"model": f"bedrock/{flagged_arn}", "aws_region_name": "us-east-1"}, + "model_info": {"supports_prompt_cache_breakpoint": True}, + }, + { + "model_name": "mistral-profile", + "litellm_params": {"model": f"bedrock/{unflagged_arn}", "aws_region_name": "us-east-1"}, + }, + ] + ) + + def user_content(model: str, ttl: dict) -> list: + return litellm.AmazonConverseConfig()._transform_request( + model=model, + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit", **ttl}}], + } + ], + optional_params={}, + litellm_params={}, + )["messages"][0]["content"] + + assert user_content(flagged_arn, {}) == [{"text": "Hello"}, {"cachePoint": {"type": "default"}}] + assert user_content(unflagged_arn, {}) == [{"text": "Hello"}] + assert user_content(unflagged_arn, {"ttl": "30m"}) == [{"text": "Hello"}] + + + +def test_converse_prompt_cache_breakpoint_follows_the_arn_flag_over_its_base_model(local_model_cost_map): + from litellm import Router + + opted_in_arn: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude-next" + opted_out_arn: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude-opted-out" + Router( + model_list=[ + { + "model_name": "opted-in", + "litellm_params": {"model": f"bedrock/{opted_in_arn}", "aws_region_name": "us-east-1"}, + "model_info": {"supports_prompt_cache_breakpoint": True}, + }, + { + "model_name": "opted-out", + "litellm_params": {"model": f"bedrock/{opted_out_arn}", "aws_region_name": "us-east-1"}, + "model_info": {"supports_prompt_cache_breakpoint": False}, + }, + ] + ) + + def user_content(model: str, base_model: str, ttl: dict) -> list: + return litellm.AmazonConverseConfig()._transform_request( + model=model, + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit", **ttl}}], + } + ], + optional_params={}, + litellm_params={"base_model": base_model}, + )["messages"][0]["content"] + + assert user_content(opted_in_arn, "anthropic.claude-unmapped-v1:0", {}) == [ + {"text": "Hello"}, + {"cachePoint": {"type": "default"}}, + ] + assert user_content(opted_out_arn, "anthropic.claude-sonnet-5-5", {"ttl": "30m"}) == [{"text": "Hello"}] + + +def test_converse_translates_prompt_cache_breakpoint_only_for_families_that_take_cache_points(local_model_cost_map): + litellm.register_model( + {"zai.glm-next-v1:0": {"litellm_provider": "bedrock_converse", "mode": "chat", "supports_prompt_caching": True}} + ) + + def user_content(model: str, mark: dict) -> list: + return litellm.AmazonConverseConfig()._transform_request( + model=model, + messages=[{"role": "user", "content": [{"type": "text", "text": "Hello", **mark}]}], + optional_params={}, + litellm_params={}, + )["messages"][0]["content"] + + marker: Final = {"prompt_cache_breakpoint": {"mode": "explicit"}} + client_mark: Final = {"cache_control": {"type": "ephemeral"}} + cache_point: Final = {"cachePoint": {"type": "default"}} + assert user_content("zai.glm-next-v1:0", marker) == [{"text": "Hello"}] + assert user_content("zai.glm-next-v1:0", client_mark) == [{"text": "Hello"}, cache_point] + assert user_content("us.anthropic.claude-next-v1:0", marker) == [{"text": "Hello"}, cache_point] + assert user_content("us.amazon.nova-next-v1:0", marker) == [{"text": "Hello"}, cache_point] + + +def test_converse_prompt_cache_breakpoint_on_guarded_user_text_keeps_its_cache_point(local_model_cost_map): + guardrail: Final = {"guardrailConfig": {"guardrailIdentifier": "gr-1", "guardrailVersion": "DRAFT"}} + + def request(system_marks: int) -> dict: + return litellm.AmazonConverseConfig()._transform_request( + model="anthropic.claude-sonnet-5-5", + messages=[ + *( + { + "role": "system", + "content": [{"type": "text", "text": f"Rule {i}", "cache_control": {"type": "ephemeral"}}], + } + for i in range(system_marks) + ), + { + "role": "user", + "content": [{"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit"}}], + }, + ], + optional_params=dict(guardrail), + litellm_params={}, + ) + + assert request(0)["messages"][0]["content"] == [ + {"guardContent": {"text": {"text": "Hello"}}}, + {"cachePoint": {"type": "default"}}, + ] + with pytest.raises(litellm.UnsupportedParamsError): + request(4) + + client_marked: Final = litellm.AmazonConverseConfig()._transform_request( + model="anthropic.claude-sonnet-5-5", + messages=[ + *( + { + "role": "system", + "content": [{"type": "text", "text": f"Rule {i}", "cache_control": {"type": "ephemeral"}}], + } + for i in range(4) + ), + {"role": "user", "content": [{"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}}]}, + ], + optional_params=dict(guardrail), + litellm_params={}, + ) + assert sum("cachePoint" in block for block in client_marked["system"]) == 4 + assert client_marked["messages"][0]["content"] == [{"guardContent": {"text": {"text": "Hello"}}}] + + +def test_converse_guarded_user_text_keeps_no_cache_point_on_a_model_not_known_to_take_one(local_model_cost_map): + request: Final = litellm.AmazonConverseConfig()._transform_request( + model="arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/unmapped1", + messages=[{"role": "user", "content": [{"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}}]}], + optional_params={"guardrailConfig": {"guardrailIdentifier": "gr-1", "guardrailVersion": "DRAFT"}}, + litellm_params={}, + ) + + assert request["messages"][0]["content"] == [{"guardContent": {"text": {"text": "Hello"}}}] + + +def _converse_request_with_30m_breakpoint(model: str, litellm_params: dict) -> dict: + return litellm.AmazonConverseConfig()._transform_request( + model=model, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}} + ], + } + ], + optional_params={}, + litellm_params=litellm_params, + ) + + +def test_converse_rejects_30m_breakpoint_ttl_only_where_a_cache_point_is_written(): + claude: Final = "anthropic.claude-sonnet-5-5" + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + _converse_request_with_30m_breakpoint(claude, {}) + + assert _converse_request_with_30m_breakpoint(claude, {"drop_params": True})["messages"] == [ + {"role": "user", "content": [{"text": "Hello"}, {"cachePoint": {"type": "default"}}]} + ] + assert _converse_request_with_30m_breakpoint("meta.llama3-1-8b-instruct-v1:0", {})["messages"] == [ + {"role": "user", "content": [{"text": "Hello"}]} + ] + + @pytest.mark.parametrize( "model, expected_supports_tool_call", [ @@ -11097,6 +11667,70 @@ def test_a_deployment_with_a_batch_control_plane_endpoint_keeps_it_out_of_chat_r assert "aws_bedrock_control_plane_endpoint" not in route.calls[0].request.content.decode() +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def _converse_request_for_tools( + model: str, non_default_params: Mapping[str, object], drop_params: bool +) -> Mapping[str, object]: + config: Final = AmazonConverseConfig() + return config._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params=config.map_openai_params( + non_default_params=dict(non_default_params), optional_params={}, model=model, drop_params=drop_params + ), + messages=[{"role": "user", "content": "hi"}], + ) + + +@pytest.mark.parametrize( + "model", ["us.openai.gpt-oss-120b-1:0", "us.anthropic.claude-haiku-5-5", "us.amazon.nova-2-lite-v1:0"] +) +def test_a_hosted_mcp_tool_answers_400_on_converse_without_drop_params(monkeypatch, model): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + _converse_request_for_tools(model, {"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, drop_params=False) + assert raised.value.status_code == 400 + assert model in str(raised.value) + assert "`deepwiki`" in str(raised.value) + assert 'server_url: "litellm_proxy/mcp/"' in str(raised.value) + assert "drop_params: true" in str(raised.value) + + +def test_converse_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + model: Final = "us.openai.gpt-oss-120b-1:0" + with caplog.at_level("WARNING"): + request: Final = _converse_request_for_tools( + model, {"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, drop_params=True + ) + assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and f"model={model} on bedrock (" in record.getMessage() + for record in caplog.records + ) + + +def test_converse_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + request: Final = _converse_request_for_tools( + "us.openai.gpt-oss-120b-1:0", {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "auto"}, drop_params=True + ) + assert "toolConfig" not in request + assert "toolChoice" not in json.dumps(request) + assert "tool_choice" not in json.dumps(request) + + def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: return tuple( r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index 8ac35caba9d..d5b57e42146 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -1,6 +1,8 @@ import base64 import json +import pytest + from litellm.llms.bedrock.count_tokens.transformation import ( DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS, BedrockCountTokensConfig, @@ -53,6 +55,22 @@ def test_detect_input_type_converse_blocks_route_to_converse(): assert config._detect_input_type(request) == "converse" +@pytest.mark.parametrize("content", [[{"text": "hi"}], "hi"]) +@pytest.mark.parametrize( + "context_management", [None, {"edits": [{"type": "clear_tool_uses_20250919"}]}] +) +def test_a_pass_through_body_with_context_management_keeps_its_converse_input(content, context_management): + request = { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": content}], + "context_management": context_management, + } + + transformed = BedrockCountTokensConfig().transform_anthropic_to_bedrock_count_tokens(request) + + assert transformed["input"]["converse"]["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + + def test_transform_to_invoke_model_format_base64_encodes_body(): """The CountTokens API expects invokeModel.body as a base64-encoded blob; Anthropic Messages bodies additionally need anthropic_version/max_tokens diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index cec7950b859..30063936b2b 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -9,6 +9,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.utils import CountContextManagement RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com" MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws" @@ -144,6 +145,104 @@ async def test_mantle_failure_is_reported_with_its_status() -> None: assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) +@pytest.mark.asyncio +async def test_runtime_rejection_of_an_edit_stands_when_mantle_does_not_serve_the_model() -> None: + rejection: Final = "'claude-haiku-4-5-20251001' does not support the 'compact_20260112' context management strategy" + bedrock: Final = _Bedrock( + runtime=(400, {"message": rejection}), + mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), + ) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-haiku-4-5", + context_management=CountContextManagement( + spec={"edits": [{"type": "compact_20260112"}]}, drop_invalid_edits=False + ), + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert result.tokenizer_type == "bedrock_api" + assert rejection in (result.error_message or "") + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + + +@pytest.mark.asyncio +async def test_runtime_400_without_an_edit_still_reports_the_mantle_failure() -> None: + bedrock: Final = _Bedrock( + runtime=(400, {"message": "The provided model identifier is invalid."}), + mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), + ) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-haiku-4-5", + ) + + assert result is not None + assert result.status_code == 404 + assert result.tokenizer_type == "bedrock_mantle_api" + assert "does not exist" in (result.error_message or "") + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + + +@pytest.mark.asyncio +async def test_plain_text_history_with_an_edit_is_counted_through_invoke_carrying_the_edit() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 57}), mantle=(200, {"input_tokens": 1})) + spec: Final[dict[str, JsonValue]] = {"edits": [{"type": "clear_tool_uses_20250919"}]} + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + context_management=CountContextManagement(spec=spec, drop_invalid_edits=False), + ) + + assert result is not None + assert result.total_tokens == 57 + assert bedrock.posted_hosts() == (RUNTIME_HOST,) + invoke_body: Final = _invoke_body(bedrock.requests[0]) + assert invoke_body["context_management"] == spec + assert invoke_body["anthropic_beta"] == ["context-management-2025-06-27"] + + +@pytest.mark.asyncio +async def test_an_empty_string_turn_with_an_edit_is_sent_as_an_empty_string_not_an_empty_text_block() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 16}), mantle=(200, {"input_tokens": 1})) + spec: Final[dict[str, JsonValue]] = {"edits": [{"type": "clear_tool_uses_20250919"}]} + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": ""}, + {"role": "user", "content": "third question here"}, + ], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + context_management=CountContextManagement(spec=spec, drop_invalid_edits=False), + ) + + assert result is not None + assert result.total_tokens == 16 + assert _invoke_body(bedrock.requests[0])["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "first"}]}, + {"role": "assistant", "content": ""}, + {"role": "user", "content": [{"type": "text", "text": "third question here"}]}, + ] + + @pytest.mark.asyncio async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None: bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py index 0bd480dcdbc..07fabbf2393 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py @@ -4175,6 +4175,43 @@ async def test_bedrock_messages_400_is_not_retried_unless_it_rejects_opt_in_nati assert len(bedrock.bodies) == 1 +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("model", "model_info"), + [ + ("us.anthropic.claude-opus-5-5", None), + ("global.anthropic.claude-fable-5-1", None), + ("us.anthropic.claude-opus-4-8", None), + ("us.anthropic.claude-opus-5-5", {"supports_bedrock_native_structured_output": True}), + ], + ids=["opus_5_5_inlined", "fable_5_1_inlined", "mapped_opt_in", "deployment_opt_in"], +) +def test_bedrock_messages_offers_a_native_output_retry_only_when_it_sends_native_output( + model: str, model_info: Mapping[str, object] | None +) -> None: + from litellm.types.router import GenericLiteLLMParams + + config: Final = AmazonAnthropicClaudeMessagesConfig() + litellm_params: Final = GenericLiteLLMParams(model_info=model_info) + optional_params: Final = { + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": _CONSTRAINED_SCHEMA}}, + } + + body: Final = REASONING_BODY_ADAPTER.validate_python( + config.transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params=optional_params, + litellm_params=litellm_params, + headers={}, + ) + ) + retry_params: Final = config.native_output_fallback_litellm_params(model, optional_params, litellm_params) + + assert (retry_params is not None) is ("format" in REASONING_BODY_ADAPTER.validate_python(body.get("output_config", {}))) + + # Bedrock's 400 bodies for requests past Claude's strict tool limits, seen live on us.anthropic.claude-sonnet-4-6 on 2026-10-29 _STRICT_TOOLS_REJECTIONS: Final = ( { diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 2ac31b6887d..cf2f25faaa9 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -1,12 +1,16 @@ import asyncio +import copy import json import sys import types +from collections.abc import Mapping +from datetime import datetime, timezone from types import SimpleNamespace -from typing import Final +from typing import Final, Literal from unittest.mock import MagicMock import pytest +from pydantic import TypeAdapter import litellm from litellm.constants import ( @@ -14,10 +18,15 @@ from litellm.constants import ( REALTIME_SESSION_SUCCESS_LOGGED_KEY, WEBSOCKET_CLOSE_REASON_MAX_BYTES, ) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.realtime_errors import websocket_close_reason from litellm.llms.bedrock.common_utils import BedrockError -from litellm.llms.bedrock.realtime.handler import BedrockRealtime +from litellm.llms.bedrock.realtime.handler import BedrockRealtime, SessionEndSignals +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig from litellm.realtime_api import main as realtime_main +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import StandardLoggingPayload, StandardLoggingPayloadErrorInformation from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault @@ -36,6 +45,66 @@ def _isolate_host_aws_config(monkeypatch, tmp_path): monkeypatch.delenv(env_var, raising=False) +def _closing_input() -> SessionEndSignals: + signals: Final = SessionEndSignals() + signals.input_closing.set() + return signals + + +_SPEND_ROW_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_ERROR_INFORMATION_ADAPTER: Final = TypeAdapter(StandardLoggingPayloadErrorInformation | None) +_BEDROCK_INPUT_CHUNK_ADAPTER: Final = TypeAdapter(Mapping[str, Mapping[str, object]]) +_STANDARD_LOGGING_PAYLOAD_ADAPTER: Final = TypeAdapter(StandardLoggingPayload) + + +def _spend_row( + logging_obj: "FakeLogging | LiteLLMLogging", +) -> tuple[Literal["success", "failure"], StandardLoggingPayloadErrorInformation | None]: + return _spend_row_of(logging_obj.model_call_details) + + +def _spend_row_of( + model_call_details: dict[str, object], +) -> tuple[Literal["success", "failure"], StandardLoggingPayloadErrorInformation | None]: + now: Final = datetime.now(timezone.utc) + row: Final = get_logging_payload(kwargs=model_call_details, response_obj=None, start_time=now, end_time=now) + metadata: Final = _SPEND_ROW_METADATA_ADAPTER.validate_json(row["metadata"]) + return row["status"], _ERROR_INFORMATION_ADAPTER.validate_python(metadata.get("error_information")) + + +class LoggedEventRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: Final[list[tuple[Literal["success", "failure"], dict[str, object]]]] = [] + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.events.append(("success", dict(kwargs))) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.events.append(("failure", dict(kwargs))) + + def payload(self, event: Literal["success", "failure"]) -> StandardLoggingPayload: + (logged,) = (kwargs for logged_event, kwargs in self.events if logged_event == event) + return _STANDARD_LOGGING_PAYLOAD_ADAPTER.validate_python(logged["standard_logging_object"]) + + +class HeldSuccessLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.entered: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.entered.set() + await self.release.wait() + + class FakePayloadPart: def __init__(self, bytes_): self.bytes_ = bytes_ @@ -85,6 +154,19 @@ class HangingCloseInputStream(FakeInputStream): await asyncio.Future() +class BodyWriterInputStream(FakeInputStream): + """Models the SDK input stream, whose request-body writer task only finishes once the stream is closed""" + + def __init__(self): + super().__init__() + self._close_requested = asyncio.Event() + self.body_writer = asyncio.get_running_loop().create_task(self._close_requested.wait()) + + async def close(self): + await super().close() + self._close_requested.set() + + class FakeBedrockStream: def __init__(self, input_stream=None): self.input_stream = input_stream if input_stream is not None else FakeInputStream() @@ -124,6 +206,11 @@ class FakeLogging: self.litellm_trace_id = trace_id self.model_call_details = {} + def attempt_scoped_copy(self): + attempt = copy.copy(self) + attempt.model_call_details = dict(self.model_call_details) + return attempt + class DisconnectingClientWS: def __init__(self, messages): @@ -159,12 +246,15 @@ class EndedBedrockStream: class RealtimeClientWS: - def __init__(self): + def __init__(self, messages=()): self.closed = False self.sent_to_client = [] self.scope = {} + self._messages = list(messages) async def receive_text(self): + if self._messages: + return self._messages.pop(0) raise RuntimeError("client disconnected") async def send_text(self, message): @@ -178,8 +268,7 @@ class ConnectedClientWS(RealtimeClientWS): """Client that sends its scripted messages and then stays connected until the server closes it""" def __init__(self, messages): - super().__init__() - self._messages = list(messages) + super().__init__(messages) self._closed_event = asyncio.Event() async def receive_text(self): @@ -251,6 +340,77 @@ class ScriptedBedrockStream: return (None, self._receiver) +def _sent_event_names(input_stream: FakeInputStream) -> list[str]: + return [ + next(iter(_BEDROCK_INPUT_CHUNK_ADAPTER.validate_json(chunk.value.bytes_)["event"])) + for chunk in input_stream.sent + ] + + +class SessionEndingInputStream(FakeInputStream): + def __init__(self): + super().__init__() + self.session_ended = asyncio.Event() + + async def send(self, event): + await super().send(event) + if "sessionEnd" in json.loads(event.value.bytes_.decode("utf-8")).get("event", {}): + self.session_ended.set() + + async def close(self): + await super().close() + self.session_ended.set() + + +class SlowClosingInputStream(SessionEndingInputStream): + async def close(self): + await asyncio.Event().wait() + + +class OutputEndsAfterSessionEndReceiver(ScriptedBedrockReceiver): + def __init__(self, payloads, session_ended): + super().__init__(payloads) + self._session_ended = session_ended + + async def receive(self): + if not self._payloads: + await self._session_ended.wait() + return await super().receive() + + +class SessionBoundBedrockStream: + def __init__(self, payloads=(), input_stream_type=SessionEndingInputStream): + self.input_stream = input_stream_type() + self._receiver = OutputEndsAfterSessionEndReceiver(payloads, self.input_stream.session_ended) + + async def await_output(self): + return (None, self._receiver) + + +class ConnectionResetOnAudioInputStream(FakeInputStream): + def __init__(self, connection_reset): + super().__init__() + self._connection_reset = connection_reset + + async def send(self, event): + await asyncio.sleep(0) + if self._connection_reset.is_set() or "audioInput" in json.loads(event.value.bytes_.decode("utf-8"))["event"]: + self._connection_reset.set() + raise RuntimeError("bedrock connection reset") + self.sent.append(event) + + +class OutputEndsAfterConnectionResetReceiver(ScriptedBedrockReceiver): + def __init__(self, payloads, connection_reset): + super().__init__(payloads) + self._connection_reset = connection_reset + + async def receive(self): + if not self._payloads: + await self._connection_reset.wait() + return await super().receive() + + class FakeAWSCredentialsIdentity: def __init__(self, access_key_id, secret_access_key, session_token=None): self.access_key_id = access_key_id @@ -352,7 +512,7 @@ def stub_aws_sdk_client(monkeypatch): raise stream captured["open_stream"] = stream return stream - stream = ScriptedBedrockStream(captured.get("scripted_payloads", [])) + stream = SessionBoundBedrockStream(captured.get("scripted_payloads", [])) captured["open_stream"] = stream return stream @@ -407,7 +567,9 @@ class TestBedrockRealtimeHandler: [json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})] ) - await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) + await handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals() + ) sent_events = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent] event_names = [next(iter(event["event"])) for event in sent_events] @@ -422,7 +584,7 @@ class TestBedrockRealtimeHandler: stream = FakeBedrockStream() await handler._forward_client_to_bedrock( - DisconnectingClientWS([]), stream, config, "amazon.nova-sonic-v1:0", {} + DisconnectingClientWS([]), stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals() ) assert stream.input_stream.sent == [] @@ -437,9 +599,13 @@ class TestBedrockRealtimeHandler: [json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})] ) - with pytest.raises(RuntimeError, match="bedrock send failed"): - await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) + with pytest.raises(BedrockError) as failure: + await handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals() + ) + assert failure.value.status_code == 503 + assert str(failure.value.__cause__) == "bedrock send failed" assert stream.input_stream.closed @pytest.mark.asyncio @@ -452,7 +618,10 @@ class TestBedrockRealtimeHandler: ) await asyncio.wait_for( - handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}), 2 + handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals() + ), + 2, ) assert not stream.input_stream.closed @@ -466,7 +635,9 @@ class TestBedrockRealtimeHandler: [json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})] ) - await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) + await handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals() + ) sent_events = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent] event_names = [next(iter(event["event"])) for event in sent_events] @@ -495,6 +666,7 @@ class TestBedrockRealtimeHandler: "amazon.nova-sonic-v1:0", FakeLogging(), {}, + signals=_closing_input(), ) ] @@ -526,6 +698,7 @@ class TestBedrockRealtimeHandler: "amazon.nova-sonic-v1:0", FakeLogging(), {}, + signals=_closing_input(), ) ] @@ -597,6 +770,7 @@ class TestBedrockRealtimeHandler: "amazon.nova-sonic-v1:0", MagicMock(), {}, + signals=_closing_input(), ): pass @@ -656,7 +830,9 @@ class TestBedrockRealtimeSessionLifecycle: [json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}})] ) - await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging()) + await handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging(), signals=SessionEndSignals() + ) acked = [json.loads(message) for message in client_ws.sent_to_client] updated = [event for event in acked if event["type"] == "session.updated"] @@ -670,7 +846,9 @@ class TestBedrockRealtimeSessionLifecycle: stream = FakeBedrockStream() client_ws = DisconnectingClientWS([json.dumps({"type": "session.update", "session": {"instructions": "hi"}})]) - await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) + await handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals() + ) assert client_ws.sent_to_client == [] @@ -693,7 +871,9 @@ class TestBedrockRealtimeSessionLifecycle: stream = FakeBedrockStream() client_ws = DisconnectingClientWS([update([tool]), update([other_tool]), update([tool])]) - await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging()) + await handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging(), signals=SessionEndSignals() + ) acked: Final = [json.loads(message) for message in client_ws.sent_to_client] assert [event["type"] for event in acked] == ["session.updated", "error", "session.updated"] @@ -720,10 +900,10 @@ class TestBedrockRealtimeSessionLifecycle: async def test_malformed_setting_is_rejected_or_dropped_and_the_session_still_answers( self, stub_aws_sdk_client, session, drop_params, expected_errors ): - stream: Final = ScriptedBedrockStream([]) + stream: Final = SessionBoundBedrockStream() stub_aws_sdk_client["streams"] = [stream] user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} - websocket: Final = ConnectedClientWS( + websocket: Final = RealtimeClientWS( [ json.dumps({"type": "session.update", "session": session}), json.dumps({"type": "conversation.item.create", "item": user_text}), @@ -761,8 +941,8 @@ class TestBedrockRealtimeProviderFailurePropagation: @pytest.mark.asyncio async def test_readiness_failure_escapes_and_fallback_replays_session_update(self, stub_aws_sdk_client): handler = BedrockRealtime() - websocket = ConnectedClientWS([self.SESSION_UPDATE]) - healthy_stream = ScriptedBedrockStream([]) + websocket = RealtimeClientWS([self.SESSION_UPDATE]) + healthy_stream = SessionBoundBedrockStream() eager_failure = ServiceUnavailableException("fault injected before the stream was returned") stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), eager_failure, healthy_stream] @@ -848,9 +1028,6 @@ class TestBedrockRealtimeProviderFailurePropagation: await spend_dispatch["coro"] assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client] - flushed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent] - assert [next(iter(event["event"])) for event in flushed][-2:] == ["promptEnd", "sessionEnd"] - assert stream.input_stream.closed with pytest.raises(BedrockError) as replay: await handler.async_realtime( @@ -877,8 +1054,8 @@ class TestBedrockRealtimeProviderFailurePropagation: @pytest.mark.asyncio async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client): handler = BedrockRealtime() - websocket = ConnectedClientWS([self.USER_TEXT]) - healthy_stream = ScriptedBedrockStream([]) + websocket = RealtimeClientWS([self.USER_TEXT]) + healthy_stream = SessionBoundBedrockStream() stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream] with pytest.raises(BedrockError) as failure: @@ -907,7 +1084,7 @@ class TestBedrockRealtimeProviderFailurePropagation: [self.SESSION_UPDATE, json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})] ) - with pytest.raises(RuntimeError, match="bedrock input stream reset"): + with pytest.raises(BedrockError) as failure: await BedrockRealtime().async_realtime( model="amazon.nova-sonic-v1:0", websocket=websocket, @@ -915,6 +1092,8 @@ class TestBedrockRealtimeProviderFailurePropagation: **self.AWS_PARAMS, ) + assert failure.value.status_code == 503 + assert str(failure.value.__cause__) == "bedrock input stream reset", "the proxy log keeps the reset's cause" assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client] await spend_dispatch["coro"] assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] @@ -923,10 +1102,10 @@ class TestBedrockRealtimeProviderFailurePropagation: async def test_success_dispatch_stamps_the_ownership_marker_only_when_spend_was_logged( self, stub_aws_sdk_client, spend_dispatch ): - stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)] + stub_aws_sdk_client["streams"] = [SessionBoundBedrockStream(self.TEXT_TURN)] await BedrockRealtime().async_realtime( model="amazon.nova-sonic-v1:0", - websocket=ConnectedClientWS([self.SESSION_UPDATE]), + websocket=RealtimeClientWS([self.SESSION_UPDATE]), logging_obj=spend_dispatch["logging_obj"], **self.AWS_PARAMS, ) @@ -935,15 +1114,315 @@ class TestBedrockRealtimeProviderFailurePropagation: assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True idle_logging = FakeLogging() - stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([])] + stub_aws_sdk_client["streams"] = [SessionBoundBedrockStream()] await BedrockRealtime().async_realtime( model="amazon.nova-sonic-v1:0", - websocket=ConnectedClientWS([self.SESSION_UPDATE]), + websocket=RealtimeClientWS([self.SESSION_UPDATE]), logging_obj=idle_logging, **self.AWS_PARAMS, ) assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in idle_logging.model_call_details + @pytest.mark.asyncio + async def test_bedrock_stream_ending_mid_session_fails_the_session_and_keeps_partial_spend( + self, stub_aws_sdk_client, spend_dispatch + ): + handler: Final = BedrockRealtime() + websocket: Final = ConnectedClientWS([self.SESSION_UPDATE]) + stream: Final = ScriptedBedrockStream(self.TEXT_TURN) + stub_aws_sdk_client["streams"] = [stream] + logging_obj: Final = spend_dispatch["logging_obj"] + request_metadata: Final = {"user_api_key": "hashed-key"} + logging_obj.model_call_details["litellm_params"] = {"litellm_metadata": request_metadata} + + with pytest.raises(BedrockError) as failure: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=logging_obj, + **self.AWS_PARAMS, + ) + + assert failure.value.status_code == 503 + assert not websocket.closed, "the proxy route owns the client-facing error event and 1011 close" + assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client] + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + status, error_information = _spend_row(logging_obj) + assert status == "failure", "the Logs page must show the dropped session as failed while it bills its turn" + assert error_information is not None + assert error_information.get("error_code") == "503" + assert error_information.get("error_message") == failure.value.message + assert request_metadata == {"user_api_key": "hashed-key"}, "the request's own metadata stays as it was sent" + + with pytest.raises(BedrockError) as replay: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=logging_obj, + **self.AWS_PARAMS, + ) + + assert replay.value.status_code == 400, "a router retry must not restart the dropped session" + assert str(replay.value).startswith(failure.value.message), "the app's error event leads with the drop" + assert websocket_close_reason(str(replay.value), fallback="Internal server error").startswith( + failure.value.message + ), "the length-capped 1011 close reason must still carry the drop" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("receiver_type", "status_code"), + [(ScriptedBedrockReceiver, 503), (BreakingBedrockReceiver, 424)], + ids=["output-ends", "output-raises"], + ) + async def test_a_failed_bedrock_output_reaches_the_app_without_waiting_on_the_input_close( + self, stub_aws_sdk_client, spend_dispatch, receiver_type, status_code: int + ): + stream: Final = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=receiver_type) + stream.input_stream = HangingCloseInputStream() + stub_aws_sdk_client["streams"] = [stream] + + with pytest.raises(BedrockError) as failure: + await asyncio.wait_for( + BedrockRealtime(close_timeout=30).async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ), + timeout=5, + ) + + assert failure.value.status_code == status_code + assert stub_aws_sdk_client["client_closed"], "closing the SDK client still tears the dead stream down" + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("receiver_type", "status_code"), + [(ScriptedBedrockReceiver, 503), (BreakingBedrockReceiver, 424)], + ids=["output-ends", "output-raises"], + ) + async def test_a_failed_bedrock_output_leaves_no_pending_input_writer( + self, stub_aws_sdk_client, spend_dispatch, receiver_type, status_code: int + ): + stream: Final = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=receiver_type) + stream.input_stream = BodyWriterInputStream() + stub_aws_sdk_client["streams"] = [stream] + + with pytest.raises(BedrockError) as failure: + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + + assert failure.value.status_code == status_code + await asyncio.wait_for(stream.input_stream.body_writer, timeout=1) + await spend_dispatch["coro"] + + @pytest.mark.asyncio + async def test_a_dropped_session_logs_one_success_usage_record_and_a_failed_spend_row( + self, stub_aws_sdk_client, spend_dispatch + ): + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)] + recorder: Final = LoggedEventRecorder() + logging_obj: Final = LiteLLMLogging( + model="amazon.nova-sonic-v1:0", + messages=[], + stream=False, + call_type="_arealtime", + start_time=datetime.now(), + litellm_call_id="nova-sonic-drop", + function_id="nova-sonic-drop", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_from_kwargs( + kwargs={"litellm_metadata": {"user_api_key": "hashed-key"}}, + model="amazon.nova-sonic-v1:0", + optional_params={}, + litellm_params={"_arealtime": True}, + custom_llm_provider="bedrock", + ) + + with pytest.raises(BedrockError) as failure: + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=logging_obj, + **self.AWS_PARAMS, + ) + await spend_dispatch["coro"] + + assert failure.value.status_code == 503 + ((logged_event, usage_kwargs),) = recorder.events + assert logged_event == "success" + usage_record: Final = recorder.payload("success") + assert usage_record["status"] == "success", "the router's failure event is each logger's one error record" + assert [event["type"] for event in usage_record["response"]["results"]] == ["response.done"] + assert logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + status, error_information = _spend_row_of(usage_kwargs) + assert status == "failure" + assert error_information is not None + assert error_information.get("error_code") == "503" + + @pytest.mark.asyncio + async def test_a_dropped_session_bills_its_usage_to_loggers_while_the_failure_event_runs( + self, stub_aws_sdk_client, spend_dispatch + ): + usage_event: Final = json.dumps( + {"event": {"usageEvent": {"totalInputTokens": 1000, "totalOutputTokens": 50, "totalTokens": 1050}}} + ) + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream((*self.TEXT_TURN, usage_event))] + held: Final = HeldSuccessLogger() + recorder: Final = LoggedEventRecorder() + logging_obj: Final = LiteLLMLogging( + model="amazon.nova-2-5-sonic", + messages=[], + stream=False, + call_type="_arealtime", + start_time=datetime.now(), + litellm_call_id="nova-sonic-drop-order", + function_id="nova-sonic-drop-order", + dynamic_async_success_callbacks=[held, recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_from_kwargs( + kwargs={"litellm_metadata": {"user_api_key": "hashed-key"}}, + model="amazon.nova-2-5-sonic", + optional_params={}, + litellm_params={"_arealtime": True}, + custom_llm_provider="bedrock", + ) + + with pytest.raises(BedrockError) as failure: + await BedrockRealtime().async_realtime( + model="amazon.nova-2-5-sonic", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=logging_obj, + **self.AWS_PARAMS, + ) + success_dispatch: Final = asyncio.create_task(spend_dispatch["coro"]) + await asyncio.wait_for(held.entered.wait(), timeout=5) + logging_obj.failure_handler(failure.value, "", logging_obj.start_time, datetime.now(), for_async_request=True) + await logging_obj.async_failure_handler(failure.value, "", logging_obj.start_time, datetime.now()) + held.release.set() + await success_dispatch + + usage_record: Final = recorder.payload("success") + assert recorder.payload("failure")["status"] == "failure" + assert usage_record["status"] == "success" + assert usage_record["total_tokens"] == 1050 + assert usage_record["response_cost"] > 0 + ((_, usage_kwargs),) = (logged for logged in recorder.events if logged[0] == "success") + status, _ = _spend_row_of(usage_kwargs) + assert status == "failure" + + @pytest.mark.asyncio + async def test_a_dropped_session_without_request_metadata_still_fails_and_bills_its_turn( + self, stub_aws_sdk_client, spend_dispatch + ): + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)] + logging_obj: Final = spend_dispatch["logging_obj"] + logging_obj.model_call_details["litellm_params"] = {"litellm_metadata": None, "metadata": {}} + + with pytest.raises(BedrockError) as failure: + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=logging_obj, + **self.AWS_PARAMS, + ) + await spend_dispatch["coro"] + + assert failure.value.status_code == 503 + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("receiver_type", "status_code"), + [(ScriptedBedrockReceiver, 503), (BreakingBedrockReceiver, 424)], + ids=["output-ends", "output-raises"], + ) + async def test_a_session_bedrock_fails_before_any_logged_event_dispatches_no_success( + self, stub_aws_sdk_client, spend_dispatch, receiver_type, status_code: int + ): + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream((), receiver_type=receiver_type)] + logging_obj: Final = spend_dispatch["logging_obj"] + logging_obj.model_call_details["litellm_params"] = {"litellm_metadata": {"user_api_key": "hashed-key"}} + + with pytest.raises(BedrockError) as failure: + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=logging_obj, + **self.AWS_PARAMS, + ) + + assert failure.value.status_code == status_code + assert "coro" not in spend_dispatch, "a failed session with no usage gets no success event, as on base" + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logging_obj.model_call_details + + @pytest.mark.asyncio + async def test_bedrock_output_ending_right_after_session_end_is_a_clean_finish( + self, stub_aws_sdk_client, spend_dispatch + ): + stream: Final = SessionBoundBedrockStream(self.TEXT_TURN, input_stream_type=SlowClosingInputStream) + stub_aws_sdk_client["streams"] = [stream] + websocket: Final = RealtimeClientWS([self.SESSION_UPDATE]) + spend_dispatch["logging_obj"].model_call_details["litellm_params"] = { + "litellm_metadata": {"user_api_key": "hashed-key"} + } + + await asyncio.wait_for( + BedrockRealtime(close_timeout=30).async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ), + timeout=5, + ) + + assert _sent_event_names(stream.input_stream)[-1] == "sessionEnd" + assert not stream.input_stream.closed, "the input close must still be pending when the output ends" + assert websocket.closed + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + assert _spend_row(spend_dispatch["logging_obj"]) == ("success", None) + + @pytest.mark.asyncio + async def test_bedrock_output_ending_after_the_input_side_failed_still_fails_the_session( + self, stub_aws_sdk_client, spend_dispatch + ): + connection_reset: Final = asyncio.Event() + stream: Final = ScriptedBedrockStream( + self.TEXT_TURN, + receiver_type=lambda payloads: OutputEndsAfterConnectionResetReceiver(payloads, connection_reset), + ) + stream.input_stream = ConnectionResetOnAudioInputStream(connection_reset) + stub_aws_sdk_client["streams"] = [stream] + websocket: Final = ConnectedClientWS([self.SESSION_UPDATE, self.USER_AUDIO]) + + with pytest.raises(BedrockError) as failure: + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + + assert (failure.value.status_code, failure.value.message) == ( + 503, + "Bedrock realtime stream ended before the session closed", + ), "the app sees one cause whichever side of the stream reports the reset first" + assert not websocket.closed, "the proxy route owns the client-facing error event and 1011 close" + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + @pytest.mark.asyncio async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client): stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver) @@ -982,7 +1461,13 @@ class TestBedrockRealtimeProviderFailurePropagation: with pytest.raises(ServiceUnavailableException): await handler._forward_client_to_bedrock( - client_ws, stream, BedrockRealtimeConfig(), "amazon.nova-sonic-v1:0", {}, FakeLogging() + client_ws, + stream, + BedrockRealtimeConfig(), + "amazon.nova-sonic-v1:0", + {}, + FakeLogging(), + signals=SessionEndSignals(), ) assert client_ws.sent_to_client == [] @@ -1128,7 +1613,6 @@ class TestBedrockRealtimeSdkLifecycle: ) assert stub_aws_sdk_client["client_closed"] - assert stub_aws_sdk_client["input_closed_before_client_close"] @pytest.mark.asyncio async def test_client_without_close_completes_session(self, monkeypatch): diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 38c69e0ce27..80715105b7f 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -1,5 +1,7 @@ import asyncio, importlib, litellm, pytest +import respx +from httpx import Response from typing import Final @@ -581,6 +583,188 @@ def test_deployment_model_info_breakpoint_flag_covers_an_unmapped_arn(local_mode assert bedrock_model_accepts_cache_points(unflagged_arn) is True + +_CONVERSE_REPLY: Final = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}, +} +_CACHE_MARK: Final = {"type": "ephemeral"} +_MARKED_WEATHER_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}, + "cache_control": _CACHE_MARK, +} +_COMPUTER_TOOL: Final = { + "type": "computer_20250124", + "name": "computer", + "display_width_px": 1024, + "display_height_px": 768, +} +_MARKED_TOOL_CONVERSATION: Final = ( + {"role": "system", "content": "stable"}, + {"role": "user", "content": [{"type": "text", "text": "Weather in Paris and Rome?", "cache_control": _CACHE_MARK}]}, + { + "role": "assistant", + "content": [{"type": "text", "text": "Checking both.", "cache_control": _CACHE_MARK}], + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}{"city": "Rome"}'}, + "cache_control": _CACHE_MARK, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": [{"type": "text", "text": "Sunny", "cache_control": _CACHE_MARK}]}, + {"role": "tool", "tool_call_id": "call_1_1", "content": "Cloudy"}, + {"role": "assistant", "content": "Paris is sunny, Rome is cloudy.", "cache_control": _CACHE_MARK}, + {"role": "user", "content": "Thanks"}, +) + + +@pytest.mark.parametrize( + ("model", "flagged", "declared", "tools"), + [ + pytest.param("moonshotai.kimi-k3", "opted-in", True, (_MARKED_WEATHER_TOOL,), id="kimi-opted-in"), + pytest.param( + "us.anthropic.claude-sonnet-5-5", + "opted-out", + False, + (_COMPUTER_TOOL, _MARKED_WEATHER_TOOL), + id="claude-opted-out-next-to-computer-use", + ), + ], +) +@pytest.mark.parametrize("entrypoint", ["completion", "acompletion"]) +@pytest.mark.parametrize("flagless_registers_first", [False, True]) +async def test_a_bedrock_deployment_sends_cache_points_by_its_own_flag_not_a_siblings_on_the_same_model( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, + model: str, + flagged: str, + declared: bool, + tools: tuple[dict[str, object], ...], + entrypoint: str, + flagless_registers_first: bool, +) -> None: + from litellm import Router + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + converse: Final = respx_mock.post(f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model}/converse").mock( + return_value=Response(200, json=_CONVERSE_REPLY) + ) + flags: Final = {flagged: {"supports_prompt_cache_breakpoint": declared}} + registration_order: Final = ("flagless", flagged) if flagless_registers_first else (flagged, "flagless") + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": f"bedrock/{model}", + "aws_access_key_id": "bedrock-test-key", + "aws_secret_access_key": "bedrock-test-secret", + "aws_region_name": "us-east-1", + "cache_control_injection_points": [{"location": "message", "role": "system"}], + }, + "model_info": {"id": f"{name}-id", **flags.get(name, {})}, + } + for name in registration_order + ], + num_retries=0, + ) + messages: Final = [dict(message) for message in _MARKED_TOOL_CONVERSATION] + + async def sends_a_cache_point(name: str) -> bool: + match entrypoint: + case "completion": + router.completion(model=name, messages=messages, tools=list(tools)) + case "acompletion": + await router.acompletion(model=name, messages=messages, tools=list(tools)) + return b"cachePoint" in converse.calls.last.request.content + + assert {name: await sends_a_cache_point(name) for name in registration_order} == { + flagged: declared, + "flagless": not declared, + } + +async def test_a_routed_bedrock_deployment_sends_cache_points_when_register_model_opted_its_model_in( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None +) -> None: + from litellm import Router + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + converse: Final = respx_mock.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/moonshotai.kimi-k3/converse" + ).mock(return_value=Response(200, json=_CONVERSE_REPLY)) + litellm.register_model({"bedrock/moonshotai.kimi-k3": {"supports_prompt_cache_breakpoint": True}}) + router: Final = Router( + model_list=[ + { + "model_name": "flagless", + "litellm_params": { + "model": "bedrock/moonshotai.kimi-k3", + "aws_access_key_id": "bedrock-test-key", + "aws_secret_access_key": "bedrock-test-secret", + "aws_region_name": "us-east-1", + "cache_control_injection_points": [{"location": "message", "role": "system"}], + }, + "model_info": {"id": "flagless-id"}, + } + ], + num_retries=0, + ) + + await router.acompletion(model="flagless", messages=[{"role": "system", "content": "stable"}, {"role": "user", "content": "hi"}]) + + assert b"cachePoint" in converse.calls.last.request.content + + +@pytest.mark.parametrize( + ("catalog_caching", "deployment_caching", "expected"), + [(None, True, True), (True, False, False), (True, None, True), (None, None, False)], +) +def test_a_routed_deployments_own_prompt_caching_flag_decides_cache_points_for_a_catalogued_model( + catalog_caching: bool | None, deployment_caching: bool | None, expected: bool, local_model_cost_map: None +) -> None: + from litellm.llms.bedrock.common_utils import bedrock_model_accepts_cache_points + + model: Final = f"qa-catalogued-model-{catalog_caching}-{deployment_caching}" + catalog_flag: Final = {} if catalog_caching is None else {"supports_prompt_caching": catalog_caching} + deployment_flag: Final = {} if deployment_caching is None else {"supports_prompt_caching": deployment_caching} + litellm.register_model({model: {"litellm_provider": "bedrock", "mode": "chat", **catalog_flag}}) + + assert bedrock_model_accepts_cache_points(model, model_info={"id": "qa-deployment", **deployment_flag}) is expected + + +@pytest.mark.parametrize( + ("declared", "catalog_caching", "expected"), + [ + ("true", False, True), + (1, False, True), + (" TRUE ", False, True), + ("false", True, False), + (0, True, False), + ("maybe", True, True), + ("maybe", False, False), + ], +) +def test_a_deployments_prompt_caching_flag_written_as_a_string_or_number_decides_cache_points_as_the_bool_it_spells( + declared: object, catalog_caching: bool, expected: bool, local_model_cost_map: None +) -> None: + from litellm.llms.bedrock.common_utils import bedrock_model_accepts_cache_points + + model: Final = f"qa-catalogued-model-lax-{catalog_caching}" + litellm.register_model( + {model: {"litellm_provider": "bedrock", "mode": "chat", "supports_prompt_caching": catalog_caching}} + ) + + assert ( + bedrock_model_accepts_cache_points(model, model_info={"id": "qa-deployment", "supports_prompt_caching": declared}) + is expected + ) + def test_merge_bedrock_aws_request_params_strips_caller_identity_when_deployment_has_static_credentials(): from litellm.llms.bedrock.common_utils import merge_bedrock_aws_request_params diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 0ae6a73968d..f8f18852543 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -1922,3 +1922,68 @@ class TestBedrockMantleResponsesPricing: def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models + + +def _mantle_marked_part( + litellm_params: GenericLiteLLMParams, marker: object, model: str = "openai.gpt-6.1-sol" +) -> object: + body: Final = BedrockMantleResponsesAPIConfig().transform_responses_api_request( + model=model, + input=[ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": marker}], + } + ], + response_api_optional_request_params={}, + litellm_params=litellm_params, + headers={}, + ) + return body["input"][0]["content"][0] + + +@pytest.mark.parametrize( + "marker, wire_part", + [ + ( + {"mode": "explicit", "ttl": "30m"}, + {"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ), + ({"mode": "implicit"}, {"type": "input_text", "text": "Stable prefix"}), + ], + ids=("ttl-dropped", "implicit-dropped"), +) +def test_mantle_drop_params_cuts_a_responses_breakpoint_down_to_its_mode( + monkeypatch: pytest.MonkeyPatch, marker: dict[str, str], wire_part: dict[str, object] +) -> None: + """Bedrock Mantle Responses answered 400 for a 30m breakpoint ttl and 200 for a bare explicit mode on 2026-10-29.""" + monkeypatch.setattr(litellm, "drop_params", False) + assert _mantle_marked_part(GenericLiteLLMParams(drop_params=True), marker) == wire_part + assert _mantle_marked_part(GenericLiteLLMParams(), marker) == { + "type": "input_text", + "text": "Stable prefix", + "prompt_cache_breakpoint": marker, + } + + +@pytest.mark.parametrize( + "marker, wire_part", + [ + ({"mode": "explicit", "ttl": "1h"}, {"type": "input_text", "text": "Stable prefix"}), + ({"mode": "explicit", "ttl": "30m"}, {"type": "input_text", "text": "Stable prefix"}), + ( + {"mode": "explicit"}, + {"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ), + ], + ids=("1h-dropped", "30m-dropped", "explicit-kept"), +) +def test_mantle_drop_params_sends_an_unlisted_model_only_a_bare_explicit_breakpoint( + monkeypatch: pytest.MonkeyPatch, marker: dict[str, str], wire_part: dict[str, object] +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + assert ( + _mantle_marked_part(GenericLiteLLMParams(drop_params=True), marker, model="openai.not-in-the-cost-map") + == wire_part + ) diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index f252d9886d3..6005acffc0a 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -1096,6 +1096,25 @@ class TestBedrockMantleClaudeChatRoute: ) assert response._hidden_params["response_cost"] != pytest.approx(_row_cost("anthropic.claude-opus-5-5", 10, 5)) + def test_claude_completion_without_max_tokens_sends_the_mantle_rows_output_limit(self, monkeypatch, local_cost_map): + from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + mapped_cap: Final = litellm.model_cost["bedrock_mantle/anthropic.claude-haiku-4-5"]["max_output_tokens"] + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "mantle-key") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + handler: Final = Mock(side_effect=_anthropic_message) + + litellm.completion( + model="bedrock_mantle/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "hello"}], + aws_region_name="us-east-1", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + ) + + assert json.loads(handler.call_args.args[0].content)["max_tokens"] == mapped_cap + def test_claude_streaming_completion_uses_native_messages_endpoint(self, monkeypatch, local_cost_map): from litellm.llms.custom_httpx.http_handler import HTTPHandler diff --git a/tests/unit/llms/cohere/chat/test_cohere_transformation.py b/tests/unit/llms/cohere/chat/test_cohere_transformation.py index d32b3bcbeca..fcb3d3f1cc0 100644 --- a/tests/unit/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/unit/llms/cohere/chat/test_cohere_transformation.py @@ -576,6 +576,54 @@ def test_cohere_v2_tool_calling(respx_mock: respx.MockRouter) -> None: assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ ("get_weather_k9q2xw1n", "get_weather", '{"location":"New York, NY"}') ] + assert response.choices[0].finish_reason == "tool_calls" + + +@pytest.mark.usefixtures("_cohere_httpx_transport") +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + ("cohere_finish_reason", "expected_finish_reason"), + [("COMPLETE", "stop"), ("MAX_TOKENS", "length"), ("STOP_SEQUENCE", "stop")], +) +def test_cohere_v2_maps_finish_reason( + respx_mock: respx.MockRouter, cohere_finish_reason: str, expected_finish_reason: str +) -> None: + respx_mock.post(COHERE_V2_CHAT_URL).mock( + return_value=httpx.Response( + 200, + json=_cohere_v2_response( + {"content": [{"type": "text", "text": "Hello"}]}, + finish_reason=cohere_finish_reason, + ), + ) + ) + + response: Final = litellm.completion( + model="cohere_chat/v2/command-a-03-2025", + messages=[{"role": "user", "content": "Say hello"}], + api_key="cohere-test-key", + ) + + assert response.choices[0].finish_reason == expected_finish_reason + + +@pytest.mark.usefixtures("_cohere_httpx_transport") +@pytest.mark.respx(assert_all_called=True) +def test_cohere_v2_missing_finish_reason_defaults_to_stop(respx_mock: respx.MockRouter) -> None: + body_without_finish_reason: Final = { + key: value + for key, value in _cohere_v2_response({"content": [{"type": "text", "text": "Hello"}]}).items() + if key != "finish_reason" + } + respx_mock.post(COHERE_V2_CHAT_URL).mock(return_value=httpx.Response(200, json=body_without_finish_reason)) + + response: Final = litellm.completion( + model="cohere_chat/v2/command-a-03-2025", + messages=[{"role": "user", "content": "Say hello"}], + api_key="cohere-test-key", + ) + + assert (response.choices[0].message.content, response.choices[0].finish_reason) == ("Hello", "stop") @pytest.mark.usefixtures("_cohere_httpx_transport") diff --git a/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py new file mode 100644 index 00000000000..f5cb874c6ef --- /dev/null +++ b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py @@ -0,0 +1,35 @@ +from typing import Final + +import httpx +import pytest +import respx + +import litellm + +BEDROCK_COHERE_INVOKE_URL: Final = ( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/cohere.embed-english-v3/invoke" +) +COHERE_EMBED_RESPONSE: Final = {"embeddings": [[0.1, 0.2]], "meta": {"billed_units": {"input_tokens": 1}}} + + +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_bedrock_cohere_embedding_sends_the_caller_timeout( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post(BEDROCK_COHERE_INVOKE_URL).mock( + return_value=httpx.Response(200, json=COHERE_EMBED_RESPONSE) + ) + request: Final = { + "model": "bedrock/cohere.embed-english-v3", + "input": ["hello"], + "timeout": 45, + "aws_region_name": "us-east-1", + "aws_access_key_id": "mockaws_access_key_id", + "aws_secret_access_key": "mockaws_secret_access_key", + } + + response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request) + + assert route.calls.last.request.extensions["timeout"]["read"] == 45 + assert response.data[0]["embedding"] == [0.1, 0.2] 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 a8f6772a76e..01760572f33 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -1668,7 +1668,7 @@ _BOTH_FOUNDRY_BETAS: Final = ("context-1m-2025-08-07,context-management-2025-06- pytest.param( None, {"anthropic-beta": "context-1m-2025-08-07"}, - "claude-code-20250219", + "mcp-servers-2025-12-04", ("context-1m-2025-08-07",), id="unsupported-client-beta-is-dropped-and-deployment-beta-kept", ), @@ -6243,6 +6243,142 @@ async def test_non_chatgpt_responses_preserves_caller_headers_after_validation( client.client.close() +_BETA_THE_TABLE_DOES_NOT_LIST: Final = "not-in-the-beta-table-2099-01-01" + + +def _anthropic_messages_logging_obj(): + from datetime import datetime, timezone + + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="foundry-beta-row", + function_id="foundry-beta-row", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("custom_llm_provider", "messages_config", "foundry_api_base", "forwarded_betas"), + [ + pytest.param( + "anthropic", + "anthropic", + "https://my-resource.services.ai.azure.com/anthropic", + f"thinking-binding-controls-2026-08-01,{_BETA_THE_TABLE_DOES_NOT_LIST}", + id="anthropic-drops-what-the-azure-ai-row-rejects", + ), + pytest.param( + "anthropic", + "anthropic", + "https://my-resource.cognitiveservices.azure.com/anthropic", + f"thinking-binding-controls-2026-08-01,{_BETA_THE_TABLE_DOES_NOT_LIST}", + id="anthropic-on-a-cognitiveservices-host", + ), + pytest.param( + "azure_ai", + "azure_ai", + "https://my-resource.services.ai.azure.com/anthropic", + "thinking-binding-controls-2026-08-01", + id="azure-ai-keeps-dropping-unlisted-betas", + ), + pytest.param( + "openai", + "pass_through", + "https://my-resource.services.ai.azure.com/anthropic", + f"advisor-tool-2026-03-01,thinking-binding-controls-2026-08-01,{_BETA_THE_TABLE_DOES_NOT_LIST}", + id="messages-pass-through-forwards-as-sent", + ), + ], +) +async def test_async_anthropic_messages_handler_picks_the_beta_row_by_deployment_on_a_foundry_api_base( + monkeypatch: pytest.MonkeyPatch, + custom_llm_provider: str, + messages_config: str, + foundry_api_base: str, + forwarded_betas: str, +) -> None: + from litellm import anthropic_beta_headers_manager + from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig + from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig + + monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") + monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None) + assert _BETA_THE_TABLE_DOES_NOT_LIST not in anthropic_beta_headers_manager._load_beta_headers_config()["azure_ai"] + configs: Final = { + "anthropic": AnthropicMessagesConfig, + "azure_ai": AzureAnthropicMessagesConfig, + "pass_through": OpenAILikeAnthropicMessagesConfig, + } + upstream: Final = _AnthropicMessagesUpstream() + + await BaseLLMHTTPHandler().async_anthropic_messages_handler( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_provider_config=configs[messages_config](), + anthropic_messages_optional_request_params={"max_tokens": 8}, + custom_llm_provider=custom_llm_provider, + litellm_params=GenericLiteLLMParams(api_base=foundry_api_base), + logging_obj=_anthropic_messages_logging_obj(), + api_key="sk-test", + api_base=foundry_api_base, + client=AsyncHTTPHandler(transport=httpx.MockTransport(upstream.respond)), + kwargs={ + "extra_headers": { + "anthropic-beta": ( + f"advisor-tool-2026-03-01,thinking-binding-controls-2026-08-01,{_BETA_THE_TABLE_DOES_NOT_LIST}" + ) + } + }, + ) + + assert set(",".join(upstream.anthropic_beta).split(",")) == set(forwarded_betas.split(",")) + + +@pytest.mark.asyncio +async def test_first_foundry_request_with_betas_fetches_the_beta_table_off_the_event_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm import anthropic_beta_headers_manager, anthropic_beta_headers_url + from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + monkeypatch.delenv("LITELLM_OFFLINE_MODE", raising=False) + monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None) + bundled_table: Final = anthropic_beta_headers_manager.GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() + foundry_api_base: Final = "https://my-resource.services.ai.azure.com/anthropic" + upstream: Final = _AnthropicMessagesUpstream() + + with respx.mock(assert_all_called=True) as router: + table_route: Final = router.get(anthropic_beta_headers_url).mock( + side_effect=lambda request: httpx.Response( + 200, json=bundled_table, headers={"x-fetch-thread": str(threading.get_ident())} + ) + ) + await BaseLLMHTTPHandler().async_anthropic_messages_handler( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_provider_config=AnthropicMessagesConfig(), + anthropic_messages_optional_request_params={"max_tokens": 8}, + custom_llm_provider="anthropic", + litellm_params=GenericLiteLLMParams(api_base=foundry_api_base), + logging_obj=_anthropic_messages_logging_obj(), + api_key="sk-test", + api_base=foundry_api_base, + client=AsyncHTTPHandler(transport=httpx.MockTransport(upstream.respond)), + kwargs={"extra_headers": {"anthropic-beta": "advisor-tool-2026-03-01,compact-2026-01-12"}}, + ) + + assert table_route.calls.last.response.headers["x-fetch-thread"] != str(threading.get_ident()) + assert upstream.anthropic_beta == ("compact-2026-01-12",) + + @pytest.mark.asyncio async def test_a_stream_answered_with_one_json_body_is_not_replayed_for_a_provider_that_does_not_opt_in( httpx_transport: None, diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 7a6f9c7dc18..d2b78235c29 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -1,14 +1,14 @@ import copy import json from collections.abc import AsyncIterator, Iterator -from typing import Final +from typing import Final, cast from unittest.mock import MagicMock, patch import httpx import pytest import respx from openai.lib.streaming.chat import ChatCompletionStreamState -from openai.types.chat import ChatCompletionChunk +from openai.types.chat import ChatCompletionChunk, ChatCompletionMessage from pydantic import TypeAdapter from fastapi.testclient import TestClient from pydantic import BaseModel, ConfigDict @@ -21,6 +21,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.databricks.common_utils import DatabricksException from litellm.llms.databricks.chat.transformation import ( DatabricksChatResponseIterator, DatabricksConfig, @@ -539,13 +540,6 @@ def test_transform_request_preserves_unity_model_service_name(): [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], id="cache-control", ), - pytest.param( - "databricks-meta-llama-3-3-70b-instruct", - [{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}], - {"response_format": {"type": "json_object"}}, - [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}], - id="multiple-text-blocks", - ), pytest.param( "databricks-meta-llama-3-3-70b-instruct", [ @@ -603,6 +597,283 @@ def test_transform_request_collapses_single_plain_text_block_for_non_claude_mode assert outgoing_messages == [{**messages[0], "content": expected_content}] +_TWO_TEXT_BLOCKS: Final = [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}] +_JSON_OBJECT: Final = {"response_format": {"type": "json_object"}} +_IMAGE_BLOCK: Final = {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}} +_CACHED_BLOCK: Final = {"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}} + + +def _outgoing_messages( + model: str, messages: list[AllMessageValues], optional_params: dict[str, object], drop_params: bool = False +) -> list[dict[str, object]]: + request: Final = TypeAdapter(dict[str, object]).validate_python( + DatabricksConfig().transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={"drop_params": True} if drop_params else {}, + headers={}, + ) + ) + return TypeAdapter(list[dict[str, object]]).validate_python(request["messages"]) + + +@pytest.mark.parametrize( + "model,messages,optional_params,expected_content", + [ + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + "Reply in JSON", + id="json-object-two-blocks-joined", + ), + pytest.param( + "databricks-qwen35-122b-a10b", + [{"role": "user", "content": [{"type": "text", "text": "A"}, {"type": "text", "text": "B json"}] * 2}], + _JSON_OBJECT, + "AB jsonAB json", + id="json-object-four-blocks-joined-as-sent", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + {}, + _TWO_TEXT_BLOCKS, + id="no-response-format-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + {"response_format": {"type": "json_schema", "json_schema": {"name": "r", "schema": {"type": "object"}}}}, + _TWO_TEXT_BLOCKS, + id="json-schema-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply"}, _CACHED_BLOCK], + id="json-object-cache-control-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + id="json-object-message-cache-control-stays-block", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK]}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK], + id="json-object-text-and-image-stays-list", + ), + pytest.param( + "databricks-claude-sonnet-5", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + _TWO_TEXT_BLOCKS, + id="json-object-claude-stays-list", + ), + pytest.param( + "system.ai.deepseek-v4-flash-0731", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + "Reply in JSON", + id="json-object-ai-gateway-model-joined", + ), + pytest.param( + "openai-gpt-external", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + _TWO_TEXT_BLOCKS, + id="json-object-external-endpoint-stays-list", + ), + ], +) +def test_transform_request_joins_plain_text_blocks_only_for_hosted_non_claude_json_object( + model: str, + messages: list[AllMessageValues], + optional_params: dict[str, object], + expected_content: object, +) -> None: + assert _outgoing_messages(model, messages, optional_params) == [{"role": "user", "content": expected_content}] + + +@pytest.mark.parametrize("response_format", [None, {"type": "json_object"}], ids=["no-response-format", "json-object"]) +def test_transform_request_rejects_a_text_block_whose_text_is_a_truthy_non_string( + response_format: dict[str, str] | None, +) -> None: + with pytest.raises(DatabricksException) as error: + _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, {"type": "text", "text": 7}]}], + {} if response_format is None else {"response_format": response_format}, + ) + + assert error.value.status_code == 400 + + +def test_transform_request_accepts_an_openai_message_object_in_history() -> None: + earlier_answer: Final = ChatCompletionMessage(role="assistant", content="Earlier answer") + + outgoing: Final = _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [ + {"role": "user", "content": "Earlier question"}, + cast(AllMessageValues, earlier_answer), # cast-ok: SDK callers pass response messages back as history + {"role": "user", "content": "Next question"}, + ], + {}, + ) + + assert [message.get("content") for message in outgoing] == ["Earlier question", "Earlier answer", "Next question"] + + +def test_transform_request_json_object_keeps_earlier_turns_as_sent() -> None: + earlier_turn: Final[AllMessageValues] = { + "role": "user", + "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}], + } + reply: Final[AllMessageValues] = {"role": "assistant", "content": "Got it."} + + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [earlier_turn, reply, {"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + ) == [earlier_turn, reply, {"role": "user", "content": "Reply in JSON"}] + + +def test_transform_request_json_object_joins_a_trailing_system_run_and_keeps_its_break() -> None: + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": _TWO_TEXT_BLOCKS}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}, + ] + + +_CACHED_TEXT_AND_IMAGE: Final = [{**_CACHED_BLOCK, "text": "Reply in JSON"}, _IMAGE_BLOCK] + + +@pytest.mark.parametrize( + "message,expected", + [ + pytest.param( + {"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}, + {"role": "user", "content": "Reply in JSON"}, + id="block-cache-control-dropped-and-joined", + ), + pytest.param( + {"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Reply in JSON"}, + id="message-cache-control-dropped", + ), + pytest.param( + {"role": "user", "content": _CACHED_TEXT_AND_IMAGE}, + {"role": "user", "content": _CACHED_TEXT_AND_IMAGE}, + id="image-keeps-cache-control", + ), + ], +) +def test_transform_request_json_object_drops_last_message_cache_control_under_drop_params( + message: AllMessageValues, expected: dict[str, object] +) -> None: + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT, drop_params=True) == [ + expected + ] + + +_THINKING_BLOCK: Final = {"type": "thinking", "thinking": "Plan the reply.", "signature": "sig"} + + +@pytest.mark.parametrize( + "dropped_turn", + [ + pytest.param({"role": "assistant", "content": "", "thinking_blocks": [_THINKING_BLOCK]}, id="thinking-only"), + pytest.param({"role": "assistant", "content": None, "reasoning_content": "Plan."}, id="reasoning-only"), + pytest.param({"role": "assistant", "content": [{"type": "text", "text": " "}]}, id="blank-text-part"), + ], +) +def test_transform_request_json_object_joins_the_message_before_a_dropped_assistant_turn( + dropped_turn: AllMessageValues, +) -> None: + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}, dropped_turn], + _JSON_OBJECT, + ) == [{"role": "user", "content": "Reply in JSON"}] + + +def test_transform_request_json_object_joins_a_system_run_split_by_a_dropped_assistant_turn() -> None: + earlier_turn: Final[AllMessageValues] = { + "role": "user", + "content": [{"type": "text", "text": "Say"}, {"type": "text", "text": " hi"}], + } + messages: Final[list[AllMessageValues]] = [ + earlier_turn, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "assistant", "content": "", "thinking_blocks": [_THINKING_BLOCK]}, + {"role": "system", "content": _TWO_TEXT_BLOCKS}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + earlier_turn, + {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}, + ] + + +@pytest.mark.parametrize( + "content,expected_content", + [ + pytest.param( + [{"type": "text", "text": "Reply in JSON"}, {"type": "text", "text": "\n\n"}], + "Reply in JSON", + id="one-non-blank-part-goes-out-as-before", + ), + pytest.param( + [{"type": "text", "text": "Reply"}, {"type": "text", "text": " "}, {"type": "text", "text": "in JSON"}], + "Reply in JSON", + id="blank-separator-between-parts-kept", + ), + ], +) +def test_transform_request_json_object_joins_only_two_or_more_non_blank_parts( + content: list[object], expected_content: str +) -> None: + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", [{"role": "user", "content": content}], _JSON_OBJECT + ) == [{"role": "user", "content": expected_content}] + + +def test_transform_request_json_object_keeps_message_cache_control_parts_as_sent_without_drop_params() -> None: + message: Final[AllMessageValues] = { + "role": "user", + "content": _TWO_TEXT_BLOCKS, + "cache_control": {"type": "ephemeral"}, + } + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], {}) == _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT + ) + + +def test_transform_request_json_object_reads_the_global_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}], + _JSON_OBJECT, + ) == [{"role": "user", "content": "Reply in JSON"}] + + def test_transform_request_strips_thinking_blocks_and_reasoning_content(): """Regression for LIT-6762: replaying an assistant turn that litellm decorated with `thinking_blocks` / `reasoning_content` made Databricks 400 with @@ -2775,6 +3046,74 @@ def test_claude_json_schema_stream_split_across_json_tools_keeps_the_non_stream_ assert [choice.finish_reason for choice in choices if choice.finish_reason] == [answer.finish_reason] +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def test_databricks_claude_answers_400_for_a_hosted_mcp_tool_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + DatabricksConfig().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "databricks-claude-opus-5-5 on databricks " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_databricks_claude_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert [tool["function"]["name"] for tool in mapped["tools"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and " on databricks (" in record.getMessage() for record in caplog.records + ) + + +def test_databricks_claude_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert "tools" not in mapped + assert "tool_choice" not in mapped + + +def test_databricks_claude_still_sends_max_completion_tokens_as_max_tokens_when_dropping_a_hosted_mcp_tool( + monkeypatch, +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL], "max_completion_tokens": 64}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert mapped["max_tokens"] == 64 + assert "max_completion_tokens" not in mapped + assert [tool["function"]["name"] for tool in mapped["tools"]] == ["get_weather"] + + @pytest.mark.parametrize( ("model", "base_model"), [("databricks-claude-haiku-5-5", None), ("support-bot-endpoint", "databricks/databricks-claude-haiku-5-5")], diff --git a/tests/unit/llms/deepgram/test_deepgram_common_utils.py b/tests/unit/llms/deepgram/test_deepgram_common_utils.py index 4c1b4603076..a1672e965bb 100644 --- a/tests/unit/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/unit/llms/deepgram/test_deepgram_common_utils.py @@ -285,6 +285,9 @@ NOVA_2_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-2" [ pytest.param(NOVA_3_URL, (), True, id="streaming entry present"), pytest.param(f"{NOVA_3_URL}&language=multi", (), True, id="multilingual entry present"), + pytest.param( + "wss://api.deepgram.com/v1/listen?model=nova-3-medical", (), True, id="medical streaming entry present" + ), pytest.param(NOVA_2_URL, (), False, id="only the pre-recorded entry"), pytest.param(f"{NOVA_2_URL}&language=multi", ("deepgram/streaming/nova-2",), False, id="needs multilingual"), pytest.param("wss://api.deepgram.com/v1/listen?model=nova-99-unmapped", (), False, id="nothing priced"), @@ -295,7 +298,8 @@ NOVA_2_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-2" def test_deepgram_listen_is_priced( monkeypatch: pytest.MonkeyPatch, upstream_url: str, extra_rows: tuple[str, ...], expected: bool ): - """The bundled map prices only nova-3 for streaming; nova-2 has a pre-recorded row, which must never count.""" + """The bundled map prices nova-3 and nova-3-medical for streaming; nova-2 has only a pre-recorded row, which must + never count.""" monkeypatch.delitem(litellm.model_cost, "deepgram/streaming/nova-2", raising=False) assert "deepgram/nova-2" in litellm.model_cost for row in extra_rows: diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index 56af47aeebe..160ed3950b3 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -955,6 +955,51 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} + + +def test_a_hosted_mcp_tool_answers_400_on_deepseek_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + DeepSeekChatConfig().map_openai_params( + {"tools": [_function_tool("shell"), _HOSTED_MCP_TOOL]}, {}, "deepseek-v4-flash", False + ) + assert raised.value.status_code == 400 + assert "deepseek-v4-flash on deepseek " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_deepseek_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + result: Final = DeepSeekChatConfig().map_openai_params( + {"tools": [_HOSTED_MCP_TOOL, _function_tool("shell")], "tool_choice": "auto"}, {}, "deepseek-v4-flash", True + ) + assert result["tools"] == [_function_tool("shell")] + assert result["tool_choice"] == "auto" + assert any( + "`deepwiki`" in record.getMessage() and " on deepseek (" in record.getMessage() for record in caplog.records + ) + + +def test_deepseek_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + result: Final = DeepSeekChatConfig().map_openai_params( + {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required", "parallel_tool_calls": True}, + {}, + "deepseek-v4-flash", + True, + ) + assert "tools" not in result + assert "tool_choice" not in result + assert "parallel_tool_calls" not in result + + def _budget_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: return tuple( r.getMessage() diff --git a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index 7fc6ac2c436..6538ccf1cac 100644 --- a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -172,10 +172,10 @@ def test_responses_call_forwards_previous_response_id_and_store() -> None: def test_responses_call_folds_developer_items_into_instructions() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) with patch(HTTPX_CLIENT_FACTORY, return_value=client): litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", input=[ {"role": "user", "content": "Hi there"}, {"role": "developer", "content": "Answer with exactly one word."}, @@ -214,6 +214,68 @@ def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_acc ) +@respx.mock +@pytest.mark.parametrize( + "model", + [ + "fireworks_ai/accounts/fireworks/models/kimi-k3", + "fireworks_ai/accounts/fireworks/models/glm-5p3", + "fireworks_ai/accounts/fireworks/routers/kimi-k3-fast", + "fireworks_ai/glm-5p3-us", + "fireworks_ai/models/kimi-k3", + "fireworks_ai/routers/kimi-k3-fast", + ], +) +def test_responses_call_keeps_developer_items_in_place_for_a_model_the_cost_map_declares_accepts_them( + model: str, +) -> None: + route: Final = _fireworks_responses_route(model.removeprefix("fireworks_ai/")) + litellm.responses( + model=model, + instructions="You are Codex", + input=[ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body["instructions"] == "You are Codex" + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Answer with exactly one word.", "type": "message"}, + {"role": "user", "content": "What is the capital of France?"}, + ) + + +@respx.mock +def test_responses_call_folds_developer_items_when_the_loaded_cost_map_turns_the_flag_off( + declare_mid_conversation_system: Callable[[str, bool], None], +) -> None: + declare_mid_conversation_system("fireworks_ai/accounts/fireworks/models/kimi-k3", False) + route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3") + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input=[ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": "Paris"}, + {"role": "developer", "content": "Reply with PONG."}, + ], + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body["instructions"] == "Answer with exactly one word." + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": "Paris"}, + {"role": "user", "content": "Reply with PONG.", "type": "message"}, + ) + + @respx.mock @pytest.mark.parametrize( ("model_info", "closing_role"), @@ -222,9 +284,9 @@ def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_acc def test_responses_call_sends_the_closing_developer_item_as_a_user_turn_only_for_a_declared_system_first_deployment( model_info: Mapping[str, object] | None, closing_role: str ) -> None: - route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3") + route: Final = _fireworks_responses_route("accounts/fireworks/models/qwen3p8-2p4t-a95b") litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", input=[ {"role": "user", "content": "Hi there"}, {"role": "developer", "content": "Answer with exactly one word."}, @@ -402,7 +464,7 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched( developer_item: Final = {"role": "developer", "content": "Answer with exactly one word."} user_item: Final = {"role": "user", "content": "What is the capital of France?"} request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request( - model="accounts/fireworks/models/kimi-k3", + model="accounts/fireworks/models/qwen3p8-2p4t-a95b", input=cast(ResponseInputParam, [developer_item, user_item]), response_api_optional_request_params={"instructions": ["not", "a", "string"]}, litellm_params=GenericLiteLLMParams(), @@ -416,7 +478,7 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched( def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) pydantic_input: Final = cast( ResponseInputParam, [ @@ -435,7 +497,9 @@ def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_outpu ) with patch(HTTPX_CLIENT_FACTORY, return_value=client): litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key" + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", + input=pydantic_input, + api_key="fw-test-key", ) _, _, body = _sent_request(client) assert body["instructions"] == "Answer with exactly one word." diff --git a/tests/unit/llms/gemini/test_gemini_tts.py b/tests/unit/llms/gemini/test_gemini_tts.py index d0ef04e6372..831808fdc96 100644 --- a/tests/unit/llms/gemini/test_gemini_tts.py +++ b/tests/unit/llms/gemini/test_gemini_tts.py @@ -298,11 +298,11 @@ class TestGeminiTTSTransformation: # Test invalid format non_default_params = { - "audio": {"voice": "Kore", "format": "wav"} # Invalid format + "audio": {"voice": "Kore", "format": "mp3"} # Invalid format } optional_params = {} - with pytest.raises(ValueError, match="Unsupported audio format for Gemini TTS models"): + with pytest.raises(litellm.UnsupportedParamsError, match="audio format 'mp3'"): config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, diff --git a/tests/unit/llms/gigachat/chat/test_transformation.py b/tests/unit/llms/gigachat/chat/test_transformation.py index 9116a4d1979..4d8a916ea87 100644 --- a/tests/unit/llms/gigachat/chat/test_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_transformation.py @@ -7,10 +7,12 @@ Run with: pytest tests/unit/llms/gigachat/chat/test_transformation.py -v import asyncio import importlib +from typing import Final import pytest import litellm +from litellm.llms.gigachat.chat.transformation import GigaChatConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -575,3 +577,56 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), "cohere_key": getattr(litellm, "cohere_key", None), } + + +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def test_a_hosted_mcp_tool_answers_400_on_gigachat_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + GigaChatConfig().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "GigaChat-2-Max on gigachat " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_gigachat_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = GigaChatConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=True, + ) + assert [function["name"] for function in mapped["functions"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and " on gigachat (" in record.getMessage() for record in caplog.records + ) + + +def test_gigachat_sends_no_functions_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = GigaChatConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "auto"}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=True, + ) + assert "functions" not in mapped + assert "function_call" not in mapped diff --git a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index fc5ae3631ff..94f1873faba 100644 --- a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -3,14 +3,15 @@ import json from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest - from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, ) from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig +from litellm.types.utils import ModelResponse def test_hosted_vllm_chat_transformation_file_url(): @@ -401,6 +402,64 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema(): assert tools[0]["function"]["parameters"] == input_schema +def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResponse: + payload: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "glm", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "logprobs": None, + "finish_reason": "stop", + **choice_extras, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + raw_response: Final = httpx.Response( + 200, json=payload, request=httpx.Request("POST", "http://vllm.local/v1/chat/completions") + ) + return HostedVLLMChatConfig().transform_response( + model="hosted_vllm/glm", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + "choice_extras", + [ + {"stop_reason": ", 5", "token_ids": [16, 11, 220]}, + {"stop_reason": 128009}, + {"stop_reason": None}, + ], +) +def test_hosted_vllm_keeps_vllm_choice_fields_on_the_choice(choice_extras: dict[str, object]): + response: Final = _transform_vllm_chat_response(choice_extras) + + dumped_choice: Final[dict[str, dict[str, object]]] = json.loads(response.model_dump_json())["choices"][0] + for field, value in choice_extras.items(): + assert dumped_choice[field] == value + assert dumped_choice["provider_specific_fields"][field] == value + + +def test_hosted_vllm_does_not_invent_vllm_choice_fields(): + response: Final = _transform_vllm_chat_response({}) + + dumped_choice: Final[dict[str, object]] = json.loads(response.model_dump_json())["choices"][0] + assert "stop_reason" not in dumped_choice + assert "token_ids" not in dumped_choice + + def _strict_weather_tool() -> dict[str, object]: return { "type": "function", diff --git a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py index ea1e04291ca..7040fbe5fcc 100644 --- a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -1,12 +1,12 @@ import json from typing import Final -from unittest.mock import patch, MagicMock, AsyncMock +from unittest.mock import AsyncMock, MagicMock, patch - -import litellm +import httpx import pytest import respx +import litellm MOCK_EMBEDDING_RESPONSE = [[0.1, 0.2, 0.3, 0.4, 0.5]] @@ -52,12 +52,10 @@ def mock_embedding_async_http_handler(): class TestHuggingFaceEmbedding: @pytest.fixture(autouse=True) def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler, mock_hf_config_fetch): - self.mock_get_task_patcher = patch( - "litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model" - ) + self.mock_get_task_patcher = patch("litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model") self.mock_get_task = self.mock_get_task_patcher.start() - def mock_get_task_side_effect(model, task_type, api_base): + def mock_get_task_side_effect(model, task_type, api_base, timeout=None): if task_type is not None: return task_type return "sentence-similarity" @@ -180,7 +178,9 @@ async def test_hf_embedding_sentence_sim_sends_one_similarity_request( monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("HUGGINGFACE_API_KEY", "hf-test-token") with respx.mock(assert_all_called=False) as hf: - hf.get(url__regex=r"https://huggingface\.co/api/models/.*").respond(json={"pipeline_tag": "sentence-similarity"}) + hf.get(url__regex=r"https://huggingface\.co/api/models/.*").respond( + json={"pipeline_tag": "sentence-similarity"} + ) hf.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512}) inference: Final = hf.post( "https://router.huggingface.co/hf-inference/pipeline/sentence-similarity/sentence-transformers/TaylorAI/bge-micro-v2" @@ -197,3 +197,33 @@ async def test_hf_embedding_sentence_sim_sends_one_similarity_request( } assert inference.calls.last.request.headers["Authorization"] == "Bearer hf-test-token" assert response.data == [{"object": "embedding", "index": 0, "embedding": 0.7708950042724609}] + + +HF_MODEL_LOOKUP_URL: Final = "https://huggingface.co/api/models/BAAI/bge-m3" +HF_EMBED_URL: Final = "https://router.huggingface.co/hf-inference/pipeline/feature-extraction/BAAI/bge-m3" + + +@pytest.mark.respx(assert_all_called=False) +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_huggingface_embedding_sends_the_caller_timeout( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512}) + lookup: Final = respx_mock.get(HF_MODEL_LOOKUP_URL).mock( + return_value=httpx.Response(200, json={"pipeline_tag": "feature-extraction"}) + ) + route: Final = respx_mock.post(HF_EMBED_URL).mock(return_value=httpx.Response(200, json=MOCK_EMBEDDING_RESPONSE)) + request: Final = { + "model": "huggingface/BAAI/bge-m3", + "input": ["hello"], + "api_key": "hf_test", + "timeout": 45, + } + + response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request) + + assert all(call.request.extensions["timeout"]["read"] == 45 for call in lookup.calls) + assert lookup.called + assert route.calls.last.request.extensions["timeout"]["read"] == 45 + assert response.data[0]["embedding"] == MOCK_EMBEDDING_RESPONSE[0] diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 6a80f29ac43..98ca5aee407 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -6,6 +6,7 @@ import sys from typing import Final, cast import pytest +import respx from pydantic import BaseModel sys.path.insert( @@ -1211,6 +1212,72 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) +_DEEPWIKI_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}, +} + + +def _ollama_chat_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("http://ollama.test:11434/api/chat").respond( + json={ + "model": "qwen3:0.6b", + "message": {"role": "assistant", "content": "Sunny"}, + "done": True, + "prompt_eval_count": 5, + "eval_count": 3, + } + ) + + +def test_ollama_chat_answers_400_for_a_hosted_mcp_tool_it_cannot_run( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = _ollama_chat_route(respx_mock) + + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.completion( + model="ollama_chat/qwen3:0.6b", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], + api_base="http://ollama.test:11434", + ) + + assert raised.value.status_code == 400 + assert "qwen3:0.6b on ollama_chat cannot run the hosted MCP tool `deepwiki`" in str(raised.value) + assert not route.called + + +@pytest.mark.parametrize( + ("tools", "sent_tools"), + [([_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], [_WEATHER_TOOL]), ([_DEEPWIKI_MCP_TOOL], None)], + ids=["beside-a-function", "alone"], +) +def test_ollama_chat_drop_params_sends_only_the_tools_ollama_can_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tools: list[object], sent_tools: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = _ollama_chat_route(respx_mock) + + response: Final = litellm.completion( + model="ollama_chat/qwen3:0.6b", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=tools, + api_base="http://ollama.test:11434", + drop_params=True, + ) + + assert json.loads(route.calls.last.request.content).get("tools") == sent_tools + assert response.choices[0].message.content == "Sunny" + + def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: return tuple( r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index edea39bb10c..26addf1b24c 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -993,6 +993,77 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): == "http://gateway.example.com/ollama/api/generate" ) + +_DEEPWIKI_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, +} + + +def _ollama_generate_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("http://ollama.test:11434/api/generate").respond( + json={"model": "llama3.2", "response": "Sunny", "done": True, "prompt_eval_count": 5, "eval_count": 3} + ) + + +def test_ollama_answers_400_for_a_hosted_mcp_tool_instead_of_pasting_it_into_the_prompt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = _ollama_generate_route(respx_mock) + + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.completion( + model="ollama/llama3.2", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], + api_base="http://ollama.test:11434", + ) + + assert raised.value.status_code == 400 + assert "cannot run the hosted MCP tool `deepwiki`" in str(raised.value) + assert not route.called + + +@pytest.mark.parametrize( + ("tools", "prompt_names_weather"), + [([_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], True), ([_DEEPWIKI_MCP_TOOL], False)], + ids=["beside-a-function", "alone"], +) +def test_ollama_drop_params_keeps_a_hosted_mcp_tool_out_of_the_prompt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tools: list[object], prompt_names_weather: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = _ollama_generate_route(respx_mock) + + litellm.completion( + model="ollama/llama3.2", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=tools, + api_base="http://ollama.test:11434", + drop_params=True, + ) + + body: Final = json.loads(route.calls.last.request.content) + prompt: Final = body["prompt"] + assert (body.get("format") == "json") is prompt_names_weather + assert "deepwiki" not in prompt + assert ("get_weather" in prompt) is prompt_names_weather + assert prompt.startswith("### User:\nweather in Paris?") + assert prompt_names_weather or prompt == "### User:\nweather in Paris?\n\n" + + def _reasoning_effort_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: return tuple( r.getMessage() diff --git a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py index 79fe7061297..263c2f7618d 100644 --- a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py @@ -1659,6 +1659,125 @@ class TestOpenAIPromptCacheBreakpointChatPath: ] +_CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM: Final = ( + ({"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}), + ({"mode": "explicit", "ttl": "30m"}, {"mode": "explicit"}), + ({"mode": "explicit"}, {"mode": "explicit"}), + ({"mode": "implicit"}, None), + ({}, None), + ("explicit", None), +) +_CLIENT_BREAKPOINT_IDS: Final = ("ttl-1h", "ttl-30m", "explicit", "implicit", "empty", "string") + + +class TestChatCompletionsClientBreakpointUnderDropParams: + """ + Live OpenAI Chat Completions (gpt-6.1-sol, gpt-6.1-terra, gpt-6-sol, gpt-5.6, 2026-10-28) answers 400 on a + content part's prompt_cache_breakpoint carrying any ttl or a mode other than "explicit". + """ + + @staticmethod + def _request(marker: object, drop_params: bool) -> dict[str, object]: + return OpenAIGPTConfig().transform_request( + model="gpt-5.6", + messages=[ + {"role": "system", "content": "Answer briefly"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this", "prompt_cache_breakpoint": marker}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + {"type": "text", "text": "and nothing else"}, + ], + }, + ], + optional_params={}, + litellm_params={"custom_llm_provider": "openai", "drop_params": drop_params}, + headers={}, + ) + + @pytest.mark.parametrize( + ("marker", "wire_marker"), _CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM, ids=_CLIENT_BREAKPOINT_IDS + ) + def test_drop_params_cuts_the_breakpoint_down_to_what_chat_completions_accepts( + self, marker: object, wire_marker: dict[str, str] | None + ): + text_part: Final = {"type": "text", "text": "Describe this"} + + assert self._request(marker, drop_params=True)["messages"] == [ + {"role": "system", "content": "Answer briefly"}, + { + "role": "user", + "content": [ + text_part if wire_marker is None else {**text_part, "prompt_cache_breakpoint": wire_marker}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + {"type": "text", "text": "and nothing else"}, + ], + }, + ] + + @pytest.mark.parametrize( + "marker", [marker for marker, _ in _CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM], ids=_CLIENT_BREAKPOINT_IDS + ) + def test_without_drop_params_the_client_breakpoint_reaches_openai_unchanged(self, marker: object): + assert self._request(marker, drop_params=False)["messages"][1]["content"][0] == { + "type": "text", + "text": "Describe this", + "prompt_cache_breakpoint": marker, + } + + def test_cache_control_still_marks_a_part_whose_client_breakpoint_was_dropped(self): + request: Final = OpenAIGPTConfig().transform_request( + model="gpt-5.6", + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Describe this", + "cache_control": {"type": "ephemeral"}, + "prompt_cache_breakpoint": {"mode": "implicit"}, + } + ], + } + ], + optional_params={}, + litellm_params={"custom_llm_provider": "openai", "drop_params": True}, + headers={}, + ) + + assert request["messages"] == [ + { + "role": "user", + "content": [{"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}}], + } + ] + + + @pytest.mark.parametrize("stray_part", ["Answer briefly", None], ids=["string", "null"]) + def test_drop_params_leaves_a_part_that_is_not_an_object_for_openai_to_reject(self, stray_part: object): + assert OpenAIGPTConfig().transform_request( + model="gpt-5.6", + messages=[ + { + "role": "system", + "content": [ + stray_part, + {"type": "text", "text": "and nothing else", "prompt_cache_breakpoint": {"mode": "implicit"}}, + ], + }, + {"role": "user", "content": "Describe this"}, + ], + optional_params={}, + litellm_params={"custom_llm_provider": "openai", "drop_params": True}, + headers={}, + )["messages"] == [ + {"role": "system", "content": [stray_part, {"type": "text", "text": "and nothing else"}]}, + {"role": "user", "content": "Describe this"}, + ] + + _ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' @@ -2054,3 +2173,45 @@ async def test_openai_request_rewrites_media_parts_without_touching_the_callers_ {"type": "file", "file": {"file_data": "data:application/pdf;base64,JVBERi0=", "filename": "my_file.pdf"}}, ] assert messages == messages_before + + +_OPENAI_DEPLOYMENT_FLAGS: Final = {"opted-in": {"supports_prompt_cache_breakpoint": True}} + + +@pytest.mark.parametrize("registration_order", [("opted-in", "flagless"), ("flagless", "opted-in")]) +async def test_chat_sends_a_breakpoint_by_the_deployments_own_flag_not_a_siblings_on_the_same_model( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, registration_order: tuple[str, str] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + catalog_row_without_a_flag: Final = {"litellm_provider": "openai", "mode": "chat"} + monkeypatch.setitem(litellm.model_cost, "qa-openai-shared", catalog_row_without_a_flag) + chat: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=Response(200, json=_OPENAI_CHAT_REPLY) + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": "openai/qa-openai-shared", + "api_base": "https://api.openai.com/v1", + "api_key": "openai-test-key", + "cache_control_injection_points": [{"location": "message", "role": "system"}], + }, + "model_info": {"id": f"{name}-id", **_OPENAI_DEPLOYMENT_FLAGS.get(name, {})}, + } + for name in registration_order + ], + num_retries=0, + ) + + async def sends_a_breakpoint(name: str) -> bool: + await router.acompletion( + model=name, messages=[{"role": "system", "content": "stable"}, {"role": "user", "content": "hi"}] + ) + return b"prompt_cache_breakpoint" in chat.calls.last.request.content + + assert {name: await sends_a_breakpoint(name) for name in registration_order} == { + "opted-in": True, + "flagless": False, + } diff --git a/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py b/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py index eddfe578256..186ee545323 100644 --- a/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py +++ b/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py @@ -1,10 +1,16 @@ from collections.abc import Mapping +from dataclasses import replace from typing import Final +import httpx import pytest from pydantic import TypeAdapter -from litellm.llms.base_llm.decisions.transformation import ir_to_systemone_response, systemone_request_to_ir +from litellm.llms.base_llm.decisions.transformation import ( + DecisionsResponseValidationError, + ir_to_systemone_response, + systemone_request_to_ir, +) from litellm.llms.openai.decisions.transformation import ( OpenAIDecisionsConfig, ir_to_openai_request, @@ -106,6 +112,35 @@ def test_an_openai_response_reaches_the_caller_unchanged() -> None: assert ir_to_openai_response(parsed, ir, "requested").model_dump(mode="json") == _OPENAI_RESPONSE +@pytest.mark.parametrize("answer_type", ["choice", "score"]) +def test_a_choice_or_score_answer_without_confidence_stays_a_provider_error(answer_type: str) -> None: + answers: Final = [ + {key: value for key, value in answer.items() if answer["type"] != answer_type or key != "confidence"} + for answer in _OPENAI_RESPONSE["answers"] + ] + with pytest.raises(DecisionsResponseValidationError) as raised: + OpenAIDecisionsConfig().transform_decisions_response( + "gpt-6-luna", + "openai", + httpx.Response(200, json={**_OPENAI_RESPONSE, "answers": answers}), + _openai_ir(_OPENAI_REQUEST), + ) + assert raised.value.status_code == 500 + + +def test_provider_extras_stay_off_the_openai_response_unless_the_provider_names_them() -> None: + ir: Final = _openai_ir(_OPENAI_REQUEST) + parsed: Final = replace( + OpenAIDecisionsConfig().parse_response(_OPENAI_RESPONSE, ir), extra={"provider": "upstream", "citations": []} + ) + + assert ir_to_openai_response(parsed, ir, "requested").model_dump(mode="json") == _OPENAI_RESPONSE + assert ir_to_openai_response(parsed, ir, "requested", frozenset({"citations"})).model_dump(mode="json") == { + **_OPENAI_RESPONSE, + "citations": [], + } + + def test_answers_openai_did_not_return_are_refusals() -> None: ir: Final = _openai_ir(_OPENAI_REQUEST) payload: Final = {**_OPENAI_RESPONSE, "answers": [_PREDICATE_ANSWER]} diff --git a/tests/unit/llms/openai/live/test_handler.py b/tests/unit/llms/openai/live/test_handler.py index ec1130376d0..bb7cda99e79 100644 --- a/tests/unit/llms/openai/live/test_handler.py +++ b/tests/unit/llms/openai/live/test_handler.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import Coroutine +from collections.abc import Callable, Coroutine from datetime import datetime from typing import Final, NoReturn, cast @@ -13,11 +13,13 @@ from websockets.frames import Close from websockets.http11 import Response import litellm +from litellm import Router from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.cost_calculator import handle_realtime_stream_cost_calculation from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.openai.live import handler as live_handler from litellm.llms.openai.live.handler import LiveSessionStart, OpenAILiveSessions, parse_live_session_start +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject, Usage _LIVE_EVENT_FRAME_ADAPTER: Final = TypeAdapter(dict[str, object]) @@ -1206,3 +1208,91 @@ async def test_async_live_session_reports_redacted_server_error_and_bounds_close finally: if worker.coroutine is not None: await worker.coroutine + + +_LIVE_SESSION_STARTED: Final = '{"type":"session.started","session":{"model":"gpt-live-1"}}' +_LIVE_SESSION_CLOSED: Final = '{"type":"session.closed","usage":{"seconds":4}}' + + +def _upstream_closing_the_session() -> _FakeUpstream: + return _FakeUpstream(client_frames=(), upstream_frames=(_LIVE_SESSION_STARTED, _LIVE_SESSION_CLOSED)) + + +def _upstream_failing_after_a_frame() -> _FakeUpstream: + return _FakeUpstream( + client_frames=(), + upstream_frames=(_LIVE_SESSION_STARTED,), + close_error=ConnectionClosedError(Close(code=1011, reason="upstream failed"), None), + ) + + +def _upstream_dropping_after_a_frame() -> _FakeUpstream: + return _FakeUpstream( + client_frames=(), upstream_frames=(_LIVE_SESSION_STARTED,), close_error=ConnectionClosedError(None, None) + ) + + +def _upstream_failing_a_client_send() -> _FakeUpstream: + return _FakeUpstream( + client_frames=(), + upstream_frames=(_LIVE_SESSION_STARTED,), + close_response=_LIVE_SESSION_CLOSED, + fail_on_send_type="conversation.item.create", + send_error=OSError("upstream socket write failed"), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_factory", "counted_successes"), + [ + (_upstream_closing_the_session, 1), + (_upstream_failing_after_a_frame, 0), + (_upstream_dropping_after_a_frame, 0), + (_upstream_failing_a_client_send, 0), + ], + ids=["session-closed", "upstream-1011-after-a-frame", "upstream-1006-after-a-frame", "upstream-send-failed"], +) +async def test_live_session_counts_a_router_success_only_when_it_ended_cleanly( + upstream_factory: Callable[[], _FakeUpstream], counted_successes: int +) -> None: + router: Final = Router( + model_list=[ + {"model_name": "live-voice", "litellm_params": {"model": "openai/gpt-live-1"}, "model_info": {"id": "live-a"}} + ], + num_retries=0, + ) + worker: Final = _FakeLoggingWorker() + + async def _arealtime(litellm_logging_obj: Logging, **kwargs: object) -> None: + await OpenAILiveSessions( + logging_worker=worker, websocket_connector=_FakeConnector(upstream_factory()) + ).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=('{"type":"conversation.item.create"}',)), + logging_obj=litellm_logging_obj, + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + await asyncio.wait_for( + router._ageneric_api_call_with_fallbacks( + model="live-voice", + original_function=_arealtime, + litellm_logging_obj=Logging( + model="gpt-live-1", + messages=[], + stream=False, + call_type="_arealtime", + start_time=datetime.now(), + litellm_call_id="live-session", + function_id="live-session", + ), + ), + timeout=5, + ) + + if worker.coroutine is not None: + worker.coroutine.close() + assert get_deployment_successes_for_current_minute(router, "live-a") == counted_successes diff --git a/tests/unit/llms/openai/realtime/test_openai_realtime_handler.py b/tests/unit/llms/openai/realtime/test_openai_realtime_handler.py index b31db62589d..8ed65ef925c 100644 --- a/tests/unit/llms/openai/realtime/test_openai_realtime_handler.py +++ b/tests/unit/llms/openai/realtime/test_openai_realtime_handler.py @@ -1,7 +1,11 @@ +import asyncio import json +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pydantic import TypeAdapter from litellm.llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from litellm.types.realtime import RealtimeQueryParams @@ -526,3 +530,83 @@ async def test_realtime_query_params_preserve_missing_model(monkeypatch): called_kwargs = mock_async_realtime.call_args.kwargs assert called_kwargs["query_params"] == {"intent": "transcription"} + + +class _BackendConnection: + def __init__(self, connection): + self.connection = connection + + async def __aenter__(self): + return self.connection + + async def __aexit__(self, exc_type, exc, tb): + return None + + +def test_translation_url_uses_dedicated_path(): + from litellm.llms.openai.realtime.handler import OpenAIRealtime + + handler = OpenAIRealtime() + url = handler._construct_url( + api_base="https://api.openai.com/v1", + query_params={"model": "gpt-realtime-translate"}, + realtime_mode="translation", + ) + assert url == "wss://api.openai.com/v1/realtime/translations?model=gpt-realtime-translate" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("beta_client", (False, True)) +async def test_translation_websocket_uses_direct_transport(beta_client: bool) -> None: + from litellm.llms.openai.realtime.handler import OpenAIRealtime + + backend: Final = AsyncMock() + pending: Final = asyncio.Event() + + async def receive_backend() -> str: + await pending.wait() + return "" + + backend.recv.side_effect = receive_backend + websocket: Final = MagicMock() + websocket.scope = {"headers": [(b"openai-beta", b"realtime=v1")] if beta_client else []} + websocket.receive_text = AsyncMock( + side_effect=( + json.dumps( + { + "type": "session.update", + "session": {"type": "translation", "audio": {"output": {"language": "fr"}}}, + } + ), + RuntimeError("connection closed"), + ) + ) + websocket.close = AsyncMock() + logging_obj: Final = MagicMock(dispatch_success_handlers=AsyncMock(), dispatch_failure_handlers=AsyncMock()) + handler: Final = OpenAIRealtime() + expected_url: Final = "wss://api.openai.com/v1/realtime/translations?model=gpt-realtime-translate" + assert ( + handler._construct_url( + api_base="https://api.openai.com/v1", + query_params={"model": "gpt-realtime-translate"}, + realtime_mode="translation", + ) + == expected_url + ) + + with patch("websockets.connect", return_value=_BackendConnection(backend)) as connect: + await handler.async_realtime( + model="gpt-realtime-translate", + websocket=websocket, + logging_obj=logging_obj, + api_base="https://api.openai.com/v1", + api_key="sk-test", + query_params={"model": "gpt-realtime-translate"}, + realtime_mode="translation", + ) + + connect.assert_called_once() + assert connect.call_args.args[0] == expected_url + assert "OpenAI-Beta" not in connect.call_args.kwargs["additional_headers"] + sent: Final = TypeAdapter(Mapping[str, object]).validate_json(backend.send.call_args_list[0].args[0]) + assert sent["session"] == {"audio": {"output": {"language": "fr"}}} diff --git a/tests/unit/llms/openai/realtime/test_transcription_sessions.py b/tests/unit/llms/openai/realtime/test_transcription_sessions.py index 3bb9a4ab1ab..f637271d667 100644 --- a/tests/unit/llms/openai/realtime/test_transcription_sessions.py +++ b/tests/unit/llms/openai/realtime/test_transcription_sessions.py @@ -10,7 +10,7 @@ from unittest.mock import AsyncMock, MagicMock import httpx import pytest - +import litellm from litellm.llms.azure.realtime.http_transformation import AzureRealtimeHTTPConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler @@ -21,9 +21,7 @@ from litellm.types.realtime import RealtimeTranscriptionSessionRequest def test_openai_transcription_session_url(): cfg = OpenAIRealtimeHTTPConfig() assert ( - cfg.get_transcription_session_url( - api_base="https://api.openai.com", model="gpt-realtime-whisper" - ) + cfg.get_transcription_session_url(api_base="https://api.openai.com", model="gpt-realtime-whisper") == "https://api.openai.com/v1/realtime/transcription_sessions" ) @@ -32,24 +30,36 @@ def test_openai_transcription_session_url_strips_trailing_v1(): """A /v1 suffix must not be duplicated in the path.""" cfg = OpenAIRealtimeHTTPConfig() assert ( - cfg.get_transcription_session_url( - api_base="https://api.openai.com/v1", model="gpt-realtime-whisper" - ) + cfg.get_transcription_session_url(api_base="https://api.openai.com/v1", model="gpt-realtime-whisper") == "https://api.openai.com/v1/realtime/transcription_sessions" ) -def test_azure_transcription_session_url_uses_deployment_and_api_version(): +@pytest.mark.parametrize( + "api_base", + ["https://my.openai.azure.com", "https://my.openai.azure.com/openai", "https://my.openai.azure.com/openai/v1/"], +) +def test_azure_transcription_session_url_uses_deployment_and_api_version(api_base: str): cfg = AzureRealtimeHTTPConfig() url = cfg.get_transcription_session_url( - api_base="https://my.openai.azure.com", + api_base=api_base, model="whisper-deploy", api_version="2025-04-01-preview", ) - assert ( - url - == "https://my.openai.azure.com/openai/realtime/transcription_sessions?api-version=2025-04-01-preview" - ) + assert url == "https://my.openai.azure.com/openai/realtime/transcription_sessions?api-version=2025-04-01-preview" + + +@pytest.mark.parametrize("api_version", [None, "v1", "latest", "preview"]) +@pytest.mark.parametrize("model", ["gpt-realtime-2.1", "gpt-realtime-2", "gpt-realtime-2-2026-05-06"]) +def test_azure_ga_realtime_http_urls(api_version, model: str, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + cfg = AzureRealtimeHTTPConfig() + base = "https://my.openai.azure.com" + + assert cfg.get_complete_url(base, model, api_version) == (f"{base}/openai/v1/realtime/client_secrets") + assert cfg.get_realtime_calls_url(base, model, api_version) == (f"{base}/openai/v1/realtime/calls") + litellm.get_model_info.cache_clear() def test_request_resolves_model_returns_none_when_both_absent(): @@ -141,6 +151,27 @@ async def test_client_secret_handler_still_targets_client_secrets_url(): assert kwargs["url"] == "https://api.openai.com/v1/realtime/client_secrets" +@pytest.mark.asyncio +async def test_azure_client_secret_routes_on_a_provider_qualified_session_model(): + mock_response = MagicMock(spec=httpx.Response) + mock_client = MagicMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=mock_response) + + result = await litellm.acreate_realtime_client_secret( + session={"type": "realtime", "model": "azure/gpt-realtime-2.1"}, + api_base="https://my.openai.azure.com", + api_key="azure-test-key", + api_version="v1", + client=mock_client, + ) + + assert result is mock_response + request = mock_client.post.call_args.kwargs + assert request["url"] == "https://my.openai.azure.com/openai/v1/realtime/client_secrets" + assert request["headers"]["api-key"] == "azure-test-key" + assert request["json"]["session"]["model"] == "gpt-realtime-2.1" + + @pytest.mark.asyncio async def test_sdk_fn_routes_openai_transcription_session(monkeypatch): """ @@ -169,18 +200,14 @@ async def test_sdk_fn_routes_openai_transcription_session(monkeypatch): assert kwargs["url"].endswith("/v1/realtime/transcription_sessions") # The litellm-only routing hint must not be forwarded upstream. assert "model" not in kwargs["json"] - assert kwargs["json"]["input_audio_transcription"] == { - "model": "gpt-realtime-whisper" - } + assert kwargs["json"]["input_audio_transcription"] == {"model": "gpt-realtime-whisper"} def test_append_query_params_skips_existing_keys(): from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler url = "wss://example.com/v1/realtime?model=gpt-4o" - result = BaseLLMHTTPHandler._append_query_params( - url, {"model": "ignored", "intent": "transcription"} - ) + result = BaseLLMHTTPHandler._append_query_params(url, {"model": "ignored", "intent": "transcription"}) assert "model=ignored" not in result assert "intent=transcription" in result diff --git a/tests/unit/llms/openai/realtime/test_translation.py b/tests/unit/llms/openai/realtime/test_translation.py new file mode 100644 index 00000000000..525782704c1 --- /dev/null +++ b/tests/unit/llms/openai/realtime/test_translation.py @@ -0,0 +1,185 @@ +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from litellm.llms.azure.realtime.http_transformation import AzureRealtimeHTTPConfig +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.openai.realtime.http_transformation import OpenAIRealtimeHTTPConfig +from litellm.types.realtime import RealtimeSessionConfig +from litellm.types.router import GenericLiteLLMParams + + +def test_azure_realtime_http_headers_use_ad_token_without_api_key(): + config = AzureRealtimeHTTPConfig() + params = GenericLiteLLMParams(azure_ad_token="test-ad-token") + + assert config.get_extra_headers(params, "", {"X-Trace": "trace"}) == { + "X-Trace": "trace", + "Authorization": "Bearer test-ad-token", + } + assert config.get_extra_headers(params, "test-api-key", {"X-Trace": "trace"}) == {"X-Trace": "trace"} + + +def test_realtime_session_config_supports_translation_and_live_transcription_fields(): + session = RealtimeSessionConfig( + type="translation", + model="gpt-realtime-translate", + audio={ + "input": { + "transcription": { + "model": "gpt-live-transcribe", + "delay": "minimal", + "languages": ["en", "fr"], + "keywords": ["LiteLLM"], + } + }, + "output": {"language": "es"}, + }, + ) + + assert session.audio is not None + assert session.audio.input is not None + assert session.audio.input.transcription is not None + assert session.audio.input.transcription.delay == "minimal" + assert session.audio.input.transcription.languages == ["en", "fr"] + assert session.audio.input.transcription.keywords == ["LiteLLM"] + assert session.audio.output is not None + assert session.audio.output.language == "es" + + +@pytest.mark.parametrize( + "api_base,expected", + [ + ( + "https://api.openai.com", + "https://api.openai.com/v1/realtime/translations/client_secrets", + ), + ( + "https://api.openai.com/v1", + "https://api.openai.com/v1/realtime/translations/client_secrets", + ), + ], +) +def test_openai_translation_client_secret_url(api_base: str, expected: str): + config = OpenAIRealtimeHTTPConfig() + assert config.get_translation_client_secret_url(api_base, "gpt-realtime-translate") == expected + + +def test_openai_translation_calls_url(): + config = OpenAIRealtimeHTTPConfig() + assert ( + config.get_translation_calls_url("https://api.openai.com/v1", "gpt-realtime-translate") + == "https://api.openai.com/v1/realtime/translations/calls" + ) + + +@pytest.mark.parametrize("suffix", ("", "/", "/openai", "/openai/v1/")) +def test_azure_translation_urls_use_ga_paths(suffix: str): + config: Final = AzureRealtimeHTTPConfig() + api_base: Final = f"https://example.openai.azure.com{suffix}" + assert ( + config.get_translation_client_secret_url(api_base, "translate-deployment") + == "https://example.openai.azure.com/openai/v1/realtime/translations/client_secrets" + ) + assert ( + config.get_translation_calls_url(api_base, "translate-deployment") + == "https://example.openai.azure.com/openai/v1/realtime/translations/calls" + ) + + +@pytest.mark.asyncio +async def test_translation_client_secret_uses_custom_translation_path(): + client = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json={"value": "ek_test"}, + request=httpx.Request("POST", "https://api.openai.com/v1/realtime/translations/client_secrets"), + ) + ) + logging_obj = MagicMock() + handler = BaseLLMHTTPHandler() + request_data = {"session": {"type": "translation", "model": "gpt-realtime-translate"}} + + response = await handler.async_realtime_translation_client_secret_handler( + api_base="https://api.openai.com", + api_key="sk-test", + request_data=request_data, + logging_obj=logging_obj, + timeout=10, + provider_config=OpenAIRealtimeHTTPConfig(), + model="gpt-realtime-translate", + client=client, + ) + + assert response.status_code == 200 + call = client.post.call_args.kwargs + assert call["url"] == "https://api.openai.com/v1/realtime/translations/client_secrets" + assert call["json"] == request_data + + +@pytest.mark.asyncio +async def test_azure_translation_client_secret_supports_entra_bearer_auth(): + client = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json={"value": "ek_test"}, + request=httpx.Request( + "POST", + "https://example.openai.azure.com/openai/v1/realtime/translations/client_secrets", + ), + ) + ) + handler = BaseLLMHTTPHandler() + + await handler.async_realtime_translation_client_secret_handler( + api_base="https://example.openai.azure.com", + api_key="", + request_data={"session": {"type": "translation", "model": "translate-deployment"}}, + logging_obj=MagicMock(), + timeout=10, + provider_config=AzureRealtimeHTTPConfig(), + model="translate-deployment", + extra_headers={"Authorization": "Bearer entra-token"}, + client=client, + ) + + headers = client.post.call_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer entra-token" + assert "api-key" not in headers + + +@pytest.mark.asyncio +async def test_translation_calls_use_translation_session_and_path(): + client = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock( + return_value=httpx.Response( + 201, + content=b"v=0\r\n", + request=httpx.Request("POST", "https://api.openai.com/v1/realtime/translations/calls"), + ) + ) + logging_obj = MagicMock() + handler = BaseLLMHTTPHandler() + + response = await handler.async_realtime_calls_handler( + api_base="https://api.openai.com", + openai_ephemeral_key="ek_test", + sdp_body=b"v=0\r\n", + logging_obj=logging_obj, + timeout=10, + provider_config=OpenAIRealtimeHTTPConfig(), + model="gpt-realtime-translate", + client=client, + translation=True, + ) + + assert response.status_code == 201 + call = client.post.call_args.kwargs + assert call["url"] == "https://api.openai.com/v1/realtime/translations/calls" + assert call["headers"]["Content-Type"] == "application/sdp" + assert call["content"] == "v=0\r\n" diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index d80a5e835c9..d8954c4072a 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -4547,6 +4547,34 @@ class TestBuildBlockSseChunks: assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} + def test_continuation_numbers_every_event_after_the_last_one_relayed(self): + yielded: Final = [ + {"type": "response.created", "sequence_number": 0, "response": {"id": "resp_live", "output": []}}, + { + "type": "response.output_item.added", + "sequence_number": 1, + "output_index": 0, + "item": {"id": "msg_orig", "type": "message", "role": "assistant"}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 2, + "item_id": "msg_orig", + "output_index": 0, + "content_index": 0, + "delta": "partial", + }, + ] + + payloads: Final = self._payloads( + OpenAIResponsesHandler().build_block_sse_chunks( + self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded + ) + ) + + assert payloads[-1]["type"] == "response.completed" + assert [payload.get("sequence_number") for payload in payloads] == list(range(3, 3 + len(payloads))) + def test_continuation_reads_usage_from_typed_completed_event(self): from litellm.types.llms.openai import ( ResponseCompletedEvent, diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 41c54bf6cc9..6b36a683df7 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -4393,8 +4393,9 @@ async def test_extra_body_merges_with_request_data(extra_body_mock_response_data assert request_body["custom_field"] == "custom_value" -_MALFORMED_MARKERS: Final = ("yes", 1, {"mode": "bogus"}, {"mode": "explicit", "ttl": "1h"}) -_MALFORMED_MARKER_IDS: Final = ("string", "int", "bogus-mode", "bad-ttl") +_MALFORMED_MARKERS: Final = ("yes", 1, {"mode": "bogus"}, {"mode": "bogus", "ttl": "30m"}) +_MALFORMED_MARKER_IDS: Final = ("string", "int", "bogus-mode", "bogus-mode-with-ttl") +_UNSUPPORTED_TTL_MARKER: Final = {"mode": "explicit", "ttl": "1h"} _VALID_MARKERS: Final = ({"mode": "explicit"}, {"mode": "explicit", "ttl": "30m"}) _VALID_MARKER_IDS: Final = ("explicit", "explicit-30m") _PARTS: Final[Mapping[str, Mapping[str, object]]] = { @@ -4480,6 +4481,76 @@ def test_valid_prompt_cache_breakpoint_reaches_the_wire_under_both_settings( assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=drop_params)) == input +@pytest.mark.parametrize("config", _CONFIGS, ids=_CONFIG_IDS) +@pytest.mark.parametrize("part_type", tuple(_PARTS)) +def test_an_unsupported_ttl_is_dropped_and_the_explicit_mode_kept_under_drop_params( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, part_type: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message(part_type, _UNSUPPORTED_TTL_MARKER) + trimmed: Final = _transformed_input(config, input, GenericLiteLLMParams(drop_params=True)) + assert _single_part(trimmed) == {**_PARTS[part_type], "prompt_cache_breakpoint": {"mode": "explicit"}} + assert _transformed_input(config, input, GenericLiteLLMParams(drop_params=False)) == input + + +@pytest.mark.parametrize( + "config, marker, wire_marker", + [ + (OpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit", "ttl": "12h"}), + (AzureOpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "30m"}, {"mode": "explicit", "ttl": "30m"}), + (AzureOpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit"}), + (AzureAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit"}), + ], + ids=("openai-keeps-12h", "azure-keeps-30m", "azure-drops-12h", "azure-ai-drops-12h"), +) +def test_drop_params_keeps_only_a_ttl_the_target_responses_api_accepts( + monkeypatch: pytest.MonkeyPatch, + config: BaseResponsesAPIConfig, + marker: Mapping[str, str], + wire_marker: Mapping[str, str], +) -> None: + """Azure OpenAI Responses answered 400 for a 12h breakpoint ttl and 200 for 30m on 2026-10-29, and Azure AI + Foundry Responses did the same on 2026-10-30.""" + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", marker) + trimmed: Final = _transformed_input(config, input, GenericLiteLLMParams(drop_params=True)) + assert _single_part(trimmed) == {**_PARTS["input_text"], "prompt_cache_breakpoint": wire_marker} + assert _transformed_input(config, input, GenericLiteLLMParams(drop_params=False)) == input + + +@pytest.mark.parametrize( + "config, litellm_params", + [ + (XAIResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True)), + (HostedVLLMResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True, api_base="http://localhost:8000/v1")), + (OpenAIResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True, api_base="http://localhost:8000/v1")), + ], + ids=("xai", "hosted-vllm", "openai-compatible-host"), +) +def test_a_responses_api_not_known_to_take_breakpoints_gets_a_marker_whole_or_not_at_all( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER) + assert _single_part(_transformed_input(config, unsupported_ttl, litellm_params)) == _PARTS["input_text"] + openai_ttl: Final = _marked_message("input_text", {"mode": "explicit", "ttl": "12h"}) + assert _transformed_input(config, openai_ttl, litellm_params) == openai_ttl + + +def test_a_custom_host_opted_in_with_prompt_cache_options_keeps_the_mode_of_a_marker_with_an_unsupported_ttl( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + request: Final = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-6.1-sol", + input=cast("ResponseInputParam", _marked_message("input_text", _UNSUPPORTED_TTL_MARKER)), + response_api_optional_request_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params=GenericLiteLLMParams(drop_params=True, api_base="https://gateway.example/v1"), + headers={}, + ) + assert _single_part(request["input"]) == {**_PARTS["input_text"], "prompt_cache_breakpoint": {"mode": "explicit"}} + + def test_prompt_cache_breakpoint_unknown_key_is_trimmed_under_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "drop_params", False) input: Final = _marked_message("input_text", {"mode": "explicit", "note": "kept"}) @@ -4569,12 +4640,43 @@ def test_a_valid_marker_survives_drop_params_on_a_deployment_name_with_unknown_b monkeypatch.setattr(litellm, "drop_params", False) litellm_params: Final = GenericLiteLLMParams(drop_params=True) valid: Final = _marked_message("input_text", {"mode": "explicit"}) - malformed: Final = _marked_message("input_text", {"mode": "explicit", "ttl": "1h"}) + malformed: Final = _marked_message("input_text", {"mode": "bogus"}) assert _transformed_input_for("litellm-qa-sol-alias", valid, litellm_params, config) == valid dropped: Final = _transformed_input_for("litellm-qa-sol-alias", malformed, litellm_params, config) assert _single_part(dropped) == _PARTS["input_text"] +@pytest.mark.parametrize("ttl", ("30m", "12h")) +@pytest.mark.parametrize( + "config", + (OpenAIResponsesAPIConfig(), AzureOpenAIResponsesAPIConfig(), AzureAIResponsesAPIConfig()), + ids=("openai", "azure", "azure_ai"), +) +def test_a_deployment_with_unknown_breakpoint_support_gets_a_marker_whole_or_not_at_all( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, ttl: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + litellm_params: Final = GenericLiteLLMParams(drop_params=True) + supported_ttl: Final = _marked_message("input_text", {"mode": "explicit", "ttl": ttl}) + unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER) + assert _transformed_input_for("litellm-qa-sol-alias", supported_ttl, litellm_params, config) == supported_ttl + dropped: Final = _transformed_input_for("litellm-qa-sol-alias", unsupported_ttl, litellm_params, config) + assert _single_part(dropped) == _PARTS["input_text"] + + +@pytest.mark.parametrize( + "config", (AzureOpenAIResponsesAPIConfig(), AzureAIResponsesAPIConfig()), ids=("azure", "azure_ai") +) +def test_an_azure_deployment_declaring_a_breakpoint_reader_keeps_the_mode_of_a_marker_with_an_unsupported_ttl( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + litellm_params: Final = GenericLiteLLMParams(drop_params=True, base_model="gpt-6.1-sol") + unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER) + trimmed: Final = _transformed_input_for("litellm-qa-sol-alias", unsupported_ttl, litellm_params, config) + assert _single_part(trimmed) == {**_PARTS["input_text"], "prompt_cache_breakpoint": {"mode": "explicit"}} + + def _declaring_base_model(where: str, base_model: str) -> GenericLiteLLMParams: if where == "litellm_params": return GenericLiteLLMParams(drop_params=True, base_model=base_model) @@ -4615,10 +4717,10 @@ def test_a_model_that_takes_markers_sends_kept_parts_byte_for_byte_and_drops_onl {"prompt_cache_breakpoint": {"mode": "explicit"}, "type": "input_text", "text": "b"}, ] rejected: Final = {"type": "input_text", "text": "c"} - bad_ttl: Final = {"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}} + bogus_mode: Final = {"prompt_cache_breakpoint": {"mode": "bogus"}} input: Final = [ {"type": "message", "role": "user", "content": kept}, - {"type": "message", "role": "user", "content": [{**bad_ttl, **rejected}]}, + {"type": "message", "role": "user", "content": [{**bogus_mode, **rejected}]}, ] transformed: Final = _transformed_input_for("gpt-6.1-sol", input, GenericLiteLLMParams(drop_params=True)) assert json.dumps(transformed) == json.dumps( diff --git a/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py b/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py new file mode 100644 index 00000000000..9bb5fb6d300 --- /dev/null +++ b/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py @@ -0,0 +1,951 @@ +import asyncio +import io +import json +import queue +import wave +from collections.abc import Callable, Iterator +from datetime import datetime +from typing import Final, NamedTuple +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from openai import AsyncAzureOpenAI, AsyncOpenAI, AsyncStream, AzureOpenAI, OpenAI + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( + litellm_exception_mapper, + wrap_transcription_stream, +) +from litellm.litellm_core_utils.redact_messages import REDACTED_BY_LITELLM +from litellm.llms.azure.audio_transcriptions import AzureAudioTranscription +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.openai.transcriptions.gpt_transformation import ( + OpenAIGPTTranscribeAudioTranscriptionConfig, +) +from litellm.llms.openai.transcriptions.handler import OpenAIAudioTranscription +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import TranscriptionResponse +from litellm.utils import get_optional_params_transcription + + +@pytest.fixture +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +def test_gpt_transcribe_config_uses_native_parameters_and_json(): + config = OpenAIGPTTranscribeAudioTranscriptionConfig() + supported = config.get_supported_openai_params("gpt-transcribe") + assert supported == ["prompt", "response_format", "keywords", "languages", "stream"] + + audio_file = io.BytesIO(b"audio") + request = config.transform_audio_transcription_request( + model="gpt-transcribe", + audio_file=audio_file, + optional_params={"keywords": ["LiteLLM"], "languages": ["en", "fr"], "stream": True, "response_format": "json"}, + litellm_params={}, + ) + assert request.data["response_format"] == "json" + assert request.data["keywords"] == ["LiteLLM"] + assert request.data["languages"] == ["en", "fr"] + assert request.data["stream"] is True + + +def test_gpt_transcribe_optional_params_are_preserved(): + params = get_optional_params_transcription( + model="gpt-transcribe", + custom_llm_provider="openai", + keywords=["LiteLLM", "Realtime API"], + languages=["en", "fr"], + stream=True, + ) + assert params == { + "keywords": ["LiteLLM", "Realtime API"], + "languages": ["en", "fr"], + "stream": True, + } + + +def test_transcription_response_preserves_empty_languages(): + response = TranscriptionResponse(text="hello", languages=[]) + assert json.loads(response.model_dump_json())["languages"] == [] + + +def test_sync_transcription_stream_logs_final_text_and_usage_once() -> None: + def send_response(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + 'data: {"type":"transcript.text.delta","delta":"hello "}\n\n' + 'data: {"type":"transcript.text.done","text":"hello world",' + '"usage":{"type":"duration","seconds":2.5},"languages":[{"code":"en"}]}\n\n' + ), + ) + + logging_obj: Final = MagicMock() + with OpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.Client(transport=httpx.MockTransport(send_response)), + ) as client: + stream: Final = client.audio.transcriptions.create( + model="gpt-transcribe", file=("sample.webm", b"audio"), stream=True + ) + wrapped: Final = wrap_transcription_stream(stream, logging_obj, datetime(2026, 1, 1), litellm_exception_mapper({})) + received: Final = tuple(wrapped) + wrapped.close() + + assert tuple(event.type for event in received) == ("transcript.text.delta", "transcript.text.done") + logging_obj.success_handler.assert_called_once() + logged_response: Final = logging_obj.success_handler.call_args.args[0] + assert logged_response.text == "hello world" + assert logged_response.usage.model_dump(exclude_none=True) == {"type": "duration", "seconds": 2.5} + assert json.loads(logged_response.model_dump_json())["languages"] == [{"code": "en"}] + logging_obj.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_openai_handler_returns_native_typed_stream(): + async def send_response(request: httpx.Request) -> httpx.Response: + body = await request.aread() + assert b'name="keywords[]"' in body + assert b'name="languages[]"' in body + assert b'name="stream"' in body + events = ( + {"type": "transcript.text.delta", "delta": "hello "}, + { + "type": "transcript.text.done", + "text": "hello world", + "languages": [], + "usage": { + "type": "tokens", + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + }, + }, + ) + content = "".join(f"data: {json.dumps(event)}\n\n" for event in events).encode() + return httpx.Response(200, content=content, headers={"content-type": "text/event-stream"}) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(send_response)) + openai_client = AsyncOpenAI(api_key="sk-test", base_url="https://example.com/v1", http_client=http_client) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + audio_file = io.BytesIO(b"audio") + audio_file.name = "sample.wav" + + handler = OpenAIAudioTranscription() + result = handler.audio_transcriptions( + model="gpt-transcribe", + audio_file=audio_file, + optional_params={"keywords": ["LiteLLM"], "languages": ["en"], "stream": True}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=10, + max_retries=0, + logging_obj=logging_obj, + api_key="sk-test", + api_base="https://example.com/v1", + client=openai_client, + atranscription=True, + provider_config=OpenAIGPTTranscribeAudioTranscriptionConfig(), + ) + stream = await result + assert isinstance(stream, AsyncStream) + logging_obj.async_success_handler = AsyncMock() + logging_obj.async_failure_handler = AsyncMock() + wrapped_stream = wrap_transcription_stream(stream, logging_obj, datetime.now(), litellm_exception_mapper({})) + received = [event async for event in wrapped_stream] + await wrapped_stream.close() + await openai_client.close() + + assert [event.type for event in received] == ["transcript.text.delta", "transcript.text.done"] + assert received[-1].languages == [] + logging_obj.async_success_handler.assert_awaited_once() + logged_response = logging_obj.async_success_handler.await_args.kwargs["result"] + assert logged_response.text == "hello world" + assert logged_response.languages == () + + +@pytest.mark.asyncio +async def test_atranscription_stream_preserves_duration_for_callback_cost(): + async def send_response(request: httpx.Request) -> httpx.Response: + events = ( + {"type": "transcript.text.delta", "delta": "hello "}, + { + "type": "transcript.text.done", + "text": "hello world", + "usage": { + "type": "tokens", + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + }, + }, + ) + content = "".join(f"data: {json.dumps(event)}\n\n" for event in events).encode() + return httpx.Response(200, content=content, headers={"content-type": "text/event-stream"}) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(send_response)) + openai_client = AsyncOpenAI(api_key="sk-test", base_url="https://example.com/v1", http_client=http_client) + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="gpt-transcribe", messages=[], stream=True, call_type="atranscription", + start_time=datetime.now(), litellm_call_id="transcription-cost-test", function_id="transcription-cost-test", + ) + audio_file = io.BytesIO() + with wave.open(audio_file, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(16000) + wav_file.writeframes(b"\x00\x00" * 16000) + audio_file.name = "sample.wav" + + stream = await litellm.atranscription( + model="openai/gpt-transcribe", + file=audio_file, + stream=True, + client=openai_client, + litellm_logging_obj=logging_obj, + ) + received = [event async for event in stream] + await stream.close() + await openai_client.close() + + assert [event.type for event in received] == ["transcript.text.delta", "transcript.text.done"] + logged_response = logging_obj.model_call_details["async_complete_streaming_response"] + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.000075) + assert logging_obj.model_call_details["standard_logging_object"]["response_cost"] == pytest.approx(0.000075) + assert logged_response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0) + + +@pytest.mark.parametrize("model", ["gpt-transcribe", "azure/gpt-transcribe"]) +def test_gpt_transcribe_rejects_whisper_response_formats(local_model_cost_map: None, model: str) -> None: + audio_file = io.BytesIO(b"audio") + audio_file.name = "sample.wav" + with pytest.raises(litellm.UnsupportedParamsError, match="only supports response_format 'json' or 'text'"): + litellm.transcription( + model=model, + file=audio_file, + response_format="verbose_json", + api_key="sk-test", + ) + + +@pytest.mark.parametrize("drop_globally", (True, False)) +def test_gpt_transcribe_drops_whisper_response_formats_under_drop_params( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, drop_globally: bool +) -> None: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: collects what the mock upstream received + + def send_response(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.read()) + return httpx.Response(200, json={"text": "hello", "usage": {"type": "duration", "seconds": 1}}) + + monkeypatch.setattr(litellm, "drop_params", drop_globally) + audio_file: Final = io.BytesIO(b"audio") + audio_file.name = "sample.wav" + with OpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.Client(transport=httpx.MockTransport(send_response)), + ) as client: + response: Final = litellm.transcription( + model="gpt-transcribe", + file=audio_file, + response_format="verbose_json", + client=client, + **({} if drop_globally else {"drop_params": True}), + ) + + assert response.text == "hello" + assert len(sent_bodies) == 1 + assert b"verbose_json" not in sent_bodies[0] + + +def _openai_transcription_client(transport: httpx.MockTransport) -> OpenAI: + return OpenAI(api_key="sk-test", base_url="https://example.com/v1", http_client=httpx.Client(transport=transport)) + + +def _azure_transcription_client(transport: httpx.MockTransport) -> AzureOpenAI: + return AzureOpenAI( + api_key="azure-test-key", + azure_endpoint="https://example.openai.azure.com", + api_version=litellm.AZURE_DEFAULT_API_VERSION, + http_client=httpx.Client(transport=transport), + ) + + +@pytest.mark.parametrize( + ("model", "make_client"), + (("gpt-transcribe", _openai_transcription_client), ("azure/gpt-transcribe", _azure_transcription_client)), +) +def test_gpt_transcribe_sends_text_response_format_to_the_provider( + local_model_cost_map: None, model: str, make_client: Callable[[httpx.MockTransport], OpenAI | AzureOpenAI] +) -> None: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: collects what the mock upstream received + + def send_response(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.read()) + return httpx.Response(200, text="hello", headers={"content-type": "text/plain"}) + + audio_file: Final = io.BytesIO(b"audio") + audio_file.name = "sample.wav" + with make_client(httpx.MockTransport(send_response)) as client: + response: Final = litellm.transcription( + model=model, + file=audio_file, + response_format="text", + api_key="sk-test", + api_base="https://example.openai.azure.com", + api_version=litellm.AZURE_DEFAULT_API_VERSION, + client=client, + ) + + assert response.text == "hello" + assert len(sent_bodies) == 1 + assert b'name="response_format"\r\n\r\ntext\r\n' in sent_bodies[0] + + +def test_file_transcription_reaches_models_whose_cost_map_lists_other_endpoints( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + sent_requests: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock upstreams received + + def send_gemini_response(request: httpx.Request) -> httpx.Response: + sent_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "interaction-1", + "status": "completed", + "steps": [ + {"type": "model_generation", "content": [{"type": "text", "text": "hello", "annotations": []}]} + ], + }, + ) + + def send_whisper_response(request: httpx.Request) -> httpx.Response: + sent_requests.append(request) + return httpx.Response(200, json={"text": "hello"}) + + gemini_response: Final = litellm.transcription( + model="gemini/gemini-2.5-flash", + file=_one_second_wav(), + api_key="gemini-test-key", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(send_gemini_response))), + ) + monkeypatch.setattr(litellm, "client_session", httpx.Client(transport=httpx.MockTransport(send_whisper_response))) + whisper_response: Final = litellm.transcription( + model="openai/gpt-realtime-whisper", + file=_one_second_wav(), + api_key="self-hosted-whisper-key", + api_base="https://whisper.self-hosted.example/v1", + ) + + assert gemini_response.text == "hello" + assert whisper_response.text == "hello" + assert tuple(str(request.url.host) for request in sent_requests) == ( + "generativelanguage.googleapis.com", + "whisper.self-hosted.example", + ) + + +def test_azure_gpt_transcribe_uses_deployment_scoped_route(): + def send_response(request: httpx.Request) -> httpx.Response: + assert str(request.url) == ( + "https://example.openai.azure.com/openai/deployments/gpt-transcribe/audio/transcriptions" + f"?api-version={litellm.AZURE_DEFAULT_API_VERSION}" + ) + return httpx.Response( + 200, + json={"text": "hello", "languages": [{"code": "en"}], "usage": {"type": "duration", "seconds": 1}}, + ) + + http_client = httpx.Client(transport=httpx.MockTransport(send_response)) + client = AzureOpenAI( + api_key="azure-test-key", + azure_endpoint="https://example.openai.azure.com", + api_version=litellm.AZURE_DEFAULT_API_VERSION, + http_client=http_client, + ) + audio_file = io.BytesIO(b"audio") + audio_file.name = "sample.wav" + + response = AzureAudioTranscription().audio_transcriptions( + model="gpt-transcribe", + audio_file=audio_file, + optional_params={"response_format": "json"}, + logging_obj=MagicMock(), + model_response=TranscriptionResponse(), + timeout=10, + max_retries=0, + api_key="azure-test-key", + api_base="https://example.openai.azure.com", + api_version=litellm.AZURE_DEFAULT_API_VERSION, + client=client, + ) + + assert response.text == "hello" + assert response.languages is not None + assert [language.code for language in response.languages] == ["en"] + client.close() + + +@pytest.mark.asyncio +async def test_azure_gpt_transcribe_sends_language_hints_in_sdk_extra_body(): + async def send_response(request: httpx.Request) -> httpx.Response: + body = await request.aread() + assert b'name="keywords[]"' in body + assert b'name="languages[]"' in body + return httpx.Response(200, json={"text": "hello"}) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(send_response)) + client = AsyncAzureOpenAI( + api_key="azure-test-key", + azure_endpoint="https://example.openai.azure.com", + api_version="2025-04-01-preview", + http_client=http_client, + ) + audio_file = io.BytesIO(b"audio") + audio_file.name = "sample.wav" + + response = await AzureAudioTranscription().audio_transcriptions( + model="gpt-transcribe", + audio_file=audio_file, + optional_params={"keywords": ["LiteLLM"], "languages": ["en"]}, + logging_obj=MagicMock(), + model_response=TranscriptionResponse(), + timeout=10, + max_retries=0, + api_key="azure-test-key", + api_base="https://example.openai.azure.com", + api_version="2025-04-01-preview", + client=client, + atranscription=True, + ) + + assert response.text == "hello" + await client.close() + + +def _one_second_wav() -> io.BytesIO: + audio_file: Final = io.BytesIO() + with wave.open(audio_file, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(16000) + wav_file.writeframes(b"\x00\x00" * 16000) + audio_file.name = "sample.wav" + audio_file.seek(0) + return audio_file + + +class _LoggedSpend(NamedTuple): + cost: float | None + usage: dict[str, object] | None + + +def _logged_spend(kwargs: dict[str, object]) -> _LoggedSpend: + payload: Final = kwargs["standard_logging_object"] + return _LoggedSpend(payload["response_cost"], payload["metadata"]["usage_object"]) + + +class _SuccessSpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.async_spends: Final[asyncio.Queue[_LoggedSpend]] = asyncio.Queue() + self.sync_spends: Final[queue.Queue[_LoggedSpend]] = queue.Queue() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + await self.async_spends.put(_logged_spend(kwargs)) + + def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.sync_spends.put(_logged_spend(kwargs)) + + +@pytest.mark.asyncio +async def test_atranscription_bills_the_transcript_when_the_model_ignores_stream( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + async def send_response(request: httpx.Request) -> httpx.Response: + return _duration_priced_transcript(streamed=False) + + recorder: Final = _SuccessSpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + client: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(send_response)), + ) + await litellm.atranscription(model="openai/whisper-1", file=_one_second_wav(), client=client) + unstreamed_spend: Final = await asyncio.wait_for(recorder.async_spends.get(), timeout=5) + ignored_stream: Final = await litellm.atranscription( + model="openai/whisper-1", file=_one_second_wav(), stream=True, client=client + ) + ignored_stream_spend: Final = await asyncio.wait_for(recorder.async_spends.get(), timeout=5) + await client.close() + + assert isinstance(ignored_stream, TranscriptionResponse) + assert ignored_stream.text == "hello world" + assert unstreamed_spend.cost is not None + assert unstreamed_spend.cost > 0 + assert unstreamed_spend.usage + assert ignored_stream_spend == unstreamed_spend + + +def test_azure_transcription_bills_the_transcript_when_the_deployment_ignores_stream( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + def send_response(request: httpx.Request) -> httpx.Response: + return _duration_priced_transcript(streamed=False) + + recorder: Final = _SuccessSpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + client: Final = AzureOpenAI( + api_key="azure-test-key", + azure_endpoint="https://example.openai.azure.com", + api_version=litellm.AZURE_DEFAULT_API_VERSION, + http_client=httpx.Client(transport=httpx.MockTransport(send_response)), + ) + azure_params: Final = { + "client": client, + "api_key": "azure-test-key", + "api_base": "https://example.openai.azure.com", + "api_version": litellm.AZURE_DEFAULT_API_VERSION, + } + litellm.transcription(model="azure/whisper-1", file=_one_second_wav(), **azure_params) + unstreamed_spend: Final = recorder.sync_spends.get(timeout=5) + ignored_stream: Final = litellm.transcription( + model="azure/whisper-1", file=_one_second_wav(), stream=True, **azure_params + ) + ignored_stream_spend: Final = recorder.sync_spends.get(timeout=5) + client.close() + + assert isinstance(ignored_stream, TranscriptionResponse) + assert ignored_stream.text == "hello world" + assert unstreamed_spend.cost is not None + assert unstreamed_spend.cost > 0 + assert unstreamed_spend.usage + assert ignored_stream_spend == unstreamed_spend + + +def _duration_priced_transcript(streamed: bool) -> httpx.Response: + usage: Final = {"type": "duration", "seconds": 2} + if not streamed: + return httpx.Response(200, json={"text": "hello world", "usage": usage}) + events: Final = ( + {"type": "transcript.text.delta", "delta": "hello world"}, + {"type": "transcript.text.done", "text": "hello world", "usage": usage}, + ) + return httpx.Response( + 200, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(), + headers={"content-type": "text/event-stream"}, + ) + + +def test_whisper_stream_request_answered_with_json_keeps_the_unstreamed_response_shape( + local_model_cost_map: None, +) -> None: + verbose_transcript: Final = { + "task": "transcribe", + "language": "english", + "duration": 1.0, + "text": "hello world", + "usage": {"type": "duration", "seconds": 1}, + } + with _openai_transcription_client( + httpx.MockTransport(lambda request: httpx.Response(200, json=verbose_transcript)) + ) as client: + unstreamed: Final = litellm.transcription(model="openai/whisper-1", file=_one_second_wav(), client=client) + ignored_stream: Final = litellm.transcription( + model="openai/whisper-1", file=_one_second_wav(), stream=True, client=client + ) + + assert isinstance(ignored_stream, TranscriptionResponse) + assert json.loads(ignored_stream.model_dump_json()) == json.loads(unstreamed.model_dump_json()) + + +@pytest.mark.parametrize("requested_stream", ("true", "True", "1", "yes", "on")) +def test_transcription_stream_requested_as_a_form_string_streams_like_a_bool( + local_model_cost_map: None, requested_stream: str +) -> None: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: collects what the mock upstream received + + def send_response(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.read()) + return _duration_priced_transcript(streamed=_asks_for_stream(sent_bodies[-1])) + + with _openai_transcription_client(httpx.MockTransport(send_response)) as client: + stream: Final = litellm.transcription( + model="openai/gpt-transcribe", file=_one_second_wav(), stream=requested_stream, client=client + ) + events: Final = [event.type for event in stream] + + assert events == ["transcript.text.delta", "transcript.text.done"] + assert len(sent_bodies) == 1 + + +def test_transcription_stream_that_is_not_a_boolean_answers_400_without_calling_the_upstream( + local_model_cost_map: None, +) -> None: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: collects what the mock upstream received + + def send_response(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.read()) + return _duration_priced_transcript(streamed=False) + + with _openai_transcription_client(httpx.MockTransport(send_response)) as client: + with pytest.raises(litellm.BadRequestError) as raised: + litellm.transcription(model="openai/gpt-transcribe", file=_one_second_wav(), stream="maybe", client=client) + + assert raised.value.status_code == 400 + assert sent_bodies == [] + + +def test_transcription_with_an_empty_stream_form_value_runs_unstreamed_without_a_stream_field( + local_model_cost_map: None, +) -> None: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: collects what the mock upstream received + + def send_response(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.read()) + return _duration_priced_transcript(streamed=False) + + with _openai_transcription_client(httpx.MockTransport(send_response)) as client: + response: Final = litellm.transcription( + model="openai/gpt-transcribe", file=_one_second_wav(), stream="", client=client + ) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "hello world" + assert len(sent_bodies) == 1 + assert b'name="stream"' not in sent_bodies[0] + + +def test_transcription_stream_that_is_not_a_boolean_is_dropped_under_drop_params(local_model_cost_map: None) -> None: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: collects what the mock upstream received + + def send_response(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.read()) + return _duration_priced_transcript(streamed=False) + + with _openai_transcription_client(httpx.MockTransport(send_response)) as client: + response: Final = litellm.transcription( + model="openai/gpt-transcribe", file=_one_second_wav(), stream="maybe", drop_params=True, client=client + ) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "hello world" + assert len(sent_bodies) == 1 + assert b'name="stream"' not in sent_bodies[0] + + +def _asks_for_stream(request_body: bytes) -> bool: + return b'name="stream"\r\n\r\ntrue' in request_body + + +@pytest.mark.asyncio +async def test_atranscription_stream_bills_the_provider_duration_like_the_unstreamed_call( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + async def send_response(request: httpx.Request) -> httpx.Response: + return _duration_priced_transcript(streamed=_asks_for_stream(await request.aread())) + + recorder: Final = _SuccessSpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + client: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(send_response)), + ) + await litellm.atranscription(model="openai/gpt-transcribe", file=_one_second_wav(), client=client) + unstreamed_spend: Final = await asyncio.wait_for(recorder.async_spends.get(), timeout=5) + stream: Final = await litellm.atranscription( + model="openai/gpt-transcribe", file=_one_second_wav(), stream=True, client=client + ) + events: Final = [event.type async for event in stream] + streamed_spend: Final = await asyncio.wait_for(recorder.async_spends.get(), timeout=5) + await client.close() + + assert events == ["transcript.text.delta", "transcript.text.done"] + assert unstreamed_spend.cost is not None + assert unstreamed_spend.cost > 0 + assert unstreamed_spend.usage + assert streamed_spend == unstreamed_spend + + +class _TranscriptRedactor(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="transcript-redactor", event_hook=GuardrailEventHooks.post_call) + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: object, response: object + ) -> None: + if isinstance(response, TranscriptionResponse): + response.text = "[REDACTED]" + + +@pytest.mark.asyncio +async def test_atranscription_stream_under_a_post_call_guardrail_streams_the_guarded_transcript( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + async def send_response(request: httpx.Request) -> httpx.Response: + return _duration_priced_transcript(streamed=_asks_for_stream(await request.aread())) + + monkeypatch.setattr(litellm, "callbacks", [_TranscriptRedactor()]) + client: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(send_response)), + ) + stream: Final = await litellm.atranscription( + model="openai/gpt-transcribe", + file=_one_second_wav(), + stream=True, + client=client, + guardrails=["transcript-redactor"], + ) + events: Final = [(event.type, event.model_dump(exclude={"type"}, exclude_unset=True)) async for event in stream] + await client.close() + + assert events == [ + ("transcript.text.delta", {"delta": "[REDACTED]"}), + ("transcript.text.done", {"text": "[REDACTED]", "usage": {"type": "duration", "seconds": 2}}), + ] + + +class _StreamedTranscriptRecorder(CustomLogger): + def __init__(self, message_logging: bool) -> None: + super().__init__(message_logging=message_logging) + self.logged: Final[asyncio.Queue[frozenset[object]]] = asyncio.Queue() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + logged_responses: Final = ( + response_obj, + *(kwargs.get(key) for key in ("complete_streaming_response", "async_complete_streaming_response")), + ) + await self.logged.put(frozenset(response.text for response in logged_responses if response is not None)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("opt_out", ("global", "callback")) +async def test_atranscription_stream_keeps_the_transcript_out_of_callbacks_that_turned_off_message_logging( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, opt_out: str +) -> None: + async def send_response(request: httpx.Request) -> httpx.Response: + return _duration_priced_transcript(streamed=_asks_for_stream(await request.aread())) + + recorder: Final = _StreamedTranscriptRecorder(message_logging=opt_out != "callback") + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "turn_off_message_logging", opt_out == "global") + client: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(send_response)), + ) + stream: Final = await litellm.atranscription( + model="openai/gpt-transcribe", file=_one_second_wav(), stream=True, client=client + ) + streamed_text: Final = [event.delta async for event in stream if event.type == "transcript.text.delta"] + logged_texts: Final = await asyncio.wait_for(recorder.logged.get(), timeout=5) + await client.close() + + assert streamed_text == ["hello world"] + assert logged_texts == {REDACTED_BY_LITELLM} + + +def test_transcription_stream_bills_the_provider_duration_like_the_unstreamed_call( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + def send_response(request: httpx.Request) -> httpx.Response: + return _duration_priced_transcript(streamed=_asks_for_stream(request.read())) + + recorder: Final = _SuccessSpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + client: Final = OpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.Client(transport=httpx.MockTransport(send_response)), + ) + litellm.transcription(model="openai/gpt-transcribe", file=_one_second_wav(), client=client) + unstreamed_spend: Final = recorder.sync_spends.get(timeout=5) + stream: Final = litellm.transcription( + model="openai/gpt-transcribe", file=_one_second_wav(), stream=True, client=client + ) + events: Final = [event.type for event in stream] + streamed_spend: Final = recorder.sync_spends.get(timeout=5) + client.close() + + assert events == ["transcript.text.delta", "transcript.text.done"] + assert unstreamed_spend.cost is not None + assert unstreamed_spend.cost > 0 + assert unstreamed_spend.usage + assert streamed_spend == unstreamed_spend + + +class _LoggedResponseRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.responses: Final[queue.Queue[object]] = queue.Queue() + + def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.responses.put(kwargs["standard_logging_object"]["response"]) + + +@pytest.mark.parametrize( + "provider_body", + ({"text": "hello"}, {"text": "hello", "languages": ["en"]}, {"text": "hello", "languages": None}), +) +def test_transcription_without_parsed_languages_has_no_languages_key( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, provider_body: dict[str, object] +) -> None: + recorder: Final = _LoggedResponseRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + with _openai_transcription_client( + httpx.MockTransport(lambda request: httpx.Response(200, json=provider_body)) + ) as client: + response: Final = litellm.transcription(model="openai/whisper-1", file=_one_second_wav(), client=client) + logged_response: Final = recorder.responses.get(timeout=5) + + assert response.text == "hello" + assert "languages" not in json.loads(response.model_dump_json()) + assert isinstance(logged_response, dict) + assert "languages" not in logged_response + + +def test_unstreamed_transcription_sends_no_stream_field(local_model_cost_map: None) -> None: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: collects what the mock upstream received + + def send_response(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.read()) + return httpx.Response(200, json={"text": "hello"}) + + with _openai_transcription_client(httpx.MockTransport(send_response)) as client: + response: Final = litellm.transcription(model="openai/whisper-1", file=_one_second_wav(), client=client) + + assert response.text == "hello" + assert len(sent_bodies) == 1 + assert b'name="stream"' not in sent_bodies[0] + + +def test_transcription_forwards_both_language_and_languages(local_model_cost_map: None) -> None: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: collects what the mock upstream received + + def send_response(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.read()) + return httpx.Response(200, json={"text": "hello", "usage": {"type": "duration", "seconds": 1}}) + + with _openai_transcription_client(httpx.MockTransport(send_response)) as client: + response: Final = litellm.transcription( + model="openai/gpt-transcribe", file=_one_second_wav(), language="en", languages=["fr"], client=client + ) + + assert response.text == "hello" + assert len(sent_bodies) == 1 + assert b'name="language"\r\n\r\nen\r\n' in sent_bodies[0] + assert b'name="languages[]"\r\n\r\nfr\r\n' in sent_bodies[0] + + +def _gpt_transcribe_cost_per_second() -> float: + rate: Final = litellm.model_cost["gpt-transcribe"]["input_cost_per_second"] + assert isinstance(rate, float) + return rate + + +async def _spend_for_one_second_wav( + monkeypatch: pytest.MonkeyPatch, provider_body: dict[str, object], router: litellm.Router | None +) -> _LoggedSpend: + recorder: Final = _SuccessSpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + client: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json=provider_body)) + ), + ) + if router is None: + await litellm.atranscription(model="openai/gpt-transcribe", file=_one_second_wav(), client=client) + else: + await router.atranscription(model="my-transcribe", file=_one_second_wav(), client=client) + spend: Final = await asyncio.wait_for(recorder.async_spends.get(), timeout=5) + await client.close() + return spend + + +@pytest.mark.asyncio +async def test_custom_named_gpt_transcribe_deployment_bills_the_measured_upload_over_reported_seconds( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "my-transcribe", + "litellm_params": {"model": "openai/my-transcribe", "api_key": "sk-test"}, + "model_info": {"base_model": "gpt-transcribe"}, + } + ] + ) + + spend: Final = await _spend_for_one_second_wav( + monkeypatch, {"text": "hello", "usage": {"type": "duration", "seconds": 2}}, router + ) + + assert spend.cost == pytest.approx(_gpt_transcribe_cost_per_second()) + + +@pytest.mark.asyncio +async def test_gpt_transcribe_bills_the_file_length_when_the_provider_reports_zero_seconds( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + spend: Final = await _spend_for_one_second_wav( + monkeypatch, {"text": "hello", "usage": {"type": "duration", "seconds": 0}}, None + ) + + assert spend.cost == pytest.approx(_gpt_transcribe_cost_per_second()) + + +@pytest.mark.asyncio +async def test_m4a_transcription_stream_that_ends_early_logs_one_spend_row_like_the_unstreamed_call( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + async def send_response(request: httpx.Request) -> httpx.Response: + if not _asks_for_stream(await request.aread()): + return httpx.Response(200, json={"text": "hello"}) + return httpx.Response( + 200, + content=f"data: {json.dumps({'type': 'transcript.text.delta', 'delta': 'hello'})}\n\n".encode(), + headers={"content-type": "text/event-stream"}, + ) + + recorder: Final = _SuccessSpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + client: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://example.com/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(send_response)), + ) + await litellm.atranscription(model="openai/gpt-transcribe", file=("sample.m4a", b"m4a audio"), client=client) + unstreamed_spend: Final = await asyncio.wait_for(recorder.async_spends.get(), timeout=5) + stream: Final = await litellm.atranscription( + model="openai/gpt-transcribe", file=("sample.m4a", b"m4a audio"), stream=True, client=client + ) + events: Final = [event.type async for event in stream] + streamed_spend: Final = await asyncio.wait_for(recorder.async_spends.get(), timeout=5) + await asyncio.sleep(0.2) + await client.close() + + assert events == ["transcript.text.delta"] + assert streamed_spend.cost == unstreamed_spend.cost + assert recorder.async_spends.empty() diff --git a/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py index 703fa13cbc9..61798b49be0 100644 --- a/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py +++ b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py @@ -9,6 +9,9 @@ TranscriptionVerbose/Diarized type. from unittest.mock import patch +import pytest + +import litellm from litellm.cost_calculator import completion_cost from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( convert_to_model_response_object, @@ -55,10 +58,7 @@ class TestDiarizedJsonUsageParsing: assert result.usage.seconds == 295.8 def test_usage_duration_object_accepts_float_seconds(self): - assert ( - TranscriptionUsageDurationObject(type="duration", seconds=295.8).seconds - == 295.8 - ) + assert TranscriptionUsageDurationObject(type="duration", seconds=295.8).seconds == 295.8 class TestTranscriptionDurationNotInResponseBody: @@ -126,6 +126,61 @@ class TestTranscriptionDurationNotInResponseBody: class TestCostCalculatorReadsDurationFromHiddenParams: """The cost calculator should read duration from _hidden_params via completion_cost().""" + @patch( # test-quality-ok: isolates duration selection from the model pricing table + "litellm.cost_calculator.openai_cost_per_second" + ) + def test_completion_cost_prefers_measured_duration_over_provider_usage_duration(self, mock_cost_fn): + mock_cost_fn.return_value = (0.001, 0.0) + response = TranscriptionResponse( + text="test", + usage={"type": "duration", "seconds": 60.0}, + ) + response._hidden_params = { + "audio_transcription_duration": 17.5, + "custom_llm_provider": "openai", + } + + completion_cost( + completion_response=response, + model="gpt-transcribe", + call_type="atranscription", + ) + + mock_cost_fn.assert_called_once() + _, kwargs = mock_cost_fn.call_args + assert kwargs["duration"] == 17.5 + + @pytest.mark.parametrize( + ("model", "provider"), + [ + ("gpt-transcribe", "openai"), + ("azure/gpt-transcribe", "azure"), + ], + ) + def test_gpt_transcribe_provider_usage_duration_costs_without_local_decode( + self, + monkeypatch, + model: str, + provider: str, + ): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + monkeypatch.setitem(litellm.model_cost, model, {**litellm.model_cost[model], "input_cost_per_second": 0.0125}) + response = TranscriptionResponse( + text="test", + usage={"type": "duration", "seconds": 60.0}, + ) + response._hidden_params = {"custom_llm_provider": provider} + + actual = completion_cost( + completion_response=response, + model=model, + call_type="atranscription", + custom_llm_provider=provider, + ) + + assert actual == pytest.approx(60.0 * 0.0125) + @patch("litellm.cost_calculator.openai_cost_per_second") def test_completion_cost_uses_hidden_params_duration(self, mock_cost_fn): """ diff --git a/tests/unit/llms/opencode/chat/test_transformation.py b/tests/unit/llms/opencode/chat/test_transformation.py index 0a5e49ca1b2..08c80db294c 100644 --- a/tests/unit/llms/opencode/chat/test_transformation.py +++ b/tests/unit/llms/opencode/chat/test_transformation.py @@ -14,6 +14,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.opencode.chat.transformation import OpenCodeGoMessagesChatConfig from litellm.llms.opencode.common_utils import go_headers @@ -354,6 +355,17 @@ def test_gemini_custom_base_normalizes_version_and_models(go_catalog: None, suff assert response.choices[0].message.content == "hello" +def test_gemini_hosted_mcp_error_names_the_opencode_go_route(go_catalog: None) -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="fixture-gemini on opencode_go cannot run"): + litellm.completion( + model="opencode_go/fixture-gemini", + api_key="fake-go-key", + api_base="https://gateway.example/google", + messages=[{"role": "user", "content": "hello"}], + tools=[{"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}], + ) + + def test_gemini_sync_stream_preserves_tool_calls_and_usage(go_catalog: None) -> None: chunk: Final = { "candidates": [ @@ -580,3 +592,22 @@ def test_stream_chunk_builder_reads_list_cost_from_dict_chunks(): assert assembled is not None assert assembled._hidden_params["response_cost"] == GO_REPORTED_USAGE["opencode"]["list_cost"] + + +@pytest.mark.parametrize(("model", "cache_marks"), [("qwen3.8-flash", 0), ("minimax-m3", 0), ("claude-haiku-5-5", 4)]) +def test_go_messages_route_carries_prompt_cache_breakpoint_to_cache_control_only_for_claude(model: str, cache_marks: int): + marked_notes: Final = [ + {"type": "text", "text": f"Note {n}", "prompt_cache_breakpoint": {"mode": "explicit"}} for n in range(4) + ] + + request: Final = OpenCodeGoMessagesChatConfig().transform_request( + model=model, + messages=[{"role": "user", "content": [*marked_notes, {"type": "text", "text": "Answer in one line."}]}], + optional_params={"max_tokens": 64}, + litellm_params={}, + headers={}, + ) + + wire: Final = json.dumps(request) + assert wire.count('"cache_control"') == cache_marks + assert "prompt_cache_breakpoint" not in wire diff --git a/tests/unit/llms/parallel_ai/decisions/__init__.py b/tests/unit/llms/parallel_ai/decisions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/parallel_ai/decisions/test_transformation.py b/tests/unit/llms/parallel_ai/decisions/test_transformation.py new file mode 100644 index 00000000000..9cd404f61f4 --- /dev/null +++ b/tests/unit/llms/parallel_ai/decisions/test_transformation.py @@ -0,0 +1,423 @@ +import json +from collections.abc import Mapping +from dataclasses import replace +from datetime import datetime +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.base_llm.decisions.transformation import ( + DecisionsResponseValidationError, + ir_to_systemone_response, + systemone_request_to_ir, +) +from litellm.llms.openai.decisions.transformation import ir_to_openai_response, openai_request_to_ir +from litellm.llms.parallel_ai.decisions.transformation import ParallelAIDecisionsConfig +from litellm.types.decisions import ( + DecisionsIRRefusal, + DecisionsRequestBody, + OpenAIDecisionRequestBody, + UnsupportedDecisionsRequest, +) +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager, Rules, function_setup + +QUESTIONS: Final = { + "available": {"type": "noul", "instructions": "Is this product available?"}, + "tier": {"type": "choice", "instructions": "Which tier is offered?", "criteria": {"free": "Free", "paid": "Paid"}}, + "fit": {"type": "score", "instructions": "How well does it fit?", "criteria": ["Poor", "Good", "Excellent"]}, + "unknown": {"type": "noul", "instructions": "Is there enough evidence?"}, +} +REQUEST: Final = systemone_request_to_ir( + DecisionsRequestBody.model_validate({"state": {"product": "Example"}, "questions": QUESTIONS}) +) +PAYLOAD: Final = { + "mode": "fast", + "answers": { + "available": {"type": "noul", "noul": 0.9}, + "tier": {"type": "choice", "choice": "paid", "probabilities": {"free": 0.2, "paid": 0.8}}, + "fit": { + "type": "score", + "score": 1.7, + "legend": {"0": "Poor", "1": "Good", "2": "Excellent"}, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + }, + "unknown": {"type": "refusal"}, + }, + "usage": {"input_tokens": 4212, "output_tokens": 0}, + "citations": [{"url": "https://example.com/product", "title": "Product", "excerpts": ["Available now"]}], + "decision_id": "decisions_example", +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("surface", ["decisions", "systemone"]) +async def test_sdk_sends_parallel_wire_format_and_preserves_evidence( + use_async: bool, surface: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + route: Final = respx_mock.post("https://api.parallel.ai/v1/decisions").respond(json=PAYLOAD) + connection: Final = {"model": "parallel_ai/decision", "api_key": "test-key", "timeout": 5, "num_retries": 0} + fields: Final = ( + {"state": "Example", "questions": QUESTIONS} + if surface == "systemone" + else { + "input": "Example", + "questions": [{"type": "predicate", "name": "available", "instructions": "Available?"}], + } + ) + sync_call: Final = litellm.systemone if surface == "systemone" else litellm.decisions + async_call: Final = litellm.asystemone if surface == "systemone" else litellm.adecisions + response: Final = await async_call(**connection, **fields) if use_async else sync_call(**connection, **fields) + sent: Final = route.calls.last.request + body: Final = json.loads(sent.content) + assert body["mode"] == "fast" + assert "model" not in body + assert body["state"] == "Example" + assert sent.headers["x-api-key"] == "test-key" + assert "authorization" not in sent.headers + assert response.model == "decision" + assert response.model_dump()["citations"] == PAYLOAD["citations"] + assert response.model_dump()["decision_id"] == PAYLOAD["decision_id"] + await GLOBAL_LOGGING_WORKER.clear_queue() + + +def test_provider_routes_systemone_questions_to_parallel_without_a_model_field() -> None: + config: Final = ProviderConfigManager.get_provider_decisions_config("decision", LlmProviders.PARALLEL_AI) + assert config is not None + assert config.transform_decisions_request("decision", REQUEST, "parallel_ai") == { + "mode": "fast", + "state": {"product": "Example"}, + "questions": QUESTIONS, + } + assert config.get_complete_url("https://gateway.example/v1/", "decision") == "https://gateway.example/v1/decisions" + + +@pytest.mark.parametrize("suffix", ["", "/", "/v1", "/v1/", "/v1/decisions", "/v1/decisions/", "/v1/search", "/v1/search/"]) +def test_base_carrying_a_parallel_route_reaches_the_decisions_endpoint(suffix: str) -> None: + complete_url: Final = ParallelAIDecisionsConfig().get_complete_url(f"https://api.parallel.ai{suffix}", "decision") + assert complete_url == "https://api.parallel.ai/v1/decisions" + + +def test_auth_replaces_conflicting_headers_without_forwarding_bearer_credentials() -> None: + headers: Final = ParallelAIDecisionsConfig().validate_environment( + {"Authorization": "Bearer old", "X-API-Key": "old", "CONTENT-TYPE": "text/plain", "X-Trace": "test"}, + "decision", + "test-key", + {}, + ) + assert headers == {"x-api-key": "test-key", "Content-Type": "application/json", "X-Trace": "test"} + + +def test_openai_predicate_is_translated_to_a_named_noul_question() -> None: + request: Final = openai_request_to_ir( + OpenAIDecisionRequestBody.model_validate( + { + "input": "Example product", + "questions": [{"type": "predicate", "name": "available", "instructions": "Available?"}], + } + ) + ) + assert ParallelAIDecisionsConfig().transform_decisions_request("decision", request, "parallel_ai") == { + "mode": "fast", + "state": "Example product", + "questions": {"available": {"type": "noul", "instructions": "Available?"}}, + } + + +@pytest.mark.parametrize("surface", ["systemone", "decisions"]) +def test_all_answers_usage_and_evidence_survive_response_translation(surface: str) -> None: + response: Final = ParallelAIDecisionsConfig().parse_response(PAYLOAD, REQUEST) + assert isinstance(response.answers[3], DecisionsIRRefusal) + formatted: Final = ( + ir_to_systemone_response(response, REQUEST) + if surface == "systemone" + else ir_to_openai_response( + response, REQUEST, "parallel_ai/decision", ParallelAIDecisionsConfig.openai_evidence_fields + ) + ) + result: Final = formatted.model_dump(mode="json") + assert result["citations"] == PAYLOAD["citations"] + assert result["decision_id"] == PAYLOAD["decision_id"] + assert result["mode"] == "fast" + assert result["usage"]["input_tokens"] == 4212 + assert result["usage"]["output_tokens"] == 0 + if surface == "systemone": + assert result["answers"]["available"]["noul"] == 0.9 + assert result["answers"]["tier"]["choice"] == "paid" + assert result["answers"]["fit"]["score"] == 1.7 + assert "unknown" not in result["answers"] + else: + assert result["answers"][3]["type"] == "refusal" + + +def test_rejects_more_than_one_hundred_questions() -> None: + request: Final = systemone_request_to_ir( + DecisionsRequestBody.model_validate( + { + "state": "Example", + "questions": {str(index): QUESTIONS["available"] for index in range(101)}, + } + ) + ) + assert isinstance( + ParallelAIDecisionsConfig().transform_decisions_request("decision", request, "parallel_ai"), + UnsupportedDecisionsRequest, + ) + + +@pytest.mark.parametrize("model", ["fast", "unknown"]) +def test_rejects_unregistered_models_instead_of_forwarding_them_as_modes(model: str) -> None: + assert isinstance( + ParallelAIDecisionsConfig().transform_decisions_request(model, REQUEST, "parallel_ai"), + UnsupportedDecisionsRequest, + ) + + +@pytest.mark.parametrize("confidence", [None, 0.8]) +@pytest.mark.parametrize("surface", ["systemone", "decisions"]) +def test_optional_confidence_is_preserved_without_inventing_a_value(confidence: float | None, surface: str) -> None: + payload: Final = { + **PAYLOAD, + "answers": {name: {**answer, "confidence": confidence} for name, answer in PAYLOAD["answers"].items()}, + } + response: Final = ParallelAIDecisionsConfig().parse_response(payload, REQUEST) + formatted: Final = ( + ir_to_systemone_response(response, REQUEST) + if surface == "systemone" + else ir_to_openai_response( + response, REQUEST, "parallel_ai/decision", ParallelAIDecisionsConfig.openai_evidence_fields + ) + ) + answers: Final = formatted.model_dump(mode="json")["answers"] + values: Final = answers.values() if isinstance(answers, dict) else answers + assert [answer["confidence"] for answer in values if answer["type"] in {"choice", "score"}] == [confidence] * 2 + + +def test_provider_extras_cannot_replace_normalized_response_fields() -> None: + response: Final = ParallelAIDecisionsConfig().parse_response(PAYLOAD, REQUEST) + with_collision: Final = replace(response, extra={**response.extra, "model": "wrong", "answers": [], "usage": {}}) + formatted: Final = ir_to_openai_response( + with_collision, REQUEST, "parallel_ai/decision", ParallelAIDecisionsConfig.openai_evidence_fields + ) + assert formatted.model == "decision" + assert len(formatted.answers) == 4 + assert formatted.usage.input_tokens == 4212 + assert formatted.model_dump()["citations"] == PAYLOAD["citations"] + + +def test_undocumented_parallel_fields_stay_off_the_openai_response() -> None: + response: Final = ParallelAIDecisionsConfig().parse_response({**PAYLOAD, "internal_trace_id": "trace"}, REQUEST) + formatted: Final = ir_to_openai_response( + response, REQUEST, "parallel_ai/decision", ParallelAIDecisionsConfig.openai_evidence_fields + ) + assert "internal_trace_id" not in formatted.model_dump() + + +@pytest.mark.parametrize("tokens", [0, 100, 32000]) +@pytest.mark.parametrize("call_type", ["decisions", "adecisions", "systemone", "asystemone"]) +def test_request_pricing_is_independent_of_token_count(tokens: int, call_type: str) -> None: + rate: Final = litellm.model_cost["parallel_ai/decision"]["input_cost_per_query"] + cost: Final = litellm.cost_per_token( + model="parallel_ai/decision", + prompt_tokens=tokens, + completion_tokens=0, + call_type=call_type, + ) + assert cost == (rate, 0.0) + + +@pytest.mark.asyncio +async def test_deployment_query_price_replaces_the_catalog_price( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://api.parallel.ai/v1/decisions").respond(json=PAYLOAD) + price: Final = 25 * litellm.model_cost["parallel_ai/decision"]["input_cost_per_query"] + router: Final = litellm.Router( + model_list=[ + { + "model_name": "webjev", + "litellm_params": { + "model": "parallel_ai/decision", + "api_key": "test-key", + "input_cost_per_query": price, + }, + } + ] + ) + fields: Final = { + "model": "webjev", + "input": "Example", + "questions": [{"type": "predicate", "name": "available", "instructions": "Available?"}], + "litellm_call_id": "query-price-call", + } + logging_obj, data = function_setup("adecisions", Rules(), datetime.now(), **fields) + response: Final = await router.adecisions(**data, litellm_logging_obj=logging_obj) + assert response._hidden_params["response_cost"] == pytest.approx(price) + await GLOBAL_LOGGING_WORKER.clear_queue() + + +@pytest.mark.asyncio +async def test_query_price_on_another_provider_leaves_its_token_billing_alone( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + input_rate: Final = 1e-06 + tokens: Final = 1000 + model: Final = "fireworks_ai/accounts/example/models/token-billed" + monkeypatch.setitem( + litellm.model_cost, + model, + {"input_cost_per_token": input_rate, "output_cost_per_token": 0.0} + | {"litellm_provider": "fireworks_ai", "mode": "evaluation"}, + ) + respx_mock.post("https://api.fireworks.ai/inference/v1/systemone").respond( + json={ + "model": "token-billed", + "answers": {"available": {"type": "noul", "noul": 0.9}}, + "usage": {"input_tokens": tokens, "output_tokens": 0}, + } + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "token-billed", + "litellm_params": {"model": model, "api_key": "test-key", "input_cost_per_query": 0.025}, + } + ] + ) + fields: Final = { + "model": "token-billed", + "input": "Example", + "questions": [{"type": "predicate", "name": "available", "instructions": "Available?"}], + "litellm_call_id": "token-billed-call", + } + logging_obj, data = function_setup("adecisions", Rules(), datetime.now(), **fields) + response: Final = await router.adecisions(**data, litellm_logging_obj=logging_obj) + assert response._hidden_params["response_cost"] == pytest.approx(tokens * input_rate) + await GLOBAL_LOGGING_WORKER.clear_queue() + + +@pytest.mark.parametrize( + "payload", [{"answers": {}}, {"mode": "fast", "answers": {"available": {"type": "noul", "noul": "invalid"}}}] +) +def test_malformed_provider_responses_raise_a_provider_error(payload: Mapping[str, object]) -> None: + with pytest.raises(DecisionsResponseValidationError): + ParallelAIDecisionsConfig().transform_decisions_response( + "decision", "parallel_ai", httpx.Response(200, json=payload), REQUEST + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("surface", ["decisions", "systemone"]) +@pytest.mark.parametrize("key_env", ["PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"]) +async def test_untrusted_base_cannot_receive_environment_key( + use_async: bool, surface: str, key_env: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("PARALLEL_AI_API_BASE", raising=False) + monkeypatch.delenv("PARALLEL_AI_API_KEY", raising=False) + monkeypatch.delenv("PARALLEL_API_KEY", raising=False) + monkeypatch.setenv(key_env, "server-secret") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://attacker.example/v1/decisions").respond(json=PAYLOAD) + connection: Final = {"model": "parallel_ai/decision", "api_base": "https://attacker.example", "num_retries": 0} + fields: Final = ( + {"state": "Example", "questions": QUESTIONS} + if surface == "systemone" + else { + "input": "Example", + "questions": [{"type": "predicate", "name": "available", "instructions": "Available?"}], + } + ) + sync_call: Final = litellm.systemone if surface == "systemone" else litellm.decisions + async_call: Final = litellm.asystemone if surface == "systemone" else litellm.adecisions + + async def call() -> object: + return await async_call(**connection, **fields) if use_async else sync_call(**connection, **fields) + + with pytest.raises(litellm.BadRequestError, match="explicit api_key"): + await call() + assert not route.called + await GLOBAL_LOGGING_WORKER.clear_queue() + + +@pytest.mark.parametrize( + "api_base", + [ + "https://api.parallel.ai.attacker.example", + "https://api.parallel.ai@attacker.example", + "http://api.parallel.ai", + "https://api.parallel.ai:8443", + "https://attacker.example/api.parallel.ai", + "https://api.parallel.ai/other-tenant", + "https://api.parallel.ai/other-tenant/v1/search", + "https://attacker.example/v1/search", + "https://attacker.example/v1/decisions", + ], +) +def test_environment_key_rejects_untrusted_destination_variants(api_base: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PARALLEL_AI_API_BASE", raising=False) + monkeypatch.setenv("PARALLEL_AI_API_KEY", "server-secret") + with pytest.raises(ValueError, match="explicit api_key"): + ParallelAIDecisionsConfig().resolve_connection(api_base, None) + + +@pytest.mark.parametrize( + "api_base", + [ + None, + "https://api.parallel.ai", + "https://API.PARALLEL.AI:443/v1/", + "https://api.parallel.ai/v1/decisions", + "https://api.parallel.ai/v1/search", + ], +) +def test_default_destination_retains_environment_key(api_base: str | None, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PARALLEL_AI_API_BASE", raising=False) + monkeypatch.setenv("PARALLEL_AI_API_KEY", "server-secret") + resolved_base, key = ParallelAIDecisionsConfig().resolve_connection(api_base, None) + assert resolved_base == (api_base or "https://api.parallel.ai") + assert key == "server-secret" + + +@pytest.mark.parametrize( + "api_base", [None, "https://gateway.example/parallel/v1/", "https://gateway.example/parallel/v1/search"] +) +def test_operator_base_retains_environment_key(api_base: str | None, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PARALLEL_AI_API_BASE", "https://gateway.example/parallel") + monkeypatch.setenv("PARALLEL_AI_API_KEY", "server-secret") + resolved_base, key = ParallelAIDecisionsConfig().resolve_connection(api_base, None) + assert resolved_base == (api_base or "https://gateway.example/parallel") + assert key == "server-secret" + + +@pytest.mark.parametrize("api_key", [None, "caller-key"]) +@pytest.mark.parametrize("api_base", ["https://api.parallel.ai\t", "https://api.parallel.ai\r\nHost: attacker.example"]) +def test_malformed_destination_is_a_bad_request(api_base: str, api_key: str | None) -> None: + with pytest.raises(litellm.BadRequestError): + litellm.decisions( + model="parallel_ai/decision", + api_base=api_base, + api_key=api_key, + input="Example", + questions=[{"type": "predicate", "name": "available", "instructions": "Available?"}], + num_retries=0, + ) + + +def test_custom_destination_uses_explicit_key_instead_of_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PARALLEL_AI_API_KEY", "server-secret") + assert ParallelAIDecisionsConfig().resolve_connection("https://gateway.example", "caller-key") == ( + "https://gateway.example", + "caller-key", + ) diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 67733884e12..233e5d888fa 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,6 +1,5 @@ import asyncio import base64 -import logging import uuid from collections.abc import Generator from contextlib import contextmanager @@ -11,20 +10,22 @@ from typing import Final, Literal from unittest.mock import patch import pytest -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue import litellm from litellm.llms.anthropic.messages.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, truncate_tool_name, ) -from litellm.llms.vertex_ai.count_tokens.context import ( +from litellm.llms.anthropic.messages.context_management.inference_view import ( CountContextEdits, InferenceRequestView, + inference_request_view, +) +from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContext, VertexGeminiCountContextError, gemini_count_contents, - inference_request_view, unconvertible_messages_error, vertex_gemini_count_context, ) @@ -424,73 +425,6 @@ async def test_image_that_cannot_be_fetched_returns_an_error_value( assert isinstance(contents, VertexGeminiCountContextError) -@pytest.mark.parametrize( - "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") -) -async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: - view: Final = await inference_request_view( - ( - {"role": "user", "content": "first question " * 50}, - {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, - {"role": "user", "content": "latest question"}, - ), - system, - "anthropic", - _NO_EDITS, - ) - - assert isinstance(view, InferenceRequestView) - assert view.messages == [{"role": "user", "content": "latest question"}] - assert "Summary of the earlier turns." in str(view.system) - assert "Be brief." in str(view.system) - - -async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: - snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." - view: Final = await inference_request_view( - ( - {"role": "user", "content": "What is new?"}, - { - "role": "assistant", - "content": [ - {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, - { - "type": "web_search_tool_result", - "tool_use_id": "srvtoolu_01", - "content": [ - { - "type": "web_search_result", - "url": "https://docs.litellm.ai/release_notes", - "title": "Release notes", - "encrypted_content": "", - "snippet": snippet, - } - ], - }, - ], - }, - ), - None, - "anthropic", - _NO_EDITS, - ) - - assert isinstance(view, InferenceRequestView) - assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) - assert [block["type"] for block in assistant_blocks] == ["text"] - assert snippet in assistant_blocks[0]["text"] - - -async def test_openai_view_is_the_request_unchanged() -> None: - messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) - - view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS) - - assert isinstance(view, InferenceRequestView) - assert view.messages == messages - assert view.system == "Be brief." - - async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: with patch( "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", @@ -556,138 +490,30 @@ async def test_anthropic_replayed_web_search_result_is_counted_as_its_text( assert snippet in str(contents) -_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 -_READ_FILE_TOOL: Final[dict[str, object]] = { - "name": "read_file", - "description": "Read a file", - "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, -} -_CLEAR_EDIT: Final[JsonValue] = { - "type": "clear_tool_uses_20250919", - "trigger": {"type": "input_tokens", "value": 500}, - "keep": {"type": "tool_uses", "value": 1}, -} -_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]} - - -def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]: - return ( - *( - turn - for index in range(first_round, first_round + rounds) - for turn in ( - {"role": "user", "content": f"Read file {index}."}, - { - "role": "assistant", - "content": [ - {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} - ], - }, - { - "role": "user", - "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}], - }, - ) - ), # comprehension-ok: flattens each round into its three turns - {"role": "user", "content": "Summarize what you read."}, - ) - - -def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]: - return tuple( - _TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"]) - ) - - -def _edits(spec: JsonValue) -> CountContextEdits: - return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,)) - - -async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None: - view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) - - assert isinstance(view, InferenceRequestView) - assert _tool_outputs_kept(view) == (False, False, True) - - -async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None: - spec: Final[JsonValue] = { - "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT] - } - turns: Final = ( - {"role": "user", "content": "first question " * 30000}, - {"role": "assistant", "content": "first answer"}, - *_agent_turns(3), - ) - - with ( - patch( - "litellm.llms.anthropic.messages.context_management.editors.compact._read_summary_model_setting", - return_value="gemini/gemini-2.5-flash", - ), - patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary, - ): - view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec)) - - assert summary.call_count == 0 - assert isinstance(view, InferenceRequestView) - assert _tool_outputs_kept(view) == (False, False, True) - - -async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails( - caplog: pytest.LogCaptureFixture, -) -> None: - with ( - patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")), - caplog.at_level(logging.DEBUG, logger="LiteLLM"), - ): - view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) - - assert isinstance(view, InferenceRequestView) - assert _tool_outputs_kept(view) == (True, True, True) - assert any("skipping edits" in record.getMessage() for record in caplog.records) - assert not any(record.levelno >= logging.ERROR for record in caplog.records) - - -async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None: - clear_before_compact: Final[JsonValue] = { - "edits": [ - { - "type": "clear_tool_uses_20250919", - "trigger": {"type": "tool_uses", "value": 3}, - "keep": {"type": "tool_uses", "value": 1}, - }, - {"type": "compact_20260112"}, - ] - } - turns: Final = ( - *_agent_turns(3)[:-1], - {"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]}, - *_agent_turns(2, first_round=3), - ) - - with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")): - view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact)) - - assert isinstance(view, InferenceRequestView) - assert _tool_outputs_kept(view) == (True, True) - assert "Read files 0 to 2." in str(view.system) - - -async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None: - low_trigger: Final[JsonValue] = { - "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] - } - - view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger)) - - assert isinstance(view, VertexGeminiCountContextError) - assert view.status_code == 400 - - async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_trigger_cannot_tokenize_it() -> None: + clear_old_tool_results: Final = CountContextEdits( + spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + model="gemini/gemini-2.5-flash", + tools=({"name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}},), + ) turns: Final = ( - *_agent_turns(1)[:-1], + {"role": "user", "content": "Read file 0."}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_0", "name": "read_file", "input": {"path": "0"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_0", "content": "line of a long file. " * 180}], + }, { "role": "user", "content": [ @@ -701,7 +527,7 @@ async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_t }, ) - view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + view: Final = await inference_request_view(turns, None, "anthropic", clear_old_tool_results) assert isinstance(view, InferenceRequestView) blocks_error: Final = unconvertible_messages_error(view.messages, "anthropic") diff --git a/tests/unit/llms/vertex_ai/gemini/test_rejected_thought_signature_retry.py b/tests/unit/llms/vertex_ai/gemini/test_rejected_thought_signature_retry.py new file mode 100644 index 00000000000..28674ec2621 --- /dev/null +++ b/tests/unit/llms/vertex_ai/gemini/test_rejected_thought_signature_retry.py @@ -0,0 +1,496 @@ +import base64 +import json +from collections.abc import Callable, Mapping, Sequence +from itertools import chain +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.google_genai import agenerate_content, agenerate_content_stream, generate_content, generate_content_stream +from litellm.litellm_core_utils.prompt_templates.factory import get_dummy_thought_signature +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.gemini.rejected_thought_signature_retry import with_placeholder_thought_signatures +from litellm.types.llms.vertex_ai import RequestBody +from litellm.types.utils import Choices, ModelResponse, ModelResponseStream +from litellm.utils import CustomStreamWrapper + +MODEL: Final = "gemini/gemini-3.8-flash" +VERTEX_MINTED: Final = base64.b64encode(b"minted-by-vertex").decode() +VERTEX_MINTED_TEXT: Final = base64.b64encode(b"minted-by-vertex-for-text").decode() +AI_STUDIO_MINTED: Final = base64.b64encode(b"minted-by-ai-studio").decode() +HISTORY_WITH_FOREIGN_SIGNATURES: Final = ( + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"thought_signature": VERTEX_MINTED}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "18C and sunny"}, + { + "role": "assistant", + "content": "It is 18C and sunny in Paris.", + "provider_specific_fields": {"thought_signatures": [VERTEX_MINTED_TEXT]}, + }, + {"role": "user", "content": "And in Berlin?"}, +) +HISTORY_WITHOUT_SIGNATURES: Final = ( + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "18C and sunny"}, + {"role": "user", "content": "And in Berlin?"}, +) +TOOLS: Final = ( + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + }, +) +ANSWER: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Berlin is 12C."}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 40, "candidatesTokenCount": 6, "totalTokenCount": 46}, +} +# Wording observed from Google on 2026-10-08 when a signature minted by the other backend is replayed +AI_STUDIO_REJECTION: Final = "Corrupted thought signature." +VERTEX_REJECTION: Final = "Invalid thought signature." +UNRELATED_REJECTION: Final = 'Invalid JSON payload received. Unknown name "foo".' + + +def _rejected(message: str) -> httpx.Response: + return httpx.Response(400, json={"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}}) + + +def _sse(*chunks: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + content="".join(f"data: {json.dumps(chunk)}\r\n\r\n" for chunk in chunks).encode(), + headers={"content-type": "text/event-stream"}, + ) + + +def _replay( + answers: Sequence[httpx.Response], sent: list[dict[str, object]] +) -> Callable[[httpx.Request], httpx.Response]: + def handler(request: httpx.Request) -> httpx.Response: + sent.append(json.loads(request.content)) + return answers[len(sent) - 1] + + return handler + + +def _signatures_of(body: Mapping[str, object], key: str = "thoughtSignature") -> list[str]: + contents: Final = body["contents"] + assert isinstance(contents, list) + parts: Final = chain.from_iterable(content["parts"] for content in contents) + return [part[key] for part in parts if key in part] + + +def _text_of(response: ModelResponse) -> str: + choice: Final = response.choices[0] + assert isinstance(choice, Choices), choice + assert isinstance(choice.message.content, str), choice.message + return choice.message.content + + +def _delta_text_of(chunk: object) -> str: + assert isinstance(chunk, ModelResponseStream), chunk + return chunk.choices[0].delta.content or "" + + +async def _complete( + messages: Sequence[Mapping[str, object]], + backend: Callable[[httpx.Request], httpx.Response], + is_async: bool, + stream: bool, +) -> str: + kwargs: Final = { + "model": MODEL, + "messages": list(messages), + "tools": list(TOOLS), + "stream": stream, + "api_key": "unit-test-key", + } + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(backend))) + if is_async + else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(backend)))) + ) + if isinstance(response, ModelResponse): + return _text_of(response) + assert isinstance(response, CustomStreamWrapper), response + if is_async: + return "".join([_delta_text_of(chunk) async for chunk in response]) + return "".join(_delta_text_of(chunk) for chunk in response) + + +def _answer(stream: bool) -> httpx.Response: + return _sse(ANSWER) if stream else httpx.Response(200, json=ANSWER) + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("rejection", [AI_STUDIO_REJECTION, VERTEX_REJECTION]) +@pytest.mark.asyncio +async def test_signature_rejected_by_the_other_backend_is_resent_once_with_the_placeholder( + is_async: bool, stream: bool, rejection: str +): + sent: Final[list[dict[str, object]]] = [] + answer: Final = await _complete( + HISTORY_WITH_FOREIGN_SIGNATURES, _replay((_rejected(rejection), _answer(stream)), sent), is_async, stream + ) + assert answer == "Berlin is 12C." + assert len(sent) == 2 + first_signatures, second_signatures = _signatures_of(sent[0]), _signatures_of(sent[1]) + assert set(first_signatures) == {VERTEX_MINTED, VERTEX_MINTED_TEXT} + assert second_signatures == [get_dummy_thought_signature()] * len(first_signatures) + assert {**sent[1], "contents": None} == {**sent[0], "contents": None} + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_placeholder_rejected_too_is_raised_after_the_second_post(is_async: bool, stream: bool): + sent: Final[list[dict[str, object]]] = [] + with pytest.raises(litellm.BadRequestError, match="thought signature"): + await _complete( + HISTORY_WITH_FOREIGN_SIGNATURES, + _replay((_rejected(AI_STUDIO_REJECTION), _rejected(AI_STUDIO_REJECTION)), sent), + is_async, + stream, + ) + assert len(sent) == 2 + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_unrelated_400_is_raised_without_a_resend(is_async: bool, stream: bool): + sent: Final[list[dict[str, object]]] = [] + with pytest.raises(litellm.BadRequestError, match="Unknown name"): + await _complete(HISTORY_WITH_FOREIGN_SIGNATURES, _replay((_rejected(UNRELATED_REJECTION),), sent), is_async, stream) + assert len(sent) == 1 + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_signature_400_on_a_history_carrying_only_the_placeholder_is_raised_without_a_resend( + is_async: bool, stream: bool +): + sent: Final[list[dict[str, object]]] = [] + with pytest.raises(litellm.BadRequestError, match="thought signature"): + await _complete(HISTORY_WITHOUT_SIGNATURES, _replay((_rejected(AI_STUDIO_REJECTION),), sent), is_async, stream) + assert len(sent) == 1 + assert set(_signatures_of(sent[0])) <= {get_dummy_thought_signature()} + + +def test_placeholder_body_keeps_the_original_untouched(): + body: Final[RequestBody] = { + "contents": [ + {"role": "model", "parts": [{"text": "hi", "thoughtSignature": VERTEX_MINTED}, {"text": "unsigned"}]}, + {"role": "user", "parts": [{"text": "again"}]}, + ], + "generationConfig": {"temperature": 0.1}, + } + resend: Final = with_placeholder_thought_signatures(body) + assert resend is not None + assert _signatures_of(resend) == [get_dummy_thought_signature()] + assert _signatures_of(body) == [VERTEX_MINTED] + assert resend.get("generationConfig") == body.get("generationConfig") + + +def test_body_without_a_foreign_signature_has_nothing_to_resend(): + body: Final[RequestBody] = { + "contents": [ + {"role": "model", "parts": [{"text": "hi", "thoughtSignature": get_dummy_thought_signature()}]}, + {"role": "user", "parts": [{"text": "again"}]}, + ] + } + assert with_placeholder_thought_signatures(body) is None + + +CACHE_CONTROL: Final = {"type": "ephemeral"} + + +def _cached_history( + tag: str, cached_signature: str = VERTEX_MINTED, later_signatures: tuple[str, ...] = () +) -> tuple[Mapping[str, object], ...]: + later_answer: Final = ( + {"role": "assistant", "content": "Noted.", "provider_specific_fields": {"thought_signatures": [*later_signatures]}} + if later_signatures + else {"role": "assistant", "content": "Noted."} + ) + prefix: Final = " ".join(f"Reference paragraph {i} for {tag}: weather values are in Celsius." for i in range(600)) + return ( + {"role": "system", "content": [{"type": "text", "text": prefix, "cache_control": CACHE_CONTROL}]}, + { + "role": "user", + "content": [{"type": "text", "text": "What is the weather in Paris?", "cache_control": CACHE_CONTROL}], + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "Calling the tool.", "cache_control": CACHE_CONTROL}], + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"thought_signature": cached_signature}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "text", "text": "18C and sunny", "cache_control": CACHE_CONTROL}], + }, + {"role": "user", "content": [{"type": "text", "text": "Got it.", "cache_control": CACHE_CONTROL}]}, + later_answer, + {"role": "user", "content": "And in Berlin?"}, + ) + + +def _google_ai_studio( + stream: bool, generated: list[dict[str, object]], minted: frozenset[str] = frozenset() +) -> Callable[[httpx.Request], httpx.Response]: + """Rejects a generateContent call whose own contents or referenced cache carry a signature AI Studio did not mint, + the way the live backend answered on 2026-11-02.""" + caches: Final[dict[str, object]] = {} + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/cachedContents"): + if request.method == "GET": + return httpx.Response(200, json={}) + caches["cachedContents/prefix"] = json.loads(request.content) + return httpx.Response( + 200, json={"name": "cachedContents/prefix", "usageMetadata": {"totalTokenCount": 9000}} + ) + body: Final = json.loads(request.content) + generated.append(body) + cache: Final = caches.get(body.get("cachedContent"), {"contents": []}) + assert isinstance(cache, dict) + signatures: Final = {*_signatures_of(body), *_signatures_of(cache)} + if signatures - {get_dummy_thought_signature(), *minted}: + return _rejected(AI_STUDIO_REJECTION) + return _answer(stream) + + return handler + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_signature_inside_a_cached_prefix_is_resent_inline_with_the_placeholder(is_async: bool, stream: bool): + generated: Final[list[dict[str, object]]] = [] + answer: Final = await _complete( + _cached_history(f"async={is_async} stream={stream}"), + _google_ai_studio(stream, generated), + is_async, + stream, + ) + assert answer == "Berlin is 12C." + assert len(generated) == 2 + assert generated[0]["cachedContent"] == "cachedContents/prefix" + assert _signatures_of(generated[0]) == [] + assert "cachedContent" not in generated[1] + assert _signatures_of(generated[1]) == [get_dummy_thought_signature()] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_signature_after_a_cache_the_backend_accepts_is_resent_with_the_cache_kept(is_async: bool, stream: bool): + generated: Final[list[dict[str, object]]] = [] + answer: Final = await _complete( + _cached_history( + f"own cache async={is_async} stream={stream}", AI_STUDIO_MINTED, later_signatures=(VERTEX_MINTED_TEXT,) + ), + _google_ai_studio(stream, generated, minted=frozenset({AI_STUDIO_MINTED})), + is_async, + stream, + ) + assert answer == "Berlin is 12C." + assert [body.get("cachedContent") for body in generated] == ["cachedContents/prefix"] * 2 + assert _signatures_of(generated[0]) == [VERTEX_MINTED_TEXT] + assert _signatures_of(generated[1]) == [get_dummy_thought_signature()] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_signature_inside_and_after_the_cached_prefix_goes_inline_once_the_cache_kept_resend_is_rejected( + is_async: bool, stream: bool +): + generated: Final[list[dict[str, object]]] = [] + answer: Final = await _complete( + _cached_history(f"both async={is_async} stream={stream}", later_signatures=(VERTEX_MINTED_TEXT,)), + _google_ai_studio(stream, generated), + is_async, + stream, + ) + assert answer == "Berlin is 12C." + assert [body.get("cachedContent") for body in generated] == ["cachedContents/prefix"] * 2 + [None] + assert _signatures_of(generated[1]) == [get_dummy_thought_signature()] + assert _signatures_of(generated[2]) == [get_dummy_thought_signature()] * 2 + + +def _backend_validating_cached_signatures( + stream: bool, created: list[dict[str, object]], generated: list[dict[str, object]] +) -> Callable[[httpx.Request], httpx.Response]: + """Rejects creating a cache whose contents carry a signature it did not mint, the way Vertex answered on + 2026-11-02.""" + + def handler(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) if request.content else {} + if request.url.path.endswith("/cachedContents"): + if request.method == "GET": + return httpx.Response(200, json={}) + created.append(body) + else: + generated.append(body) + if set(_signatures_of(body)) - {get_dummy_thought_signature()}: + return _rejected(VERTEX_REJECTION) + if request.url.path.endswith("/cachedContents"): + return httpx.Response( + 200, json={"name": "cachedContents/prefix", "usageMetadata": {"totalTokenCount": 9000}} + ) + return _answer(stream) + + return handler + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_signature_rejected_while_creating_the_cache_is_sent_inline_with_the_placeholder( + is_async: bool, stream: bool +): + created: Final[list[dict[str, object]]] = [] + generated: Final[list[dict[str, object]]] = [] + answer: Final = await _complete( + _cached_history(f"cache creation async={is_async} stream={stream}"), + _backend_validating_cached_signatures(stream, created, generated), + is_async, + stream, + ) + assert answer == "Berlin is 12C." + assert [_signatures_of(cache) for cache in created] == [[VERTEX_MINTED]] + assert len(generated) == 1 + assert "cachedContent" not in generated[0] + assert _signatures_of(generated[0]) == [get_dummy_thought_signature()] + + +NATIVE_QUESTION: Final = {"role": "user", "parts": [{"text": "What is the weather in Paris?"}]} +NATIVE_TOOL_RESULT: Final = { + "role": "user", + "parts": [{"functionResponse": {"name": "get_weather", "response": {"result": "18C and sunny"}}}], +} + + +def _native_contents(signature_key: str) -> list[Mapping[str, object]]: + return [ + NATIVE_QUESTION, + { + "role": "model", + "parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, signature_key: VERTEX_MINTED}], + }, + NATIVE_TOOL_RESULT, + ] + + +async def _generate( + contents: list[Mapping[str, object]], + backend: Callable[[httpx.Request], httpx.Response], + is_async: bool, + stream: bool, +) -> str: + kwargs: Final = {"model": MODEL, "contents": contents, "api_key": "unit-test-key"} + sync_client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(backend))) + async_client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(backend)) + match is_async, stream: + case True, True: + chunks = await agenerate_content_stream(**kwargs, client=async_client) + return b"".join([chunk async for chunk in chunks]).decode() + case True, False: + return json.dumps((await agenerate_content(**kwargs, client=async_client)).model_dump(mode="json")) + case False, True: + return b"".join(generate_content_stream(**kwargs, client=sync_client)).decode() + case False, False: + return json.dumps(generate_content(**kwargs, client=sync_client).model_dump(mode="json")) + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("signature_key", ["thoughtSignature", "thought_signature"]) +@pytest.mark.asyncio +async def test_native_generate_content_resends_a_rejected_signature_with_the_placeholder( + is_async: bool, stream: bool, signature_key: str +): + sent: Final[list[dict[str, object]]] = [] + output: Final = await _generate( + _native_contents(signature_key), + _replay((_rejected(AI_STUDIO_REJECTION), _answer(stream)), sent), + is_async, + stream, + ) + assert "Berlin is 12C." in output + assert len(sent) == 2 + assert _signatures_of(sent[0], signature_key) == [VERTEX_MINTED] + assert _signatures_of(sent[1], signature_key) == [get_dummy_thought_signature()] + assert {**sent[1], "contents": None} == {**sent[0], "contents": None} + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_native_generate_content_raises_an_unrelated_400_without_a_resend(is_async: bool, stream: bool): + sent: Final[list[dict[str, object]]] = [] + with pytest.raises(litellm.BadRequestError, match="Unknown name"): + await _generate( + _native_contents("thoughtSignature"), _replay((_rejected(UNRELATED_REJECTION),), sent), is_async, stream + ) + assert len(sent) == 1 + + +def _history_signed_after_the_cached_prefix(tag: str) -> tuple[Mapping[str, object], ...]: + prefix: Final = " ".join(f"Reference paragraph {i} for {tag}: weather values are in Celsius." for i in range(600)) + return ( + {"role": "system", "content": [{"type": "text", "text": prefix, "cache_control": CACHE_CONTROL}]}, + {"role": "user", "content": [{"type": "text", "text": "Hello.", "cache_control": CACHE_CONTROL}]}, + *HISTORY_WITH_FOREIGN_SIGNATURES, + ) + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_signature_after_the_cached_prefix_is_resent_with_the_cache_kept(is_async: bool, stream: bool): + generated: Final[list[dict[str, object]]] = [] + answer: Final = await _complete( + _history_signed_after_the_cached_prefix(f"async={is_async} stream={stream}"), + _google_ai_studio(stream, generated), + is_async, + stream, + ) + assert answer == "Berlin is 12C." + assert len(generated) == 2 + assert generated[1]["cachedContent"] == generated[0]["cachedContent"] == "cachedContents/prefix" + assert set(_signatures_of(generated[0])) == {VERTEX_MINTED, VERTEX_MINTED_TEXT} + assert _signatures_of(generated[1]) == [get_dummy_thought_signature()] * 2 diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index c0ff931a41d..2c8c7520c38 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -2824,7 +2824,7 @@ def test_gemini_25_text_signature_survives_replay_to_gemini_3(): """gemini-2.5 history (signed text, unsigned tool call) replayed to gemini-3 keeps the real text signature; the dummy signature synthesized for the unsigned tool call must not suppress it.""" from litellm.litellm_core_utils.prompt_templates.factory import ( - _get_dummy_thought_signature, + get_dummy_thought_signature, ) from litellm.llms.vertex_ai.gemini.transformation import ( gemini_convert_messages_with_history, @@ -2848,7 +2848,7 @@ def test_gemini_25_text_signature_survives_replay_to_gemini_3(): assert parts[0]["text"] == "I will list the directory." assert parts[0]["thoughtSignature"] == "real_25_signature" assert "function_call" in parts[1] - assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() + assert parts[1]["thoughtSignature"] == get_dummy_thought_signature() def test_gemini_function_call_signature_round_trip_no_duplicate(): @@ -2990,7 +2990,7 @@ def test_thinking_block_signature_is_not_forwarded_to_gemini() -> None: def test_anthropic_messages_history_replays_to_gemini_without_claude_signature() -> None: - from litellm.litellm_core_utils.prompt_templates.factory import _get_dummy_thought_signature + from litellm.litellm_core_utils.prompt_templates.factory import get_dummy_thought_signature from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter thinking: Final = "I should look the weather up before answering." @@ -3013,6 +3013,6 @@ def test_anthropic_messages_history_replays_to_gemini_without_claude_signature() parts: Final = _parts_of(gemini_convert_messages_with_history(messages=chat_messages, model="gemini-3.8-flash")) signatures: Final = [part["thoughtSignature"] for part in parts if "thoughtSignature" in part] - assert signatures == [_get_dummy_thought_signature()] + assert signatures == [get_dummy_thought_signature()] assert [part for part in parts if part.get("text") == thinking] == [{"thought": True, "text": thinking}] - assert next(part for part in parts if "function_call" in part)["thoughtSignature"] == _get_dummy_thought_signature() + assert next(part for part in parts if "function_call" in part)["thoughtSignature"] == get_dummy_thought_signature() diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 98d1b8f1300..2eecb83e2df 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -328,7 +328,9 @@ def test_chat_audio_drop_params_keeps_every_voice_config_field_gemini_accepts(sp def test_chat_audio_reads_snake_case_wrapper_beside_null_camel_case_wrapper(): mapped: Final = VertexGeminiConfig().map_openai_params( - non_default_params={"audio": {"format": "pcm16", "speechConfig": None, "speech_config": {"voice_name": "Kore"}}}, + non_default_params={ + "audio": {"format": "pcm16", "speechConfig": None, "speech_config": {"voice_name": "Kore"}} + }, optional_params={}, model="gemini-2.5-flash-preview-tts", drop_params=False, @@ -3468,9 +3470,7 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped(): async def dispatch_failure_handlers(self, *args, **kwargs): return None - async def dispatch_attempt_failure_handlers( - self, exception: Exception, traceback_exception: str - ) -> None: + async def dispatch_attempt_failure_handlers(self, exception: Exception, traceback_exception: str) -> None: return None def attempt_scoped_copy(self) -> "DummyLogging": @@ -3600,9 +3600,7 @@ def test_grounding_annotations_preserve_all_web_sources(indices: list[int], expe ] metadata: Final = { "groundingChunks": chunks, - "groundingSupports": [ - {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": indices} - ], + "groundingSupports": [{"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": indices}], } annotations: Final = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata], "The passage") expected: Final = [ @@ -10425,6 +10423,106 @@ def test_map_openai_params_audio_voice_reaches_speech_config_mapped() -> None: assert optional_params["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" +def _top_level_format_audio(audio_format: str) -> Mapping[str, object]: + return {"voice": "Kore", "format": audio_format} + + +def _format_beside_speech_config_audio(audio_format: str) -> Mapping[str, object]: + return {"speech_config": {"voice_name": "Kore"}, "format": audio_format} + + +def _format_inside_speech_config_audio(audio_format: str) -> Mapping[str, object]: + return {"speech_config": {"voice_name": "Kore", "format": audio_format}} + + +_FORMAT_PLACEMENTS: Final = ( + _top_level_format_audio, + _format_beside_speech_config_audio, + _format_inside_speech_config_audio, +) + + +@pytest.mark.parametrize("audio_format", ("mp3", "opus", "flac", "PCM16")) +@pytest.mark.parametrize("audio", _FORMAT_PLACEMENTS) +@pytest.mark.parametrize("model", ("gemini/gemini-3.8-flash-tts", "vertex_ai/gemini-3.8-flash-tts")) +def test_chat_audio_unsupported_format_is_a_400_naming_the_model_and_format( + model: str, audio: Callable[[str], Mapping[str, object]], audio_format: str +) -> None: + with respx.mock, pytest.raises(litellm.BadRequestError) as raised: + completion( + model=model, + messages=[{"role": "user", "content": "Say hello"}], + modalities=["audio"], + audio=audio(audio_format), + api_key="test-key", + vertex_project="test-project", + vertex_location="us-central1", + ) + + assert raised.value.status_code == 400 + assert f"gemini-3.8-flash-tts does not take the audio format '{audio_format}'" in str(raised.value) + + +@pytest.mark.parametrize("audio", _FORMAT_PLACEMENTS) +@pytest.mark.parametrize("request_drop_params, global_drop_params", ((True, False), (False, True))) +def test_chat_audio_unsupported_format_is_dropped_under_drop_params( + monkeypatch: pytest.MonkeyPatch, + request_drop_params: bool, + global_drop_params: bool, + audio: Callable[[str], Mapping[str, object]], +) -> None: + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + + optional_params: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"modalities": ["audio"], "audio": audio("mp3")}, + optional_params={}, + model="gemini-3.8-flash-tts", + drop_params=request_drop_params, + ) + + assert optional_params == { + "responseModalities": ["AUDIO"], + "speechConfig": {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}}, + } + + +def test_chat_audio_unsupported_format_error_caps_the_echoed_format() -> None: + with pytest.raises(litellm.UnsupportedParamsError) as raised: + VertexGeminiConfig().map_openai_params( + non_default_params={"modalities": ["audio"], "audio": {"voice": "Kore", "format": "x" * 300}}, + optional_params={}, + model="gemini-3.8-flash-tts", + drop_params=False, + ) + + assert f"audio format '{'x' * 100}'" in str(raised.value) + assert "x" * 101 not in str(raised.value) + + +@pytest.mark.parametrize("audio_format", [5, {"codec": "mp3"}, ["x"] * 300]) +def test_chat_audio_non_string_format_error_names_the_value_as_sent(audio_format: object) -> None: + with pytest.raises(litellm.UnsupportedParamsError) as raised: + VertexGeminiConfig().map_openai_params( + non_default_params={"modalities": ["audio"], "audio": {"voice": "Kore", "format": audio_format}}, + optional_params={}, + model="gemini-3.8-flash-tts", + drop_params=False, + ) + + assert f"audio format {repr(audio_format)[:100]}. Set format" in str(raised.value) + + +def test_map_audio_params_rejects_an_unsupported_format() -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="audio format 'mp3'"): + VertexGeminiConfig().map_audio_params({"voice": "Kore", "format": "mp3"}) + + +def test_map_audio_params_maps_a_pcm16_voice() -> None: + assert VertexGeminiConfig().map_audio_params({"voice": "Kore", "format": "pcm16"}) == { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}} + } + + _CACHED_PROMPT_USAGE_METADATA = { "promptTokenCount": 1536, "candidatesTokenCount": 40, @@ -10763,6 +10861,90 @@ async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_n assert sent["includeThoughts"] is expected +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} +_GEMINI_ROUTES: Final = ( + pytest.param(VertexGeminiConfig, "vertex_ai", id="vertex_ai"), + pytest.param(GoogleAIStudioGeminiConfig, "gemini", id="gemini"), +) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_a_hosted_mcp_tool_answers_400_on_gemini_without_drop_params(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert f"gemini-3.8-flash on {route} " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + assert "litellm_proxy/mcp/" in str(raised.value) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool( + monkeypatch, caplog, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "deepwiki" not in json.dumps(mapped) + assert "get_weather" in json.dumps(mapped["tools"]) + assert any( + "`deepwiki`" in record.getMessage() and f" on {route} (" in record.getMessage() for record in caplog.records + ) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "tools" not in mapped + assert "tool_choice" not in mapped + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_keeps_tool_choice_none_for_legacy_functions_after_dropping_every_hosted_mcp_tool( + monkeypatch, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={ + "tools": [_HOSTED_MCP_TOOL], + "functions": [_WEATHER_FUNCTION_TOOL["function"]], + "tool_choice": "none", + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "deepwiki" not in json.dumps(mapped) + assert "get_weather" in json.dumps(mapped["tools"]) + assert mapped["tool_choice"] == {"functionCallingConfig": {"mode": "NONE"}} + + class _BodyInChunks(httpx.SyncByteStream, httpx.AsyncByteStream): def __init__(self, chunks: tuple[bytes, ...]) -> None: self._chunks: Final = chunks diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 3b222be6e5e..32d68249858 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,16 +1,19 @@ import copy +import threading import time from collections.abc import Iterator from contextlib import contextmanager from enum import Enum from typing import Final -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from pydantic import JsonValue, TypeAdapter import litellm +from litellm import anthropic_beta_headers_manager from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -1826,6 +1829,100 @@ async def test_vertex_ai_token_counter_forwards_system_and_tools_to_partner_requ assert without_optional_fields.total_tokens == 37 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("edit_type", "expected_beta"), + [("clear_tool_uses_20250919", "context-management-2025-06-27"), ("compact_20260112", "compact-2026-01-12")], +) +async def test_vertex_claude_count_forwards_the_context_edits_and_beta_vertex_inference_sends( + edit_type: str, expected_beta: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class FakeResponse: + status_code = 200 + + def json(self) -> dict[str, int]: + return {"input_tokens": 21} + + class FakeHttpClient: + posted: tuple[tuple[dict[str, str], dict[str, object]], ...] = () + + async def post( + self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float + ) -> FakeResponse: + self.posted = (*self.posted, (headers, json)) + return FakeResponse() + + fake_http_client: Final = FakeHttpClient() + spec: Final[dict[str, JsonValue]] = {"edits": [{"type": edit_type}]} + application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None) + monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project")) + with patch.object(handler, "get_async_httpx_client", return_value=fake_http_client): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}}, + messages_format="anthropic", + context_management=CountContextManagement(spec=spec, drop_invalid_edits=False), + ) + + assert result is not None + assert result.total_tokens == 21 + ((headers, body),) = fake_http_client.posted + assert body["context_management"] == spec + assert headers["anthropic-beta"] == expected_beta + + +@pytest.mark.asyncio +async def test_vertex_claude_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class FakeResponse: + status_code = 200 + + def json(self) -> dict[str, int]: + return {"input_tokens": 21} + + class FakeHttpClient: + async def post( + self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float + ) -> FakeResponse: + return FakeResponse() + + class FetchThreads: + seen: tuple[int, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.seen = (*self.seen, threading.get_ident()) + return httpx.Response(200, json={}) + + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None) + fetch_threads: Final = FetchThreads() + application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None) + monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project")) + with patch.object(handler, "get_async_httpx_client", return_value=FakeHttpClient()), respx.mock: + respx.get(litellm.anthropic_beta_headers_url).mock(side_effect=fetch_threads) + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}}, + messages_format="anthropic", + context_management=CountContextManagement( + spec={"edits": [{"type": "clear_tool_uses_20250919"}]}, drop_invalid_edits=False + ), + ) + + assert result is not None + assert len(fetch_threads.seen) == 1 + assert threading.get_ident() not in fetch_threads.seen + + _PARTNER_CHAT_TOOL_TURNS: Final[list[dict[str, object]]] = [ {"role": "user", "content": "Weather in SF?"}, { diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 0a582c95749..2be6794186e 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -640,36 +640,16 @@ async def test_vertex_ai_anthropic_prompt_caching_does_not_add_beta_header( assert "anthropic_beta" not in json.loads(request.content) -def test_vertex_ai_partner_models_anthropic_remove_prompt_caching_scope_beta_header(): - """ - Test that remove_unsupported_beta correctly filters out prompt-caching-scope-2026-01-05 - from the anthropic-beta headers. - """ +def test_vertex_ai_partner_models_anthropic_forwards_prompt_caching_scope_beta_header( + local_beta_headers_config: None, +) -> None: + """A live Vertex rawPredict probe on 2026-10-28 accepted prompt-caching-scope-2026-01-05, so it goes upstream + while a beta missing from the vertex_ai column is still dropped.""" + headers: Final = update_headers_with_filtered_beta( + {"anthropic-beta": "other-feature,prompt-caching-scope-2026-01-05,web-search-2025-03-05"}, "vertex_ai" + ) - # This beta header should be removed - PROMPT_CACHING_BETA_HEADER = "prompt-caching-scope-2026-01-05" - headers = { - "anthropic-beta": f"other-feature,{PROMPT_CACHING_BETA_HEADER},web-search-2025-03-05" - } - - headers = update_headers_with_filtered_beta(headers, "vertex_ai") - - beta_header = headers.get("anthropic-beta") - assert PROMPT_CACHING_BETA_HEADER not in ( - beta_header or "" - ), f"{PROMPT_CACHING_BETA_HEADER} should be filtered out" - assert "other-feature" not in ( - beta_header or "" - ), "Other non-excluded beta headers should remain" - assert "web-search-2025-03-05" in ( - beta_header or "" - ), "Other non-excluded beta headers should remain" - # If prompt-caching was the only value, header should be removed completely - headers2 = {"anthropic-beta": PROMPT_CACHING_BETA_HEADER} - headers2 = update_headers_with_filtered_beta(headers2, "vertex_ai") - assert ( - "anthropic-beta" not in headers2 - ), "Header should be removed if no supported values remain" + assert headers == {"anthropic-beta": "prompt-caching-scope-2026-01-05,web-search-2025-03-05"} def test_vertex_ai_anthropic_output_config_effort_only_forwarded(): diff --git a/tests/unit/llms/vllm/test_vllm_common_utils.py b/tests/unit/llms/vllm/test_vllm_common_utils.py new file mode 100644 index 00000000000..aa8a4852273 --- /dev/null +++ b/tests/unit/llms/vllm/test_vllm_common_utils.py @@ -0,0 +1,27 @@ +from typing import Final + +import httpx +import pytest +import respx + +from litellm.llms.vllm.common_utils import HostedVLLMModelInfo + + +@respx.mock +def test_hosted_vllm_get_models_lists_the_server_ids_sending_the_key_as_a_bearer_token() -> None: + respx.get("http://vllm-sdk.test:8000/v1/models", headers={"Authorization": "Bearer vllm-sdk-key"}).mock( + return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}, {"id": "meta-llama/Llama-3.1-8B-Instruct"}]}) + ) + + listed: Final = HostedVLLMModelInfo().get_models(api_key="vllm-sdk-key", api_base="http://vllm-sdk.test:8000/v1") + + assert listed == ["qwen2.5", "meta-llama/Llama-3.1-8B-Instruct"] + + +def test_hosted_vllm_get_models_without_a_server_names_the_setting_that_is_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + + with pytest.raises(ValueError, match="HOSTED_VLLM_API_BASE"): + HostedVLLMModelInfo().get_models() diff --git a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py index 2b682fda982..37dfdbf489c 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py @@ -364,6 +364,10 @@ async def test_unknown_invocation_target_leaves_billing_unset(monkeypatch: pytes ("/v1beta/models/gemini-model:generateContent", "POST", True), ("/v1/realtime", "GET", True), ("/v1/realtime", "POST", False), + ("/v1/realtime/translations", "GET", True), + ("/v1/realtime/translations", "POST", False), + ("/v1/realtime/translations/client_secrets", "POST", False), + ("/v1/realtime/translations/calls", "POST", False), ("/v1/realtime/client_secrets", "POST", False), ("/mcp/tools/call", "POST", True), ("/a2a/target/message/send", "POST", True), @@ -437,14 +441,21 @@ def test_managed_inference_ignores_unsupported_query_model(): ) -@pytest.mark.parametrize("route", ["/realtime", "/v1/realtime", "/openai/v1/realtime"]) +@pytest.mark.parametrize( + "route", + ( + "/realtime", + "/v1/realtime", + "/openai/v1/realtime", + ), +) def test_managed_realtime_requires_a_model_and_ignores_completion_defaults(route: str) -> None: from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_inference_request with pytest.raises(HTTPException, match="explicit or configured model"): managed_inference_request(route, {}, {"completion_model": "allowed-default"}, "cli") assert ( - managed_inference_request(route, {"model": "requested"}, {"completion_model": "allowed-default"}, "cli")[ + managed_inference_request(route, {}, {"completion_model": "allowed-default"}, "cli", query_model="requested")[ "model" ] == "requested" @@ -465,6 +476,20 @@ def test_managed_live_upgrade_defers_the_model_to_the_session_start_frame(route: ) == {"model": "frame-model"} +@pytest.mark.parametrize("route", ("/realtime/translations", "/v1/realtime/translations")) +@pytest.mark.parametrize("query_model", (None, "requested")) +def test_managed_translation_authorizes_the_same_model_as_the_endpoint(route: str, query_model: str | None) -> None: + from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_inference_request + from litellm.proxy.common_utils.http_parsing_utils import resolve_realtime_route_model + + admitted: Final = managed_inference_request( + route, {}, {"completion_model": "unrelated-default"}, "cli", query_model=query_model + ) + + assert admitted["model"] == resolve_realtime_route_model(query_model, None, is_translation=True) + assert isinstance(admitted["model"], str) and admitted["model"] + + @pytest.mark.parametrize("mode,user", [("autonomous", None), ("delegated", "verified-human")]) def test_matching_identity_revision_and_execution_mode_pass_admission(mode: str, user: str | None) -> None: context: Final = ManagedAgentContext.model_validate( @@ -573,21 +598,41 @@ async def test_ordinary_agent_admission_preserves_legacy_authentication( @pytest.mark.parametrize( "route", - tuple(dict.fromkeys( - LiteLLMRoutes.openai_routes.value - + LiteLLMRoutes.anthropic_routes.value - + LiteLLMRoutes.google_routes.value - )), + tuple( + dict.fromkeys( + LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.anthropic_routes.value + LiteLLMRoutes.google_routes.value + ) + ), ) def test_registered_inference_routes_have_an_explicit_managed_access_decision(route: str) -> None: from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_route_allowed normalized: Final = route.removeprefix("/openai").removeprefix("/v1beta").removeprefix("/v1") - unsupported: Final = normalized.startswith(( - "/videos", "/batches", "/files", "/fine_tuning", "/assistants", "/threads", "/utils/", - "/vector_stores", "/vector_store/", "/search", "/containers", "/skills", "/claude-code/", - "/interactions", "/agents", "/responses/{", "/responses/input_tokens", - "/realtime/client_secrets", "/realtime/calls", "/realtime/transcription_sessions", - )) or normalized in ("/models", "/cursor/models", "/cursor/v1/models", "/claude_code_gateway/v1/models") + unsupported: Final = normalized.startswith( + ( + "/videos", + "/batches", + "/files", + "/fine_tuning", + "/assistants", + "/threads", + "/utils/", + "/vector_stores", + "/vector_store/", + "/search", + "/containers", + "/skills", + "/claude-code/", + "/interactions", + "/agents", + "/responses/{", + "/responses/input_tokens", + "/realtime/client_secrets", + "/realtime/calls", + "/realtime/transcription_sessions", + "/realtime/translations/client_secrets", + "/realtime/translations/calls", + ) + ) or normalized in ("/models", "/cursor/models", "/cursor/v1/models", "/claude_code_gateway/v1/models") concrete: Final = route.split("?")[0].replace("{model}", "model").replace("{model_name:path}", "model") assert managed_agent_route_allowed(concrete, None) is not unsupported, route diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 604fc84a096..553c3d4a9d7 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -319,6 +319,22 @@ class TestProxyExceptionAnthropicEnvelope: assert fallback.status_code == 500 assert json.loads(fallback.body)["error"]["type"] == "api_error" + @pytest.mark.parametrize("status_code", [402, 409, 422, 504]) + def test_a_failure_inside_the_handler_keeps_the_api_error_type_for_a_status_outside_the_table( + self, status_code: int + ): + import litellm.proxy.anthropic_endpoints.endpoints as ep + from litellm.proxy._types import ProxyException + + request: Final = MagicMock(headers={}) + + response: Final = ep._anthropic_error_json_response( + ProxyException(message="scripted", type="budget_exceeded", param=None, code=status_code), request + ) + + assert response.status_code == status_code + assert json.loads(response.body)["error"] == {"type": "api_error", "message": "scripted"} + @staticmethod def _call_id_error_response(general_settings, provider_specific_fields=None): import litellm.proxy.anthropic_endpoints.endpoints as ep diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 8e9e87f370a..7a060a214f7 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -10998,139 +10998,6 @@ async def test_key_budget_error_keeps_the_masked_key_name(key_name): assert f"Key=prod-key ({key_name}) Current cost" in message -class _UntouchedPrisma: - def __getattr__(self, name: str) -> object: - raise AssertionError(f"database reached through {name}") - - -class _MissingUserPrisma: - class db: - class litellm_usertable: - @staticmethod - async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: - return None - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestricted(): - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - scopes = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="default_user_id"), - prisma_client=_MissingUserPrisma(), - user_api_key_cache=cache, - proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), - ) - - assert [list(scope) for scope in scopes] == [[], [], [], [], []] - - -class _UnreachableUserPrisma: - class db: - class litellm_usertable: - @staticmethod - async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: - raise RuntimeError("database gone") - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_surfaces_a_failed_user_lookup(): - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - with pytest.raises(ValueError, match="database gone"): - await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), - prisma_client=_UnreachableUserPrisma(), - user_api_key_cache=cache, - proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), - ) - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_reads_every_level_from_cache(): - from litellm.proxy._types import ( - LiteLLM_BudgetTable, - LiteLLM_ProjectTableCachedObj, - LiteLLM_TeamMembership, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - ) - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import ( - UserApiKeyCache, - team_membership_reservation_cache_key, - ) - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - proxy_logging_obj = ProxyLogging(user_api_key_cache=cache) - await cache.async_set_cache( - key="team_id:team-fake", value=LiteLLM_TeamTableCachedObj(team_id="team-fake", models=["gpt-4o"]) - ) - await cache.async_set_cache( - key=team_membership_reservation_cache_key(user_id="user-fake", team_id="team-fake"), - value=LiteLLM_TeamMembership( - user_id="user-fake", - team_id="team-fake", - litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["gpt-4o-mini"]), - ), - ) - await cache.async_set_cache( - key="project_id:project-fake", - value=LiteLLM_ProjectTableCachedObj(project_id="project-fake", models=["gpt-4.1"]), - ) - await cache.async_set_cache(key="user-fake", value=LiteLLM_UserTable(user_id="user-fake", models=["o3"])) - prisma_client = _UntouchedPrisma() - - team_scoped = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth( - token="hashed-fake", - models=["all-team-models"], - team_models=["gpt-4o", "gpt-4o-mini"], - user_id="user-fake", - team_id="team-fake", - project_id="project-fake", - ), - prisma_client=prisma_client, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - personal = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), - prisma_client=prisma_client, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - without_database = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth( - token="hashed-fake", - models=["gpt-4o"], - team_models=["gpt-4o-mini"], - user_id="user-fake", - team_id="team-fake", - ), - prisma_client=None, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - - assert [list(scope) for scope in team_scoped] == [ - ["gpt-4o", "gpt-4o-mini"], - ["gpt-4o"], - ["gpt-4o-mini"], - [], - ["gpt-4.1"], - ] - assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []] - assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]] - - @pytest.mark.asyncio @pytest.mark.parametrize("channel", ["team", "key"]) async def test_access_group_model_fallback_uses_the_injected_database(channel: str) -> None: diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index d0156ff0195..bc9366aa607 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -1,13 +1,25 @@ -from collections.abc import Generator, Mapping +import datetime +import ipaddress +import json +import ssl +import threading +from collections.abc import Generator, Iterator, Mapping from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path from typing import Final from unittest.mock import patch import httpx import pytest import respx +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.x509.oid import NameOID from litellm import Router +from litellm.utils import AZURE_MODEL_PLACEHOLDER def test_get_team_models_for_all_models_and_team_only_models(): @@ -430,6 +442,262 @@ def test_wildcard_credential_hydration_preserves_deployment_params( } +@respx.mock +def test_hosted_vllm_wildcard_lists_the_server_models_without_an_api_key(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + models_route: Final = respx.get("http://vllm-without-key.test:8000/v1/models").mock( + return_value=httpx.Response( + 200, + json={ + "data": [ + {"id": "meta-llama/Llama-3.1-8B-Instruct"}, + {"id": "qwen2.5"}, + {"id": "openai/gpt-oss-20b"}, + ] + }, + ) + ) + + listed: Final = get_known_models_from_wildcard( + "hosted_vllm/*", + LiteLLM_Params(model="hosted_vllm/*", api_base="http://vllm-without-key.test:8000/v1"), + ) + + assert listed == [ + "hosted_vllm/meta-llama/Llama-3.1-8B-Instruct", + "hosted_vllm/qwen2.5", + "hosted_vllm/openai/gpt-oss-20b", + ] + assert "authorization" not in models_route.calls.last.request.headers + + +@respx.mock +def test_hosted_vllm_wildcard_sends_the_configured_api_key_as_a_bearer_token(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + models_route: Final = respx.get( + "http://vllm-with-key.test:8000/v1/models", headers={"Authorization": "Bearer vllm-test-key"} + ).mock(return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}]})) + + listed: Final = get_known_models_from_wildcard( + "hosted_vllm/*", + LiteLLM_Params(model="hosted_vllm/*", api_base="http://vllm-with-key.test:8000/v1", api_key="vllm-test-key"), + ) + + assert listed == ["hosted_vllm/qwen2.5"] + assert models_route.call_count == 1 + + +@respx.mock +def test_hosted_vllm_wildcard_without_api_base_lists_the_server_chat_would_call(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://hosted-vllm-env.test:8000/v1") + monkeypatch.setenv("HOSTED_VLLM_API_KEY", "hosted-vllm-env-key") + monkeypatch.setenv("VLLM_API_BASE", "http://vllm-passthrough-env.test:8000") + passthrough_server_route: Final = respx.get("http://vllm-passthrough-env.test:8000/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "model-on-the-passthrough-server"}]}) + ) + respx.get( + "http://hosted-vllm-env.test:8000/v1/models", headers={"Authorization": "Bearer hosted-vllm-env-key"} + ).mock(return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}]})) + + listed: Final = get_known_models_from_wildcard( + "hosted_vllm/*", LiteLLM_Params(model="hosted_vllm/*", custom_llm_provider="hosted_vllm") + ) + + assert listed == ["hosted_vllm/qwen2.5"] + assert not passthrough_server_route.called + + +@respx.mock +def test_hosted_vllm_wildcard_listing_does_not_wait_the_completion_timeout_on_a_hung_server( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm.constants import NON_LLM_CONNECTION_TIMEOUT + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + models_route: Final = respx.get("http://vllm-timeout.test:8000/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}]}) + ) + + get_known_models_from_wildcard( + "hosted_vllm/*", LiteLLM_Params(model="hosted_vllm/*", api_base="http://vllm-timeout.test:8000/v1") + ) + + read_timeout: Final = models_route.calls.last.request.extensions["timeout"]["read"] + assert read_timeout is not None + assert read_timeout <= NON_LLM_CONNECTION_TIMEOUT + + +@respx.mock +def test_hosted_vllm_wildcard_listing_sends_the_deployment_extra_headers(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + respx.get( + "http://vllm-behind-gateway.test:8000/v1/models", + headers={"X-Tenant": "team-a", "Authorization": "Bearer vllm-test-key"}, + ).mock(return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}]})) + + listed: Final = get_known_models_from_wildcard( + "hosted_vllm/*", + LiteLLM_Params( + model="hosted_vllm/*", + api_base="http://vllm-behind-gateway.test:8000/v1", + api_key="vllm-test-key", + extra_headers={"X-Tenant": "team-a"}, + ), + ) + + assert listed == ["hosted_vllm/qwen2.5"] + + + +@respx.mock +def test_vllm_wildcard_never_lists_the_ids_of_a_configured_vllm_server(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setenv("VLLM_API_BASE", "http://vllm-library-env.test:8000") + monkeypatch.setenv("VLLM_API_KEY", "vllm-library-env-key") + server_route: Final = respx.get(url__regex=r"http://vllm-library(-env)?\.test:8000/.*").mock( + return_value=httpx.Response(200, json={"data": [{"id": "model-only-a-server-serves"}]}) + ) + + listed: Final = get_known_models_from_wildcard( + "vllm/*", + LiteLLM_Params(model="vllm/*", api_base="http://vllm-library.test:8000", api_key="vllm-library-key"), + ) + + assert "vllm/model-only-a-server-serves" not in listed + assert not server_route.called + +_SELF_SIGNED_CERT_NOT_VALID_BEFORE: Final = datetime.datetime(2000, 1, 1, tzinfo=datetime.UTC) +_SELF_SIGNED_CERT_NOT_VALID_AFTER: Final = datetime.datetime(9999, 12, 31, tzinfo=datetime.UTC) + + +def _self_signed_cert(directory: Path) -> tuple[Path, Path]: + key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "127.0.0.1")]) + cert: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(_SELF_SIGNED_CERT_NOT_VALID_BEFORE) + .not_valid_after(_SELF_SIGNED_CERT_NOT_VALID_AFTER) + .add_extension( + x509.SubjectAlternativeName([x509.IPAddress(ipaddress.IPv4Address("127.0.0.1"))]), critical=False + ) + .sign(key, hashes.SHA256()) + ) + cert_path: Final = directory / "cert.pem" + key_path: Final = directory / "key.pem" + cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + key_path.write_bytes( + key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + ) + return cert_path, key_path + + +@pytest.fixture +def self_signed_vllm_api_base(tmp_path: Path) -> Iterator[str]: + cert_path, key_path = _self_signed_cert(tmp_path) + + class Handler(BaseHTTPRequestHandler): + timeout = 5 + + def do_GET(self) -> None: + body: Final = json.dumps({"data": [{"id": "qwen2.5"}]}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(cert_path, key_path) + server.socket = context.wrap_socket(server.socket, server_side=True, do_handshake_on_connect=False) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"https://127.0.0.1:{server.server_address[1]}/v1" + finally: + server.shutdown() + server.server_close() + + +def test_hosted_vllm_wildcard_lists_a_self_signed_server_only_when_the_deployment_turns_off_ssl_verify( + monkeypatch: pytest.MonkeyPatch, self_signed_vllm_api_base: str +) -> None: + import litellm + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.llms.custom_httpx import http_handler + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setattr(litellm, "ssl_verify", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.setattr(http_handler, "_ssl_context_cache", {}) + monkeypatch.delenv("SSL_VERIFY", raising=False) + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + + verified: Final = get_known_models_from_wildcard( + "hosted_vllm/*", LiteLLM_Params(model="hosted_vllm/*", api_base=self_signed_vllm_api_base) + ) + unverified: Final = get_known_models_from_wildcard( + "hosted_vllm/*", + LiteLLM_Params(model="hosted_vllm/*", api_base=self_signed_vllm_api_base, ssl_verify=False), + ) + + assert verified == [] + assert unverified == ["hosted_vllm/qwen2.5"] + + +@respx.mock +def test_custom_prefix_over_a_litellm_proxy_wildcard_lists_upstream_models_under_the_prefix_alone( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + respx.get("http://upstream-proxy.test:4000/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "gpt-4o"}, {"id": "claude-sonnet"}]}) + ) + + listed: Final = get_known_models_from_wildcard( + "team-proxy/*", + LiteLLM_Params(model="litellm_proxy/*", api_base="http://upstream-proxy.test:4000", api_key="sk-upstream"), + ) + + assert listed == ["team-proxy/gpt-4o", "team-proxy/claude-sonnet"] + + def test_wildcard_custom_prefix_does_not_stack_provider_prefix(monkeypatch): """Regression test for #30358. @@ -834,6 +1102,56 @@ def test_expand_wildcard_invalid_litellm_params_passthrough(): assert result == [deployment] + +@pytest.mark.parametrize( + ("allowed_model_names", "expected_names", "expected_listings"), + [ + pytest.param(None, ["qwen-qwen2.5"], 1, id="unscoped"), + pytest.param(["qwen-qwen2.5"], ["qwen-qwen2.5"], 1, id="granted-a-listed-name"), + pytest.param(["hosted_vllm/qwen2.5"], [], 0, id="granted-only-the-provider-id"), + ], +) +@respx.mock +def test_expand_wildcard_alias_without_a_slash_lists_and_filters_by_its_own_pattern( + monkeypatch: pytest.MonkeyPatch, + request: pytest.FixtureRequest, + allowed_model_names: list[str] | None, + expected_names: list[str], + expected_listings: int, +) -> None: + import litellm + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + api_base: Final = f"http://vllm-alias-{request.node.callspec.id}.test:8000/v1" + listing: Final = respx.get(f"{api_base}/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}]}) + ) + deployment: Final = {"model_name": "qwen-*", "litellm_params": {"model": "hosted_vllm/*", "api_base": api_base}} + + rows: Final = expand_wildcard_deployments_for_model_info([deployment], allowed_model_names=allowed_model_names) + + assert ([row["model_name"] for row in rows], listing.call_count) == (expected_names, expected_listings) + + +@pytest.mark.parametrize( + ("allowed_model_names", "keeps_the_alias"), + [ + pytest.param(["gpt-4o-mini"], True, id="granted-a-name-the-alias-matches"), + pytest.param(["claude-sonnet"], False, id="granted-only-other-names"), + ], +) +def test_expand_wildcard_alias_over_one_model_is_kept_only_for_keys_with_a_name_it_matches( + allowed_model_names: list[str], keeps_the_alias: bool +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info + + deployment: Final = {"model_name": "gpt-4o*", "litellm_params": {"model": "openai/gpt-4o"}} + + rows: Final = expand_wildcard_deployments_for_model_info([deployment], allowed_model_names=allowed_model_names) + + assert rows == ([deployment] if keeps_the_alias else []) + def test_get_complete_model_list_excludes_wildcard_routes_by_default(): """Regression (LIT-4108): a wildcard with a matching router deployment leaked into /v1/models.""" from litellm import Router @@ -1327,6 +1645,40 @@ def test_key_granted_openrouter_wildcard_without_an_openrouter_deployment_lists_ assert (listed, catalog.call_count) == ([], 0) +@respx.mock +def test_key_granted_hosted_vllm_wildcard_without_a_hosted_vllm_deployment_lists_no_environment_server( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://hosted-vllm-environment-only.test:8000/v1") + environment_server: Final = respx.get("http://hosted-vllm-environment-only.test:8000/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}]}) + ) + router: Final = Router( + model_list=[ + { + "model_name": "vllm-team/*", + "litellm_params": {"model": "hosted_vllm/*", "api_base": "http://vllm-team.test:8000/v1"}, + } + ] + ) + + listed: Final = get_complete_model_list( + key_models=["hosted_vllm/*"], + team_models=[], + proxy_model_list=router.model_names, + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + + assert (listed, environment_server.call_count) == ([], 0) + + @contextmanager def _openai_catalog_behind_a_catch_all( monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str] @@ -1748,6 +2100,34 @@ _ROUTE_CF_IDS: Final = {_ROUTE_CF_ID, _ROUTE_CF_SUFFIXED_ID} pytest.param( "*", "openai/*", "openai/gpt-partial-route-9*", set(), id="provider-prefixed-key-filter-over-catch-all" ), + pytest.param( + "gpt-partial-route-*", + "openai/gpt-partial-route-*", + None, + {f"openai/{_ROUTE_OPENAI_ID}"}, + id="slashless-name-repeating-the-deployment-partial-word", + ), + pytest.param("team-*", "cloudflare/@cf/*", None, _ROUTE_CF_IDS, id="slashless-alias-over-partial-deployment"), + pytest.param( + "team-*", "openai/*", None, {f"openai/{_ROUTE_OPENAI_ID}"}, id="slashless-alias-over-provider-catch-all" + ), + pytest.param( + "team-*", "cloudflare/@cf/*/*", None, set(), id="slashless-alias-with-fewer-wildcards-than-its-deployment" + ), + pytest.param( + "*partial-route-9*", + "openai/gpt-partial-route-9*", + None, + {f"openai/{_ROUTE_OPENAI_ID}"}, + id="slashless-name-with-more-wildcards-than-its-deployment", + ), + pytest.param( + "gemini/*partial-route-9*", + "gemini/gemini-*", + None, + {_ROUTE_GEMINI_ID}, + id="slashed-name-with-more-wildcards-than-its-deployment", + ), pytest.param("*", "cloudflare/@cf/*", None, _ROUTE_CF_IDS, id="catch-all-over-partial-deployment"), pytest.param("*", "gemini/gemini-partial-route-9*", None, set(), id="catch-all-over-mid-word-deployment"), pytest.param("*", "openai/responses/gpt-partial-route-9*", None, set(), id="catch-all-over-mid-word-route"), @@ -1825,6 +2205,114 @@ def test_every_listed_partial_wildcard_id_routes_to_the_catalog_model_it_names( assert listed_model_info_names == sorted(reachable) +@pytest.mark.parametrize( + ("model_name", "expected_v1_models"), + [ + pytest.param("azure/*", [f"azure/{AZURE_MODEL_PLACEHOLDER}"], id="provider-wildcard-lists-as-before"), + pytest.param("*", [AZURE_MODEL_PLACEHOLDER], id="catch-all-lists-as-before"), + pytest.param("my-*", [], id="slashless-alias-skips-the-placeholder"), + ], +) +def test_only_slashless_azure_wildcards_skip_the_catalog_placeholder( + model_name: str, expected_v1_models: list[str] +) -> None: + from litellm import Router + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + + router: Final = Router(model_list=[{"model_name": model_name, "litellm_params": {"model": "azure/*"}}]) + + v1_models: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=[model_name], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + model_info_names: Final = [ + str(row["model_name"]) + for row in expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or []], router + ) + ] + + assert v1_models == expected_v1_models + assert model_info_names == (expected_v1_models or [model_name]) + + +@pytest.mark.parametrize( + ("model_name", "deployment_model"), + [ + pytest.param("gpt-*", "gpt-*", id="same-wildcards-as-its-deployment"), + pytest.param("*gpt-5*", "gpt-5*", id="more-wildcards-than-its-deployment"), + ], +) +def test_slashless_wildcard_over_a_bare_deployment_lists_its_custom_llm_provider_catalog( + model_name: str, deployment_model: str +) -> None: + from fnmatch import fnmatchcase + + from litellm import Router + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + from litellm.utils import get_valid_models + + router: Final = Router( + model_list=[ + { + "model_name": model_name, + "litellm_params": {"model": deployment_model, "custom_llm_provider": "openai", "api_key": "sk-test"}, + } + ] + ) + + v1_models: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=[model_name], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + model_info_rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or []], router + ) + openai_ids: Final = frozenset(model.removeprefix("openai/") for model in get_valid_models(custom_llm_provider="openai")) + routed_models: Final = { + listed_id: [deployment["litellm_params"]["model"] for deployment in router.wildcard_route_deployments(listed_id, None)] + for listed_id in v1_models + } + + assert v1_models + assert [ + listed_id for listed_id in v1_models if listed_id not in openai_ids or not fnmatchcase(listed_id, model_name) + ] == [] + assert {listed_id: models for listed_id, models in routed_models.items() if models != [listed_id]} == {} + assert sorted(str(row["model_name"]) for row in model_info_rows) == sorted(v1_models) + assert [row for row in model_info_rows if row["litellm_params"]["model"] != row["model_name"]] == [] + + +@pytest.mark.parametrize( + ("model_name", "deployment_model"), + [ + pytest.param("*", "*", id="catch-all-over-a-bare-catch-all"), + pytest.param("foo/gpt-*", "gpt-*", id="slashed-name-over-a-bare-deployment"), + pytest.param("gpt-*", "openai/gpt-*", id="slashless-name-over-a-provider-prefixed-deployment"), + ], +) +def test_custom_llm_provider_changes_no_listing_outside_slashless_names_over_bare_deployments( + model_name: str, deployment_model: str +) -> None: + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + with_provider: Final = get_known_models_from_wildcard( + model_name, LiteLLM_Params(model=deployment_model, custom_llm_provider="openai") + ) + without_provider: Final = get_known_models_from_wildcard(model_name, LiteLLM_Params(model=deployment_model)) + + assert with_provider == without_provider + + def test_inference_provider_route_lists_every_id_the_plain_provider_wildcard_lists() -> None: from litellm.proxy.auth.model_checks import get_known_models_from_wildcard from litellm.types.router import LiteLLM_Params diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index 5df3fbab9a6..efb4451f2fc 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -818,6 +818,52 @@ def test_realtime_webrtc_http_routes_classified_as_llm_api(route): assert RouteChecks.is_management_route(route=route) is False +@pytest.mark.parametrize( + "route", + [ + "/realtime/translations", + "/v1/realtime/translations", + "/realtime/translations/client_secrets", + "/v1/realtime/translations/client_secrets", + "/realtime/translations/calls", + "/v1/realtime/translations/calls", + ], +) +def test_realtime_translation_routes_allowed_for_openai_virtual_keys(route): + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["openai_routes"], + ) + + assert RouteChecks.is_llm_api_route(route=route) is True + assert RouteChecks.is_management_route(route=route) is False + assert ( + RouteChecks.is_virtual_key_allowed_to_call_route( + route=route, + valid_token=valid_token, + ) + is True + ) + + +@pytest.mark.parametrize( + "route", + [ + "/openai/v1/realtime/translations", + "/openai/v1/realtime/translations/client_secrets", + "/openai/v1/realtime/translations/calls", + ], +) +def test_openai_prefixed_realtime_translation_routes_denied_for_openai_virtual_keys(route: str) -> None: + valid_token: Final = UserAPIKeyAuth(user_id="test_user", allowed_routes=["openai_routes"]) + + with pytest.raises(HTTPException) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token) + + assert exc_info.value.status_code == 403 + assert f"Tried to call route: {route}" in str(exc_info.value.detail) + + def test_virtual_key_allowed_routes_with_litellm_routes_member_name_denied(): """Test that virtual key is denied when route is not in the allowed LiteLLMRoutes group""" diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 81ea0e3f60b..43a096e0f42 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -931,6 +931,43 @@ async def test_user_api_key_auth_live_websocket_ignores_query_model(): mock_user_api_key_auth.assert_called_once() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("auth_name", "path"), + [("user_api_key_auth_websocket", "/v1/responses"), ("user_api_key_auth_realtime_websocket", "/v1/realtime")], +) +async def test_websocket_without_a_model_is_refused_for_a_role_limited_to_other_models( + auth_name: str, path: str +) -> None: + from litellm.proxy._types import RoleBasedPermissions, UserAPIKeyAuth + from litellm.proxy.auth import user_api_key_auth as auth_module + from litellm.proxy.auth.handle_jwt import JWTAuthManager + + async def role_limited_auth(request, api_key: str) -> UserAPIKeyAuth: + JWTAuthManager.can_rbac_role_call_model( + rbac_role=LitellmUserRoles.INTERNAL_USER, + general_settings={ + "role_permissions": [RoleBasedPermissions(role=LitellmUserRoles.INTERNAL_USER, models=["chat"])] + }, + model=json.loads(await request.body())["model"], + ) + return UserAPIKeyAuth(api_key=api_key) + + websocket: Final = MagicMock(spec=WebSocket) + websocket.query_params = {} + websocket.headers = {"authorization": "Bearer some_api_key"} + websocket.scope = {"type": "websocket", "path": path, "headers": [(b"authorization", b"Bearer some_api_key")]} + websocket.url = URL(url=path) + websocket.close = AsyncMock() + + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", side_effect=role_limited_auth): + with pytest.raises(HTTPException) as raised: + await getattr(auth_module, auth_name)(websocket) + + assert raised.value.status_code == 403 + websocket.close.assert_awaited_once_with(code=status.WS_1008_POLICY_VIOLATION) + + @pytest.mark.asyncio async def test_accepted_websocket_reauthorizes_the_frame_model_without_closing(): from litellm.proxy._types import ProxyException diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 10738eabddf..17e6c004566 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -12872,6 +12872,188 @@ async def test_websocket_auth_hands_the_reservation_to_the_socket_state(): assert websocket.scope["state"]["budget_reservation"] is reservation +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ("/realtime/translations", "/v1/realtime/translations")) +@pytest.mark.parametrize("query", (b"", b"model=gpt-realtime-translate"), ids=("default_model", "explicit_model")) +async def test_translation_websocket_auth_enforces_the_default_model_budget( + monkeypatch: pytest.MonkeyPatch, path: str, query: bytes +) -> None: + from fastapi import WebSocket + + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_realtime_websocket + + model: Final = "gpt-realtime-translate" + token: Final = UserAPIKeyAuth( + token="test-key", + models=[model], + model_max_budget={model: {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, _ = await _real_model_budget_limiter(key_spend={model: 10.0}) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=token, model=model) + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return token + + attrs: Final = { + **_proxy_server_attrs_for_custom_auth(user_custom_auth=authenticate), + "prisma_client": None, + "model_max_budget_limiter": limiter, + } + for name, value in attrs.items(): + monkeypatch.setattr(proxy_server, name, value) + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + websocket: Final = WebSocket( + scope={ + "type": "websocket", + "path": path, + "headers": [(b"authorization", b"Bearer sk-test")], + "query_string": query, + }, + receive=AsyncMock(), + send=AsyncMock(), + ) + + with pytest.raises(HTTPException) as denied: + await user_api_key_auth_realtime_websocket(websocket) + + assert denied.value.status_code == 403 + assert model in str(denied.value.detail) + assert "budget" in str(denied.value.detail).lower() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ("/realtime/translations/client_secrets", "/v1/realtime/translations/client_secrets")) +async def test_translation_client_secret_auth_enforces_the_session_model_budget( + monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + model: Final = "gpt-realtime-translate" + token: Final = UserAPIKeyAuth( + token="test-key", + models=[model], + model_max_budget={model: {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, _ = await _real_model_budget_limiter(key_spend={model: 10.0}) + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return token + + attrs: Final = { + **_proxy_server_attrs_for_custom_auth(user_custom_auth=authenticate), + "prisma_client": None, + "model_max_budget_limiter": limiter, + } + for name, value in attrs.items(): + monkeypatch.setattr(proxy_server, name, value) + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + body: Final = json.dumps({"session": {"type": "translation", "model": model}}).encode() + request: Final = Request( + scope={ + "type": "http", + "method": "POST", + "path": path, + "headers": [(b"authorization", b"Bearer sk-test"), (b"content-type", b"application/json")], + "query_string": b"", + }, + receive=AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False}), + ) + + with pytest.raises(ProxyException) as denied: + await user_api_key_auth(request=request, api_key="Bearer sk-test") + + assert denied.value.type == "budget_exceeded" + assert model in denied.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ("/realtime", "/v1/realtime", "/openai/v1/realtime")) +async def test_managed_agent_transcription_websocket_without_a_model_is_refused( + monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + from fastapi import WebSocket + + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_realtime_websocket + from litellm.types.agents import AgentResponse + + token: Final = UserAPIKeyAuth(token="test-key") + token.managed_agent_policy = AgentResponse.model_validate( + {"agent_id": "agent", "agent_name": "Agent", "agent_card_params": {}, "execution_mode": "both"} + ) + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return token + + attrs: Final = {**_proxy_server_attrs_for_custom_auth(user_custom_auth=authenticate), "prisma_client": None} + for name, value in attrs.items(): + monkeypatch.setattr(proxy_server, name, value) + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + websocket: Final = WebSocket( + scope={ + "type": "websocket", + "path": path, + "headers": [(b"authorization", b"Bearer sk-test")], + "query_string": b"intent=transcription", + }, + receive=AsyncMock(), + send=AsyncMock(), + ) + + with pytest.raises(HTTPException) as denied: + await user_api_key_auth_realtime_websocket(websocket) + + assert denied.value.status_code == 403 + assert "Managed inference requires an explicit or configured model" in str(denied.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ("/openai_passthrough/v1/realtime", "/v1/responses", "/vertex_ai/live")) +async def test_non_realtime_websocket_auth_checks_no_default_transcription_model( + monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + from fastapi import WebSocket + + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket + + exhausted_model: Final = "gpt-realtime-whisper" + token: Final = UserAPIKeyAuth( + token="test-key", + model_max_budget={exhausted_model: {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, _ = await _real_model_budget_limiter(key_spend={exhausted_model: 10.0}) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=token, model=exhausted_model) + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return token + + attrs: Final = { + **_proxy_server_attrs_for_custom_auth(user_custom_auth=authenticate), + "prisma_client": None, + "model_max_budget_limiter": limiter, + } + for name, value in attrs.items(): + monkeypatch.setattr(proxy_server, name, value) + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + websocket: Final = WebSocket( + scope={ + "type": "websocket", + "path": path, + "headers": [(b"authorization", b"Bearer sk-test")], + "query_string": b"intent=transcription", + }, + receive=AsyncMock(), + send=AsyncMock(), + ) + + assert (await user_api_key_auth_websocket(websocket)).token == token.token + + @pytest.mark.asyncio async def test_admission_and_budget_reservation_read_the_key_spend_counter_with_one_redis_mget(): from fastapi import Request diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 694b817ab86..2dc98b5ca6b 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -3,6 +3,7 @@ import hashlib import json import time from collections.abc import Iterable +from typing import Final from unittest.mock import patch import pytest @@ -19,7 +20,11 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( evict_and_broadcast, publish_auth_cache_invalidation, ) -from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + AUTH_OBJECTS_TARGET, + UserApiKeyCache, + budget_windows_absent_cache_key, +) class _RecordingRedisClient(Redis): @@ -243,6 +248,28 @@ async def test_subscriber_deletes_key_object_partition_entry_on_message() -> Non assert cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None +@pytest.mark.asyncio +async def test_subscriber_evicts_a_user_whose_id_looks_like_a_no_window_marker_along_with_the_marker() -> None: + user_id: Final = budget_windows_absent_cache_key("key", "alice") + cache: Final = UserApiKeyCache() + cache.set_cache(user_id, {"user_id": user_id, "models": ["gpt-4o"]}) + cache.budget_windows_absent_cache.set_cache(key=user_id, value=True) + + pubsub: Final = _QueuePubSub(initial_messages=[_invalidation_message(user_id)]) + subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + ) + subscriber.start() + try: + assert await asyncio.wait_for(pubsub.drained.wait(), timeout=5) + finally: + await subscriber.stop() + + assert cache.get_cache(user_id) is None + assert cache.budget_windows_absent_cache.get_cache(key=user_id) is None + + @pytest.mark.asyncio async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None: """ diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index 509303d9635..e3d941a6aa4 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -1,9 +1,12 @@ import json +from collections.abc import Mapping from types import MappingProxyType +from typing import Final, Literal import pytest import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy.common_utils.codex_model_catalog import ( CODEX_CATALOG_BYTE_LIMIT, CodexCatalogRow, @@ -233,6 +236,140 @@ def test_alias_with_deployment_defaults_caps_codexs_window_at_its_listed_max_inp assert json.loads(body.json)["models"][0]["context_window"] == expected +_WIDE_STOCK = MappingProxyType( + { + "gpt-6-sol": _stock("gpt-6-sol").model_copy(update={"context_window": 272000, "max_context_window": 872000}), + "gpt-5.5": _stock("gpt-5.5").model_copy(update={"context_window": 272000}), + } +) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ((500000,), 500000), + ((100000,), 100000), + ((1050000,), 872000), + ((600000, 400000), 400000), + ((600000, None), 272000), + ((100000, None), 100000), + ((None,), 272000), + ((), 272000), + ((0,), 272000), + ((-5, 400000), 272000), + ((0, 100000), 100000), + ], +) +@pytest.mark.parametrize("alias_defaults", ["codex", "deployment"]) +def test_a_known_model_and_its_alias_take_the_smallest_window_any_deployment_takes( + configured: tuple[int | None, ...], expected: int, alias_defaults: Literal["codex", "deployment"] +) -> None: + rows: Final = ( + _row("gpt-6-sol", max_input_tokens=1050000, configured_max_input_tokens=configured), + _row( + "sol-1m", + upstream_model="openai/gpt-6-sol", + max_input_tokens=1050000, + configured_max_input_tokens=configured, + ), + ) + + body: Final = codex_models_response_json( + rows, stock=_WIDE_STOCK, instructions=_PROMPT, alias_defaults=alias_defaults + ) + + assert {entry["slug"]: entry["context_window"] for entry in json.loads(body.json)["models"]} == { + "gpt-6-sol": expected, + "sol-1m": expected, + } + + +def test_a_configured_window_goes_past_codexs_own_when_the_stock_entry_sets_no_ceiling() -> None: + (entry,) = json.loads( + codex_models_response_json( + (_row("gpt-5.5", configured_max_input_tokens=(1050000,)),), stock=_WIDE_STOCK, instructions=_PROMPT + ).json + )["models"] + + assert entry["context_window"] == 1050000 + + +def test_a_deployment_without_a_limit_keeps_a_windowless_stock_entry_as_codex_ships_it() -> None: + windowless: Final = MappingProxyType({"gpt-5.5": _stock("gpt-5.5").model_copy(update={"context_window": None})}) + + (entry,) = json.loads( + codex_models_response_json( + (_row("gpt-5.5", configured_max_input_tokens=(1050000, None)),), stock=windowless, instructions=_PROMPT + ).json + )["models"] + + assert entry["context_window"] is None + + +def test_model_list_body_sizes_a_known_models_window_off_the_deployments_the_key_team_routes_to() -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra"}, + "model_info": {"max_input_tokens": 922000}, + }, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://team-2.example"}, + "model_info": {"team_id": "team-2", "max_input_tokens": 400000}, + }, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://team-3.example"}, + "model_info": {"team_id": "team-3"}, + }, + ], + model_group_alias={"astra-1m": "gpt-6-astra"}, + ) + listing: Final = tuple( + {"id": name, "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"} + for name in ("gpt-6-astra", "astra-1m") + ) + entries: Final = (("gpt-6-astra", "gpt-6-astra"), ("astra-1m", "astra-1m")) + stock: Final = bundled_codex_models()["gpt-6-astra"] + assert stock.context_window is not None and stock.max_context_window is not None + assert stock.context_window < 400000 < stock.max_context_window < 922000 + + def windows(team_id: str) -> dict[str, int | None]: + models: Final = json.loads(codex_model_list_body(listing, entries, router, team_id))["models"] + return {entry["slug"]: entry["context_window"] for entry in models} + + assert windows("team-1") == {"gpt-6-astra": stock.max_context_window, "astra-1m": stock.max_context_window} + assert windows("team-2") == {"gpt-6-astra": 400000, "astra-1m": 400000} + assert windows("team-3") == {"gpt-6-astra": stock.context_window, "astra-1m": stock.context_window} + + +def test_model_list_body_skips_a_limit_a_saved_model_info_response_copied_from_the_cost_map() -> None: + catalog_limit: Final = GetModelCostMap.loaded_model_cost_map()["gpt-6-astra"]["max_input_tokens"] + stock: Final = bundled_codex_models()["gpt-6-astra"] + assert isinstance(catalog_limit, int) and stock.max_context_window is not None + assert stock.context_window is not None and stock.context_window < 600000 < stock.max_context_window + + def window(model_info: Mapping[str, object]) -> int | None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra"}, + "model_info": dict(model_info), + } + ] + ) + listing: Final = ({"id": "gpt-6-astra", "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"},) + (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6-astra", "gpt-6-astra"),), router))["models"] + return entry["context_window"] + + assert window({"key": "gpt-6-astra", "max_input_tokens": catalog_limit}) == stock.context_window + assert window({"key": "gpt-6-astra", "max_input_tokens": 600000}) == 600000 + assert window({"max_input_tokens": catalog_limit}) == min(catalog_limit, stock.max_context_window) + + def test_alias_with_deployment_defaults_drops_codexs_tier_and_upgrade_while_the_native_name_keeps_them(): rows = (_row("gpt-6-sol"), _row("team-sol", upstream_model="gpt-6-sol"), _row("team-55", upstream_model="gpt-5.5")) codex = {entry["slug"]: entry for entry in _models(*rows)} @@ -599,9 +736,14 @@ def test_catalog_rows_read_the_router_by_each_entry_lookup_id(): "model_name_team-1_c0ffee" ).supported_reasoning_efforts, search_tool_support=(None, None), + configured_max_input_tokens=(None, None), ), CodexCatalogRow( - id="plain", upstream_model="openai/some-unmapped-model", service_tiers=(None,), search_tool_support=(None,) + id="plain", + upstream_model="openai/some-unmapped-model", + service_tiers=(None,), + search_tool_support=(None,), + configured_max_input_tokens=(None,), ), ) @@ -631,6 +773,7 @@ def test_catalog_rows_read_an_alias_off_its_target_under_the_alias_name(): service_tiers=(["ultrafast"],), supported_reasoning_efforts=router.get_model_group_info("gpt-6-astra").supported_reasoning_efforts, search_tool_support=(None,), + configured_max_input_tokens=(None,), ) (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6", "gpt-6"),), router))["models"] assert (entry["slug"], entry["display_name"], [tier["id"] for tier in entry["service_tiers"]]) == ( diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index 62189f5b0e7..84cb62f0426 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -786,6 +786,17 @@ async def test_get_form_data(): assert result["temperature"] == "0.2" assert result["include"] == ["logprobs"] assert result["timestamp_granularities"] == ["word", "segment"] + assert tuple(result) == ( + "file", + "model", + "include", + "language", + "prompt", + "response_format", + "stream", + "temperature", + "timestamp_granularities", + ) def test_get_tags_from_request_body_with_metadata_tags(): diff --git a/tests/unit/proxy/common_utils/test_model_deprecation.py b/tests/unit/proxy/common_utils/test_model_deprecation.py index 4f9f947fadb..7e96ad591b7 100644 --- a/tests/unit/proxy/common_utils/test_model_deprecation.py +++ b/tests/unit/proxy/common_utils/test_model_deprecation.py @@ -5,14 +5,17 @@ Slack integration — so they can run without the full proxy stack. """ import logging +from collections.abc import Mapping from datetime import date, datetime, timezone +from typing import Final from unittest.mock import MagicMock - +import pytest import litellm from litellm import Router from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry from litellm.proxy.common_utils.model_deprecation import ( _classify, _parse_deprecation_date, @@ -37,12 +40,7 @@ class TestParseDeprecationDate: def test_should_return_none_for_documentation_sentinel(self): # The JSON map ships a sentinel string under the "sample_spec" key. - assert ( - _parse_deprecation_date( - "date when the model becomes deprecated in the format YYYY-MM-DD" - ) - is None - ) + assert _parse_deprecation_date("date when the model becomes deprecated in the format YYYY-MM-DD") is None def test_should_return_none_for_none(self): assert _parse_deprecation_date(None) is None @@ -51,9 +49,7 @@ class TestParseDeprecationDate: assert _parse_deprecation_date(12345) is None def test_should_narrow_datetime_to_date(self): - assert _parse_deprecation_date( - datetime(2026, 12, 31, 23, 59, tzinfo=timezone.utc) - ) == date(2026, 12, 31) + assert _parse_deprecation_date(datetime(2026, 12, 31, 23, 59, tzinfo=timezone.utc)) == date(2026, 12, 31) class TestClassify: @@ -134,9 +130,7 @@ class TestCollectModelDeprecations: ] ) - snapshot = collect_model_deprecations( - llm_router=router, warn_within_days=30, today=today - ) + snapshot = collect_model_deprecations(llm_router=router, warn_within_days=30, today=today) assert [m.model_name for m in snapshot.deprecated] == ["deprecated-alias"] assert [m.model_name for m in snapshot.imminent] == ["imminent-alias"] @@ -166,9 +160,7 @@ class TestCollectModelDeprecations: ] ) - snapshot = collect_model_deprecations( - llm_router=router, warn_within_days=30, today=today - ) + snapshot = collect_model_deprecations(llm_router=router, warn_within_days=30, today=today) assert len(snapshot.imminent) == 1 assert snapshot.imminent[0].deprecation_date == date(2026, 6, 10) @@ -195,9 +187,7 @@ class TestCollectModelDeprecations: ] ) - snapshot = collect_model_deprecations( - llm_router=router, warn_within_days=30, today=today - ) + snapshot = collect_model_deprecations(llm_router=router, warn_within_days=30, today=today) assert len(snapshot.imminent) == 1 @@ -217,9 +207,7 @@ class TestCollectModelDeprecations: ] ) - snapshot = collect_model_deprecations( - llm_router=router, warn_within_days=30, today=date(2026, 6, 1) - ) + snapshot = collect_model_deprecations(llm_router=router, warn_within_days=30, today=date(2026, 6, 1)) assert [m.litellm_model for m in snapshot.imminent] == ["gpt-4o"] @@ -244,9 +232,7 @@ class TestCollectModelDeprecations: ] ) - snapshot = collect_model_deprecations( - llm_router=router, warn_within_days=30, today=date(2026, 6, 1) - ) + snapshot = collect_model_deprecations(llm_router=router, warn_within_days=30, today=date(2026, 6, 1)) assert len(snapshot.imminent) == 1 assert len(snapshot.upcoming) == 1 @@ -268,14 +254,136 @@ class TestCollectModelDeprecations: ] ) - snapshot = collect_model_deprecations( - llm_router=router, warn_within_days=30, today=today - ) + snapshot = collect_model_deprecations(llm_router=router, warn_within_days=30, today=today) assert len(snapshot.imminent) == 1 assert snapshot.imminent[0].litellm_model == "base-thing" +class TestProviderScopedResolution: + CATALOG = { + "gpt-4o": {"deprecation_date": "2027-01-14", "litellm_provider": "openai"}, + "azure/gpt-4o": {"deprecation_date": "2027-04-14", "litellm_provider": "azure"}, + "github_copilot/gpt-4o": {"litellm_provider": "github_copilot"}, + "azure_ai/gpt-4.1": AliasedCostEntry( + {"deprecation_date": "2027-05-01", "litellm_provider": "azure_ai", "mode": "chat"}, () + ), + "custom-base": {"deprecation_date": "2026-12-01"}, + "gpt-4.1": {"deprecation_date": "2027-01-14", "litellm_provider": "openai"}, + "azure/gpt-4.1": {"deprecation_date": "2027-04-14", "litellm_provider": "azure"}, + "gpt-4o-2024-05-13": {"deprecation_date": "2026-10-23", "litellm_provider": "openai"}, + "azure/gpt-4o-2024-05-13": {"deprecation_date": "2026-12-09", "litellm_provider": "azure"}, + "azure/aliased-base": AliasedCostEntry( + {"deprecation_date": "2026-11-20"}, ("azure/aliased-base", "azure/aliased-base-v1") + ), + "aliased-base": {"deprecation_date": "2027-02-01"}, + } + COST_MAP = { + **CATALOG, + "azure_ai/gpt-4o": AliasedCostEntry( + {"deprecation_date": "2027-01-14", "litellm_provider": "azure", "mode": "chat"}, () + ), + "registered-base": AliasedCostEntry({"deprecation_date": "2026-12-15"}, (), frozenset({"deprecation_date"})), + } + + @pytest.mark.parametrize( + ("litellm_params", "model_info", "expected"), + [ + ({"model": "gpt-4o"}, {}, ("gpt-4o", date(2027, 1, 14), "openai")), + ({"model": "openai/gpt-4o"}, {}, ("gpt-4o", date(2027, 1, 14), "openai")), + ( + {"model": "gpt-4o", "custom_llm_provider": "azure"}, + {}, + ("azure/gpt-4o", date(2027, 4, 14), "azure"), + ), + ({"model": "azure/gpt-4o"}, {}, ("azure/gpt-4o", date(2027, 4, 14), "azure")), + ( + {"model": "azure/prod-chat"}, + {"base_model": "gpt-4o"}, + ("azure/gpt-4o", date(2027, 4, 14), "azure"), + ), + ( + {"model": "gpt-4o", "custom_llm_provider": "github_copilot"}, + {}, + ("gpt-4o", date(2027, 1, 14), "openai"), + ), + ({"model": "azure_ai/gpt-4o"}, {}, ("azure/gpt-4o", date(2027, 4, 14), "azure")), + ( + {"model": "azure_ai/prod-chat"}, + {"base_model": "gpt-4o"}, + ("azure/gpt-4o", date(2027, 4, 14), "azure"), + ), + ( + {"model": "azure_ai/gpt-4o"}, + {"base_model": "azure_ai/gpt-4o"}, + ("azure/gpt-4o", date(2027, 4, 14), "azure"), + ), + ( + {"model": "azure_ai/prod-chat"}, + {"base_model": "azure_ai/gpt-4o"}, + ("azure/gpt-4o", date(2027, 4, 14), "azure"), + ), + ({"model": "azure_ai/gpt-4.1"}, {}, ("azure_ai/gpt-4.1", date(2027, 5, 1), "azure_ai")), + ({"model": "azure/gpt-4o"}, {"base_model": "custom-base"}, ("custom-base", date(2026, 12, 1), None)), + ( + {"model": "azure/gpt-4o"}, + {"base_model": "registered-base"}, + ("registered-base", date(2026, 12, 15), None), + ), + ( + {"model": "azure/prod-chat"}, + {"base_model": "aliased-base"}, + ("azure/aliased-base", date(2026, 11, 20), None), + ), + ( + {"model": "openai/gpt-4o-2024-05-13"}, + {"base_model": "gpt-4o"}, + ("gpt-4o-2024-05-13", date(2026, 10, 23), "openai"), + ), + ( + {"model": "azure/gpt-4o-2024-05-13"}, + {"base_model": "gpt-4o"}, + ("azure/gpt-4o-2024-05-13", date(2026, 12, 9), "azure"), + ), + ], + ids=[ + "bare", + "openai-prefix", + "declared-azure", + "azure-prefix", + "azure-base-model", + "undated-own-row", + "azure-ai-reads-azure", + "azure-ai-base-model-reads-azure", + "azure-ai-copied-base-model-yields-to-azure", + "azure-ai-prefixed-base-model-on-custom-deployment-reads-azure", + "azure-ai-own-row-wins", + "custom-base-model-row-wins", + "registered-base-model-row-wins", + "aliased-base-model-row-wins", + "openai-snapshot-shuts-down-before-its-base-model", + "azure-snapshot-shuts-down-before-its-base-model", + ], + ) + def test_should_prefer_the_deployments_own_provider_date( + self, + monkeypatch: pytest.MonkeyPatch, + litellm_params: Mapping[str, str], + model_info: Mapping[str, str], + expected: tuple[str, date, str | None], + ) -> None: + monkeypatch.setattr(litellm, "model_cost", self.COST_MAP) + router: Final = _make_router( + [{"model_name": "alias", "litellm_params": litellm_params, "model_info": {"id": "1", **model_info}}] + ) + + snapshot: Final = collect_model_deprecations( + llm_router=router, warn_within_days=30, today=date(2026, 6, 1), loaded_catalog=lambda: self.CATALOG + ) + + assert [(m.litellm_model, m.deprecation_date, m.litellm_provider) for m in snapshot.upcoming] == [expected] + + class TestFormatDeprecationAlertMessage: def test_should_return_none_when_nothing_to_alert(self): snapshot = collect_model_deprecations(llm_router=None) @@ -321,9 +429,7 @@ class TestFormatDeprecationAlertMessage: ] ) - snapshot = collect_model_deprecations( - llm_router=router, warn_within_days=30, today=today - ) + snapshot = collect_model_deprecations(llm_router=router, warn_within_days=30, today=today) message = format_deprecation_alert_message(snapshot) assert message is not None @@ -350,9 +456,7 @@ class TestFormatDeprecationAlertMessage: ] ) - snapshot = collect_model_deprecations( - llm_router=router, warn_within_days=30, today=today - ) + snapshot = collect_model_deprecations(llm_router=router, warn_within_days=30, today=today) message = format_deprecation_alert_message(snapshot) assert message is not None diff --git a/tests/unit/proxy/common_utils/test_realtime_cache.py b/tests/unit/proxy/common_utils/test_realtime_cache.py index 6f5751fa923..893e39b309a 100644 --- a/tests/unit/proxy/common_utils/test_realtime_cache.py +++ b/tests/unit/proxy/common_utils/test_realtime_cache.py @@ -1,6 +1,9 @@ -from typing import Any, cast +import json +from collections.abc import Mapping +from typing import Any, Final, cast import pytest +from pydantic import TypeAdapter from litellm.proxy.common_utils.realtime_utils import realtime_request_body from litellm.proxy.proxy_server import _realtime_query_params_template @@ -40,6 +43,19 @@ def test_realtime_request_body_caches_each_model_separately(): assert gpt4o_body_first is not gpt4o_mini_body +@pytest.mark.parametrize("model", (None, "requested", 'model"with\\quotes', "model\nwith-newline")) +def test_realtime_request_body_preserves_model_names(model: str | None) -> None: + body: Final = TypeAdapter(Mapping[str, object]).validate_json(realtime_request_body(model)) + + assert body["model"] == model + + +def test_realtime_request_body_keeps_a_model_name_with_quotes_inside_its_own_field(): + body: Final = realtime_request_body('gpt-4o", "user": "someone-else') + + assert json.loads(body) == {"model": 'gpt-4o", "user": "someone-else'} + + def test_realtime_query_params_template_caches_each_pair_separately(): params_with_intent_first = _realtime_query_params_template("gpt-4o", "intent-a") params_with_intent_second = _realtime_query_params_template("gpt-4o", "intent-a") diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 43edc7e449f..68910f52b02 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -131,6 +131,107 @@ def test_shipped_australian_count_region_stays_with_an_australian_deployment() - assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4") +_ABSENT: Final = object() + + +@pytest.mark.parametrize( + ("key", "field"), + [ + ("token_count_tool_framing", "tool_framing"), + ("token_count_responses_framing", "responses_framing"), + ("token_count_responses_tool_framing", "responses_tool_framing"), + ], +) +@pytest.mark.parametrize( + ("configured", "default", "expected"), + [ + (_ABSENT, 151, 151), + (40, 151, 40), + (0, 151, 0), + (40.0, _ABSENT, 40), + (-1, 151, 151), + (True, 151, 151), + ("40", _ABSENT, None), + ], +) +def test_each_request_kind_framing_reads_its_value_like_framing( + key: str, field: str, configured: object, default: object, expected: int | None +) -> None: + settings: Final = resolve_token_count_settings( + configured={} if configured is _ABSENT else {key: configured}, + default={} if default is _ABSENT else {key: default}, + deployment_region=_in_us_west_2, + ) + assert getattr(settings, field) == expected + + +def _hidden_prompts(settings: TokenCountSettings) -> tuple[int, int, int, int]: + return ( + settings.hidden_prompt_tokens(carries_tools=False, responses_api=False), + settings.hidden_prompt_tokens(carries_tools=True, responses_api=False), + settings.hidden_prompt_tokens(carries_tools=False, responses_api=True), + settings.hidden_prompt_tokens(carries_tools=True, responses_api=True), + ) + + +@pytest.mark.parametrize( + ("default", "expected"), + [ + ({}, (0, 0, 0, 0)), + ({"token_count_framing": 23}, (23, 23, 23, 23)), + ({"token_count_framing": 185, "token_count_tool_framing": 214}, (185, 214, 185, 214)), + ( + { + "token_count_framing": 149, + "token_count_tool_framing": 300, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 104, + }, + (149, 300, 145, 104), + ), + ({"token_count_framing": 149, "token_count_tool_framing": 0}, (149, 0, 149, 0)), + ({"token_count_framing": 149, "token_count_responses_tool_framing": 0}, (149, 149, 149, 0)), + ({"token_count_framing": 149, "token_count_responses_framing": 145}, (149, 149, 145, 145)), + ], +) +def test_hidden_prompt_is_the_most_specific_framing_set_for_the_request_api_and_tools( + default: dict[str, object], expected: tuple[int, int, int, int] +) -> None: + settings: Final = resolve_token_count_settings(configured={}, default=default, deployment_region=_in_us_west_2) + assert _hidden_prompts(settings) == expected + + +_PER_API_COST_MAP: Final = { + "token_count_framing": 149, + "token_count_tool_framing": 300, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 104, +} + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({}, (149, 300, 145, 104)), + ({"token_count_framing": 500}, (500, 500, 500, 500)), + ({"token_count_framing": 0}, (0, 0, 0, 0)), + ({"token_count_tool_framing": 250}, (149, 250, 145, 250)), + ({"token_count_framing": 500, "token_count_tool_framing": 250}, (500, 250, 500, 250)), + ({"token_count_responses_framing": 120}, (149, 300, 120, 120)), + ({"token_count_tool_framing": 250, "token_count_responses_framing": 120}, (149, 250, 120, 120)), + ({"token_count_responses_tool_framing": 90}, (149, 300, 145, 90)), + ({"token_count_framing": -1}, (149, 300, 145, 104)), + ], +) +def test_an_operator_framing_applies_to_every_request_kind_the_cost_map_splits_more_finely( + configured: dict[str, object], expected: tuple[int, int, int, int] +) -> None: + settings: Final = resolve_token_count_settings( + configured=configured, default=_PER_API_COST_MAP, deployment_region=_in_us_west_2 + ) + assert _hidden_prompts(settings) == expected + + def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None: assert resolve_token_count_settings( configured={}, default={}, deployment_region=_in_us_west_2 diff --git a/tests/unit/proxy/common_utils/test_user_api_key_cache.py b/tests/unit/proxy/common_utils/test_user_api_key_cache.py index 09523cd1901..577e5e5a389 100644 --- a/tests/unit/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/unit/proxy/common_utils/test_user_api_key_cache.py @@ -1,6 +1,6 @@ import hashlib import json -from typing import Any +from typing import Any, Final import pytest @@ -11,6 +11,7 @@ from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + budget_windows_absent_cache_key, end_user_cache_key, get_management_object_ttl, is_user_key_cache_key, @@ -32,6 +33,17 @@ class CapturingInMemoryCache(InMemoryCache): super().set_cache(key, value, **kwargs) +class _ManualClock: + def __init__(self) -> None: + self.seconds = 1000.0 + + def __call__(self) -> float: + return self.seconds + + def advance(self, seconds: float) -> None: + self.seconds += seconds + + class FakeRedisCache(RedisCache): """ In-memory fake that enforces the UserApiKeyCache Redis payload contract. @@ -432,6 +444,29 @@ class TestUserKeyObjectPartition: assert first_key.token == tokens[0] assert cache.get_cache(end_user_cache_key("u0")) == {"user_id": "u0"} + def test_update_in_memory_max_size_applies_to_no_window_notes(self): + cache = UserApiKeyCache() + cache.update_in_memory_max_size(1) + + cache.budget_windows_absent_cache.set_cache(key=budget_windows_absent_cache_key("key", "first"), value=True) + cache.budget_windows_absent_cache.set_cache(key=budget_windows_absent_cache_key("key", "second"), value=True) + + assert cache.budget_windows_absent_cache.get_cache(key=budget_windows_absent_cache_key("key", "first")) is None + assert cache.budget_windows_absent_cache.get_cache(key=budget_windows_absent_cache_key("key", "second")) is True + + def test_update_cache_ttl_applies_to_no_window_notes(self): + clock: Final = _ManualClock() + cache: Final = UserApiKeyCache( + default_in_memory_ttl=60, budget_windows_absent_in_memory_cache=InMemoryCache(clock=clock) + ) + cache.update_cache_ttl(default_in_memory_ttl=5, default_redis_ttl=5) + note: Final = budget_windows_absent_cache_key("key", "no-windows") + + cache.budget_windows_absent_cache.set_cache(key=note, value=True) + clock.advance(6) + + assert cache.budget_windows_absent_cache.get_cache(key=note) is None + def test_update_in_memory_max_size_none_resets_key_object_partition_to_default(self): cache = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=1)) cache.update_in_memory_max_size(None) diff --git a/tests/unit/proxy/config_resolvers/test_settings_store.py b/tests/unit/proxy/config_resolvers/test_settings_store.py index 7b2cd404b46..76850825937 100644 --- a/tests/unit/proxy/config_resolvers/test_settings_store.py +++ b/tests/unit/proxy/config_resolvers/test_settings_store.py @@ -471,3 +471,34 @@ def test_settings_store_truthiness_matches_emptiness() -> None: assert bool(store) is True del store["max_parallel_requests"] assert bool(store) is False + + +def test_settings_store_ignores_stored_keys_outside_its_managed_db_keys_and_still_applies_them_from_yaml() -> None: + store: Final = SettingsStore("general_settings", managed_db_keys=frozenset({"max_parallel_requests"})) + stored_row: Final[dict[str, JsonValue]] = {"enforce_user_param": True, "max_parallel_requests": 5} + + store.apply_db_row("general_settings", stored_row) + + assert dict(store) == {"max_parallel_requests": 5} + assert store.source("enforce_user_param") == "unset" + assert store.unmanaged_db_keys() == ("enforce_user_param",) + + store.load_yaml({"enforce_user_param": True}) + store.apply_db_row("general_settings", stored_row) + + assert dict(store) == {"enforce_user_param": True, "max_parallel_requests": 5} + assert store.source("enforce_user_param") == "config" + assert store.unmanaged_db_keys() == ("enforce_user_param",) + + +def test_settings_store_managed_db_keys_leave_other_rows_and_unrestricted_stores_whole() -> None: + restricted: Final = SettingsStore("general_settings", managed_db_keys=frozenset()) + unrestricted: Final = SettingsStore("general_settings") + + restricted.apply_db_row("ui_settings", {"disable_custom_api_keys": True}) + unrestricted.apply_db_row("general_settings", {"enforce_user_param": True}) + + assert dict(restricted) == {"disable_custom_api_keys": True} + assert restricted.unmanaged_db_keys() == () + assert dict(unrestricted) == {"enforce_user_param": True} + assert unrestricted.unmanaged_db_keys() == () diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 3da77a8707e..2603cf92f30 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1694,6 +1694,48 @@ class TestResponsesStreamRestoration: assert "".join(e.delta for e in out) == "see a@example.com a[" + @pytest.mark.asyncio + async def test_a_flushed_delta_takes_the_done_events_number_and_later_events_move_up(self): + guardrail, _ = _shielded(self.VAULT) + done = OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id="msg_1", + output_index=0, + content_index=0, + text="Mail [EMAIL_1] x[", + sequence_number=3, + ) + + out = await _restore_stream( + guardrail, + [ + self._text_delta("Mail [EMA", 1), + self._text_delta("IL_1] x[", 2), + done, + self._text_delta("next", 4, content_index=1), + ], + ) + + assert [(e.type, e.sequence_number) for e in out] == [ + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 1), + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 2), + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 3), + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, 4), + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 5), + ] + assert out[2].delta == "[" + + @pytest.mark.asyncio + async def test_a_truncated_stream_flush_continues_after_the_last_number(self): + guardrail, _ = _shielded(self.VAULT) + + out = await _restore_stream( + guardrail, [self._text_delta("see [EMAIL_1] a[", 7), self._text_delta(" b", 8, content_index=1)] + ) + + assert [e.sequence_number for e in out] == [7, 8, 9] + assert out[-1].delta == "[" + @pytest.mark.asyncio async def test_completed_response_is_restored(self): """The terminal event repeats the whole reply, and clients read it as the answer.""" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py index 0e6670c26e7..43238e37f63 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -6,19 +6,24 @@ and allows requests with only registered servers. Covers both /chat/completions and /responses API paths (same pre_call_hook logic, different call_type). """ -from typing import Literal -from unittest.mock import MagicMock, patch +from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException import litellm +from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.mcp_security import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( MCPSecurityGuardrail, ) +from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.types.mcp_server.mcp_toolset import MCPToolset @pytest.fixture @@ -60,93 +65,67 @@ class TestExtractMCPServerNames: assert names == set() +def _manager_with_saved_server() -> MCPServerManager: + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="3f6b2c1e-saved-zapier", + name="zapier", + server_name="zapier", + alias="zap", + transport=MCPTransport.http, + url="https://upstream.example/mcp", + ) + manager.registry[server.server_id] = server + return manager + + class TestMCPSecurityGuardrailPreCall: @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_blocks_unregistered_server_chat_completions( - self, mock_manager, guardrail - ): - """Simulates /chat/completions with an unregistered MCP server.""" - mock_manager.get_registry.return_value = {"zapier": MagicMock()} - + @pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) + async def test_blocks_unregistered_server(self, guardrail, call_type): data = { "tools": [ {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, {"type": "mcp", "server_url": "litellm_proxy/mcp/evil_server"}, ], "model": "gpt-4", - "messages": [{"role": "user", "content": "hi"}], "guardrails": ["test-mcp-security"], } - with pytest.raises(HTTPException) as exc_info: + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", _manager_with_saved_server()), + pytest.raises(HTTPException) as exc_info, + ): await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(), cache=MagicMock(), data=data, - call_type="acompletion", + call_type=call_type, ) assert exc_info.value.status_code == 400 - assert "evil_server" in str(exc_info.value.detail) + assert exc_info.value.detail["unregistered_servers"] == ["evil_server"] @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_blocks_unregistered_server_responses_api( - self, mock_manager, guardrail - ): - """Simulates /responses with an unregistered MCP server.""" - mock_manager.get_registry.return_value = {"github": MagicMock()} - + @pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) + @pytest.mark.parametrize("reference", ["zapier", "zap", "3f6b2c1e-saved-zapier"]) + async def test_allows_saved_server_by_any_name_the_gateway_routes(self, guardrail, call_type, reference): data = { - "tools": [ - {"type": "mcp", "server_url": "litellm_proxy/mcp/unknown_server"}, - ], - "model": "gpt-4o", - "input": "What can you do?", - "guardrails": ["test-mcp-security"], - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(), - cache=MagicMock(), - data=data, - call_type="aresponses", - ) - assert exc_info.value.status_code == 400 - assert "unknown_server" in str(exc_info.value.detail) - - @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_allows_registered_servers(self, mock_manager, guardrail): - """All MCP servers are registered - request passes through.""" - mock_manager.get_registry.return_value = { - "zapier": MagicMock(), - "github": MagicMock(), - } - - data = { - "tools": [ - {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, - {"type": "mcp", "server_url": "litellm_proxy/mcp/github"}, - ], + "tools": [{"type": "mcp", "server_url": f"litellm_proxy/mcp/{reference}"}], "model": "gpt-4", - "messages": [{"role": "user", "content": "hi"}], "guardrails": ["test-mcp-security"], } - result = await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(), - cache=MagicMock(), - data=data, - call_type="acompletion", - ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", _manager_with_saved_server()), + ): + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=data, + call_type=call_type, + ) assert result == data @pytest.mark.asyncio @@ -222,11 +201,13 @@ async def test_guardrail_observes_saved_server_creation_and_deletion_on_another_ prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [])) prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcptoolsettable.find_first = AsyncMock(return_value=None) data = {"tools": [{"type": "mcp", "server_url": "litellm_proxy/mcp/peer-server"}], "guardrails": ["test-mcp-security"]} with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", manager), ): result = await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) @@ -235,3 +216,34 @@ async def test_guardrail_observes_saved_server_creation_and_deletion_on_another_ await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) assert exc.value.status_code == 400 assert exc.value.detail["unregistered_servers"] == ["peer-server"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) +async def test_allows_saved_toolset_by_name_and_blocks_unknown_name(guardrail, call_type): + saved_toolset: Final = MCPToolset(toolset_id="7c1d-research", toolset_name="research") + + async def find_toolset(*, where: dict[str, str]) -> MCPToolset | None: + return saved_toolset if where["toolset_name"] == saved_toolset.toolset_name else None + + prisma: Final = MagicMock() + prisma.writer_db = prisma.db + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcptoolsettable.find_first = AsyncMock(side_effect=find_toolset) + data: Final = { + "tools": [ + {"type": "mcp", "server_url": "litellm_proxy/mcp/research"}, + {"type": "mcp", "server_url": "litellm_proxy/mcp/evil_server"}, + ], + "guardrails": ["test-mcp-security"], + } + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", MCPServerManager()), + pytest.raises(HTTPException) as exc, + ): + await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) + assert exc.value.detail["unregistered_servers"] == ["evil_server"] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py index d2e5b407e30..1baf9504208 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py @@ -200,3 +200,78 @@ async def test_write_back_keeps_real_tool_results_under_modify_params(): assert "skipped" not in serialized assert "Please continue" not in serialized assert [m["role"] for m in written] == ["user", "assistant", "user", "user"] + + +@pytest.mark.asyncio +async def test_write_back_sends_the_client_url_blocks_as_the_client_sent_them(async_only_image_fetch): + image_link = "https://img.example/presskit/page.pdf?fm=png" + document_link = "https://files.example/reports/q3.pdf" + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu_1", "name": "fetch_report", "input": {}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [ + {"type": "document", "source": {"type": "url", "url": document_link}, "title": "Q3 attachment"} + ], + }, + {"type": "text", "text": "compare these"}, + {"type": "image", "source": {"type": "url", "url": image_link}}, + { + "type": "document", + "source": {"type": "url", "url": document_link}, + "title": "Q3 report", + "context": "Board copy", + "citations": {"enabled": True}, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "document", "source": {"type": "url", "url": document_link}, "title": "Q3 report, again"}, + ], + }, + ] + + written = await _write_back_identity([dict(m) for m in messages]) + + assert async_only_image_fetch.fetched == [] + assert [m["content"] for m in written[1:]] == [messages[1]["content"][:1], messages[1]["content"][1:]] + + +@pytest.mark.asyncio +async def test_write_back_keeps_each_url_document_fields_when_a_tool_result_after_it_has_the_same_link( + async_only_image_fetch, +): + document_link = "https://files.example/reports/q3.pdf" + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu_1", "name": "fetch_report", "input": {}}], + }, + { + "role": "user", + "content": [ + { + "type": "document", + "source": {"type": "url", "url": document_link}, + "title": "Q3 report", + "context": "Board copy", + }, + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [ + {"type": "document", "source": {"type": "url", "url": document_link}, "title": "Q3 attachment"} + ], + }, + ], + }, + ] + + written = await _write_back_identity([dict(m) for m in messages]) + + assert [m["content"] for m in written[1:]] == [messages[1]["content"][1:], messages[1]["content"][:1]] 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 bed185ba8f5..f9c234be55a 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3711,6 +3711,255 @@ async def test_failure_event_refunds_reservation_for_input_only_estimate_v3(): assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_first", [True, False]) +@pytest.mark.parametrize( + ("response_obj", "used_tokens"), + [ + ([], 0), + (ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=10, total_tokens=30)), 30), + ], +) +async def test_a_request_logging_a_failure_and_a_success_settles_its_reservation_once_v3( + failure_first: bool, response_obj: object, used_tokens: int +) -> None: + api_key: Final = hash_token("sk-failure-and-success") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + tokens_key: Final = handler.create_rate_limit_keys(key="api_key", value=api_key, rate_limit_type="tokens") + await _reserve_tpm_for_owner_call(handler, local_cache, api_key, "dropped-call") + kwargs: Final = { + "litellm_call_id": "dropped-call", + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + } + now: Final = datetime.now() + log_failure: Final = partial( + handler.async_log_failure_event, kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now + ) + log_success: Final = partial( + handler.async_log_success_event, kwargs=dict(kwargs), response_obj=response_obj, start_time=now, end_time=now + ) + + for log_event in (log_failure, log_success) if failure_first else (log_success, log_failure): + await log_event() + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == used_tokens + + +class _YieldingIncrementInMemoryCache(InMemoryCache): + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + await asyncio.sleep(0) + return await super().async_increment(key, value, **kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("success_starts_first", [True, False]) +@pytest.mark.parametrize("failure_path", ["failure_event", "post_call_failure_hook"]) +async def test_concurrent_failure_and_success_callbacks_settle_a_reservation_once_v3( + success_starts_first: bool, failure_path: str +) -> None: + api_key: Final = hash_token("sk-concurrent-failure-and-success") + model: Final = "gpt-4o-mini" + project: Final = "proj-concurrent" + local_cache: Final = DualCache(in_memory_cache=_YieldingIncrementInMemoryCache()) + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + caller: Final = UserAPIKeyAuth( + api_key=api_key, + tpm_limit=10_000, + project_id=project, + project_metadata={"model_itpm_limit": {model: 10_000}, "model_otpm_limit": {model: 10_000}}, + ) + counter_keys: Final = ( + handler.create_rate_limit_keys(key="api_key", value=api_key, rate_limit_type="tokens"), + handler.create_rate_limit_keys( + key="model_per_project_itpm", value=f"{project}:{model}", rate_limit_type="tokens" + ), + handler.create_rate_limit_keys( + key="model_per_project_otpm", value=f"{project}:{model}", rate_limit_type="tokens" + ), + ) + await handler.async_pre_call_hook( + user_api_key_dict=caller, + cache=local_cache, + data={ + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + "litellm_call_id": "concurrent-call", + }, + call_type="completion", + ) + kwargs: Final = { + "litellm_call_id": "concurrent-call", + "model": model, + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + } + now: Final = datetime.now() + log_failure: Final = ( + partial(handler.async_log_failure_event, kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now) + if failure_path == "failure_event" + else partial( + handler.async_post_call_failure_hook, + request_data=dict(kwargs), + original_exception=Exception("stream ended"), + user_api_key_dict=caller, + ) + ) + log_success: Final = partial( + handler.async_log_success_event, + kwargs=dict(kwargs), + response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=10, total_tokens=30)), + start_time=now, + end_time=now, + ) + + await asyncio.gather( + *((log_success(), log_failure()) if success_starts_first else (log_failure(), log_success())) + ) + + assert tuple([int(await local_cache.async_get_cache(key=key) or 0) for key in counter_keys]) == (30, 20, 10) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("second_call_type", "second_model"), + [("call_mcp_tool", "mcp-tool-call"), ("completion", "gpt-4o-mini")], +) +async def test_a_pre_call_after_a_failure_leaves_the_success_charging_project_io_once_v3( + second_call_type: str, second_model: str +) -> None: + api_key: Final = hash_token("sk-retried-request-with-mcp-tool") + model: Final = "gpt-4o-mini" + project: Final = "proj-retried-mcp" + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + caller: Final = UserAPIKeyAuth( + api_key=api_key, + tpm_limit=100_000, + project_id=project, + project_metadata={"model_itpm_limit": {model: 10_000}, "model_otpm_limit": {model: 10_000}}, + ) + request: Final = { + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + "litellm_call_id": "retried-call-with-mcp-tool", + } + kwargs: Final = { + "litellm_call_id": "retried-call-with-mcp-tool", + "model": model, + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + } + now: Final = datetime.now() + itpm_key: Final = handler.create_rate_limit_keys( + key="model_per_project_itpm", value=f"{project}:{model}", rate_limit_type="tokens" + ) + otpm_key: Final = handler.create_rate_limit_keys( + key="model_per_project_otpm", value=f"{project}:{model}", rate_limit_type="tokens" + ) + + await handler.async_pre_call_hook( + user_api_key_dict=caller, cache=local_cache, data=dict(request), call_type="completion" + ) + await handler.async_log_failure_event(kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now) + await handler.async_pre_call_hook( + user_api_key_dict=caller, + cache=local_cache, + data={**request, "model": second_model}, + call_type=second_call_type, + ) + await handler.async_log_success_event( + kwargs=dict(kwargs), + response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27)), + start_time=now, + end_time=now, + ) + + assert ( + int(await local_cache.async_get_cache(key=itpm_key) or 0), + int(await local_cache.async_get_cache(key=otpm_key) or 0), + ) == (20, 7) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("steps", "key_tpm_limit", "expected_counters"), + [ + (("failure", "mcp-tool-call", "failure"), 100_000, (0, 0, 0, 0)), + (("failure", "tool-model-with-io-limit", "post_call_failure"), None, (0, 0, 0, 0)), + (("failure", "mcp-tool-call", "post_call_failure", "success"), 100_000, (20, 7, 0, 0)), + (("tool-model-over-io-limit", "success"), 100_000, (20, 7, 0, 0)), + ], +) +async def test_a_nested_mcp_pre_call_settles_each_project_io_hold_once_v3( + steps: tuple[str, ...], key_tpm_limit: Optional[int], expected_counters: tuple[int, int, int, int] +) -> None: + api_key: Final = hash_token("sk-nested-mcp-hold-flags") + model: Final = "gpt-4o-mini" + tool_model: Final = "tool-model-with-io-limit" + project: Final = "proj-nested-mcp-hold-flags" + io_limits: Final = {model: 10_000, tool_model: 10_000, "tool-model-over-io-limit": 1} + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + caller: Final = UserAPIKeyAuth( + api_key=api_key, + tpm_limit=key_tpm_limit, + project_id=project, + project_metadata={"model_itpm_limit": io_limits, "model_otpm_limit": io_limits}, + ) + request: Final = { + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + "litellm_call_id": "nested-mcp-hold-flags", + } + kwargs: Final = { + "litellm_call_id": "nested-mcp-hold-flags", + "model": model, + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + } + now: Final = datetime.now() + counter_keys: Final = tuple( + handler.create_rate_limit_keys( + key=f"model_per_project_{kind}", value=f"{project}:{scope}", rate_limit_type="tokens" + ) + for scope, kind in ((model, "itpm"), (model, "otpm"), (tool_model, "itpm"), (tool_model, "otpm")) + ) + mcp_pre_call: Final = partial( + handler.async_pre_call_hook, user_api_key_dict=caller, cache=local_cache, call_type="call_mcp_tool" + ) + + await handler.async_pre_call_hook( + user_api_key_dict=caller, cache=local_cache, data=dict(request), call_type="completion" + ) + for step in steps: + match step: + case "failure": + await handler.async_log_failure_event( + kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now + ) + case "post_call_failure": + await handler.async_post_call_failure_hook( + request_data=dict(request), + original_exception=HTTPException(status_code=429, detail="upstream rate limited"), + user_api_key_dict=caller, + ) + case "success": + await handler.async_log_success_event( + kwargs=dict(kwargs), + response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27)), + start_time=now, + end_time=now, + ) + case "tool-model-over-io-limit": + with pytest.raises(ProxyRateLimitError): + await mcp_pre_call(data={**request, "model": step}) + case _: + await mcp_pre_call(data={**request, "model": step}) + + assert tuple([int(await local_cache.async_get_cache(key=key) or 0) for key in counter_keys]) == expected_counters + + @pytest.mark.asyncio async def test_post_call_failure_hook_settles_reservation_at_recovered_partial_usage_v3(): """ @@ -4198,6 +4447,125 @@ async def test_failure_event_settles_project_itpm_otpm_at_recovered_partial_usag assert int(await local_cache.async_get_cache(key=otpm_key) or 0) == 7 + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_first", [True, False]) +async def test_a_request_logging_a_failure_and_a_success_settles_its_project_itpm_otpm_once_v3( + failure_first: bool, +) -> None: + api_key: Final = hash_token("sk-failure-and-success-project-io") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key=api_key, + project_id="proj-dropped", + project_metadata={ + "model_itpm_limit": {"gpt-4o-mini": 10_000}, + "model_otpm_limit": {"gpt-4o-mini": 10_000}, + }, + ), + cache=local_cache, + data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + "litellm_call_id": "dropped-project-call", + }, + call_type="completion", + ) + itpm_key: Final = handler.create_rate_limit_keys( + key="model_per_project_itpm", value="proj-dropped:gpt-4o-mini", rate_limit_type="tokens" + ) + otpm_key: Final = handler.create_rate_limit_keys( + key="model_per_project_otpm", value="proj-dropped:gpt-4o-mini", rate_limit_type="tokens" + ) + kwargs: Final = { + "litellm_call_id": "dropped-project-call", + "model": "gpt-4o-mini", + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + } + now: Final = datetime.now() + log_failure: Final = partial( + handler.async_log_failure_event, kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now + ) + log_success: Final = partial( + handler.async_log_success_event, + kwargs=dict(kwargs), + response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27)), + start_time=now, + end_time=now, + ) + + for log_event in (log_failure, log_success) if failure_first else (log_success, log_failure): + await log_event() + + assert int(await local_cache.async_get_cache(key=itpm_key) or 0) == 20 + assert int(await local_cache.async_get_cache(key=otpm_key) or 0) == 7 + + +@pytest.mark.asyncio +async def test_a_reservation_made_after_a_released_one_in_the_same_context_settles_at_its_usage_v3() -> None: + api_key: Final = hash_token("sk-second-reservation") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict: Final = UserAPIKeyAuth( + api_key=api_key, + tpm_limit=10_000, + project_id="proj-second", + project_metadata={ + "model_itpm_limit": {"gpt-4o-mini": 10_000}, + "model_otpm_limit": {"gpt-4o-mini": 10_000}, + }, + ) + now: Final = datetime.now() + + async def reserve(call_id: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + "litellm_call_id": call_id, + }, + call_type="completion", + ) + + def logged(call_id: str) -> dict[str, object]: + return { + "litellm_call_id": call_id, + "model": "gpt-4o-mini", + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + } + + await reserve("first-call") + await handler.async_log_failure_event(kwargs=logged("first-call"), response_obj=None, start_time=now, end_time=now) + await reserve("second-call") + await handler.async_log_success_event( + kwargs=logged("second-call"), + response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27)), + start_time=now, + end_time=now, + ) + + counters: Final = { + "tokens": ("api_key", api_key), + "itpm": ("model_per_project_itpm", "proj-second:gpt-4o-mini"), + "otpm": ("model_per_project_otpm", "proj-second:gpt-4o-mini"), + } + settled: Final = { + name: int( + await local_cache.async_get_cache( + key=handler.create_rate_limit_keys(key=key, value=value, rate_limit_type="tokens") + ) + or 0 + ) + for name, (key, value) in counters.items() + } + assert settled == {"tokens": 27, "itpm": 20, "otpm": 7} + # ----------------------- Per-MCP-server rate limiting (v3) ----------------------- diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index a4e78a48a1a..21b5bd300c1 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -3491,6 +3491,40 @@ async def test_sidecar_ignores_an_undecodable_event(): # test-quality-ok: a dis mock_proxy_logging.db_spend_update_writer.update_database.assert_not_awaited() +def _spend_row_error_information(update_database_kwargs: Mapping[str, object]) -> dict[str, object]: + row: Final = get_logging_payload( + kwargs=cast(dict, update_database_kwargs["kwargs"]), + response_obj=update_database_kwargs["completion_response"], + start_time=cast(datetime, update_database_kwargs["start_time"]), + end_time=cast(datetime, update_database_kwargs["end_time"]), + ) + return json.loads(row["metadata"])["error_information"] + + +@pytest.mark.asyncio +async def test_failure_hook_spend_row_truncates_a_long_provider_error_exactly_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("MAX_STRING_LENGTH_PROMPT_IN_DB", "100") + writer: Final = MagicMock(spec=DBSpendUpdateWriter) + writer.update_database = AsyncMock() + logger: Final = ProxyDBLogger(spend_writer=lambda: writer) + provider_error: Final = "upstream failure " + "x" * 1000 + + await logger.async_post_call_failure_hook( + request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + original_exception=litellm.InternalServerError(message=provider_error, llm_provider="openai", model="gpt-4o"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + stored_message: Final = _spend_row_error_information(writer.update_database.call_args.kwargs)["error_message"] + assert isinstance(stored_message, str) + assert stored_message.count("skipped") == 1 + kept_chars: Final = len(stored_message.split("... (")[0]) + len(stored_message.rsplit(") ...", 1)[1]) + skipped_chars: Final = int(stored_message.split("skipped ", 1)[1].split(" chars", 1)[0]) + assert kept_chars + skipped_chars > len(provider_error) + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_model_rejection(): raw_model: Final = "opus-4.6 Please summarize my medical records\nPatient has diabetes" @@ -3504,9 +3538,7 @@ async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_ user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), ) - error_information: Final = writer.update_database.call_args.kwargs["kwargs"]["litellm_params"]["metadata"][ - "error_information" - ] + error_information: Final = _spend_row_error_information(writer.update_database.call_args.kwargs) assert "medical records" not in json.dumps(error_information) assert ( error_information["error_message"] diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index e43a72541c7..bfb7a9d6199 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -3060,7 +3060,7 @@ def test_retry_success_charges_released_project_io_reservations( stash.itpm_reserved_scopes = frozenset({itpm_scope}) stash.otpm_reserved_tokens = 60 stash.otpm_reserved_scopes = frozenset({otpm_scope}) - stash.reservation_released = True + stash.mark_reservation_released() operations = handler._build_io_token_reservation_ops( {}, diff --git a/tests/unit/proxy/image_endpoints/test_endpoints.py b/tests/unit/proxy/image_endpoints/test_endpoints.py index ea37b7b0213..f153a0f6990 100644 --- a/tests/unit/proxy/image_endpoints/test_endpoints.py +++ b/tests/unit/proxy/image_endpoints/test_endpoints.py @@ -715,3 +715,29 @@ def test_image_generation_uses_key_num_retries(monkeypatch: pytest.MonkeyPatch) router: Final = proxy_server.llm_router assert router is not None assert get_deployment_failures_for_current_minute(router, "image-regression-failing") == 3 + + +def test_failed_image_generation_spend_row_records_aimage_generation(monkeypatch: pytest.MonkeyPatch) -> None: + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + spend_writer: Final = MagicMock(update_database=AsyncMock()) + monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger(spend_writer=lambda: spend_writer)]) + client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth(api_key="sk-test")) + try: + with client: + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + assert (response.status_code, handler.calls) == (429, ["failing"]), response.text + db_call: Final = spend_writer.update_database.call_args.kwargs + payload: Final = get_logging_payload( + kwargs=db_call["kwargs"], + response_obj=db_call["completion_response"], + start_time=db_call["start_time"], + end_time=db_call["end_time"], + ) + assert (payload["status"], payload["call_type"]) == ("failure", "aimage_generation") diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 0198f582be7..b453f985512 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -6,7 +6,7 @@ from contextlib import ExitStack from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch from uuid import UUID, uuid4 import pytest @@ -51,7 +51,11 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + budget_windows_absent_cache_key, + project_cache_key, +) from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.management_endpoints.key_management_endpoints import ( KeyProjectBindingError, @@ -8911,7 +8915,10 @@ async def test_delete_cache_key_object_broadcasts_invalidation(monkeypatch): # Real, observable state: the cache object itself no longer holds the entry. assert real_user_api_key_cache.get_cache(key="hashed-broadcast-key") is None - mock_publish.assert_awaited_once_with(cache_key="hashed-broadcast-key") + assert mock_publish.await_args_list == [ + call(cache_key="hashed-broadcast-key"), + call(cache_key=budget_windows_absent_cache_key("key", "hashed-broadcast-key")), + ] @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index ab6cfcfebfe..44c43713438 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4896,6 +4896,9 @@ class TestModelInfoCostMapEchoFilter: [ ("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}), ("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_tool_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_responses_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_responses_tool_framing", 0), ], ) def test_token_count_setting_is_stored_only_when_the_operator_changed_it( diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index e33743bb149..2f4dad0ccd7 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -2684,7 +2684,7 @@ async def test_team_block_unblock_reach_warmed_auth_on_every_worker( from litellm.proxy._types import BlockTeamRequest from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.team_grants import team_model_aliases - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, budget_windows_absent_cache_key from litellm.proxy.management_endpoints import team_endpoints team_id: Final = "team-kill-switch" @@ -2721,7 +2721,11 @@ async def test_team_block_unblock_reach_warmed_auth_on_every_worker( seen_by_auth: Final = await get_team_object(team_id=team_id, prisma_client=prisma_client, user_api_key_cache=cache) assert seen_by_auth.blocked is blocked assert team_model_aliases(seen_by_auth) == {"fast": "gpt-4o"} - assert redis.broadcast_keys == [f"team_id:{team_id}", "team_alias:kill-switch"] + assert redis.broadcast_keys == [ + f"team_id:{team_id}", + "team_alias:kill-switch", + budget_windows_absent_cache_key("team", team_id), + ] assert (response.blocked, response.litellm_model_table, response.object_permission) == (blocked, None, None) diff --git a/tests/unit/proxy/mcp_server/test_mcp_server.py b/tests/unit/proxy/mcp_server/test_mcp_server.py index 884c82760ea..456b1994e2d 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server.py @@ -3319,3 +3319,35 @@ async def test_mcp_routing_truncated_json_body_is_not_stashed_for_auth(): assert await scope[MCP_PEEKED_BODY_SCOPE_KEY]() == b"{}" assert await _read_request_body(_admission_request(scope)) == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/mcp", "/atlassian/mcp"]) +async def test_bare_options_skips_litellm_auth(path: str): + """A bare OPTIONS probe must not run LiteLLM auth, so it never triggers the 401 OAuth challenge.""" + from litellm.proxy.mcp_server import server as mcp_server + + scope: Scope = { + "type": "http", + "method": "OPTIONS", + "path": path, + "headers": [], + "query_string": b"", + "server": ("localhost", 8000), + "scheme": "http", + } + receive = AsyncMock(return_value={"type": "http.request", "body": b"", "more_body": False}) + sent: list[dict] = [] + + async def send(message: dict) -> None: + sent.append(message) + + auth = AsyncMock(side_effect=AssertionError("LiteLLM auth ran on OPTIONS")) + with patch.object(mcp_server, "extract_mcp_auth_context", auth): + await mcp_server.handle_streamable_http_mcp(scope, receive, send) + + auth.assert_not_awaited() + start = next(m for m in sent if m["type"] == "http.response.start") + assert start["status"] == 204 + headers = {k.decode().lower(): v.decode() for k, v in start["headers"]} + assert headers["allow"] == "GET, POST, DELETE, OPTIONS" diff --git a/tests/unit/proxy/middleware/test_admission_control_middleware.py b/tests/unit/proxy/middleware/test_admission_control_middleware.py index e449ebc83bf..66549b0c351 100644 --- a/tests/unit/proxy/middleware/test_admission_control_middleware.py +++ b/tests/unit/proxy/middleware/test_admission_control_middleware.py @@ -510,3 +510,32 @@ async def test_admission_lease_cannot_be_reused_by_another_worker(state: Admissi response: Final = await _call(parent) assert response[1]["body"] == b"1" assert state.get_stats() == other_state.get_stats() == AdmissionControlStats(0, 0, 0) + + +@pytest.mark.asyncio +async def test_a_messages_request_turned_away_at_capacity_gets_the_anthropic_overloaded_error( + state: AdmissionControlState, +) -> None: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + middleware: Final = AdmissionControlMiddleware( + _handler_with_release(started, release), + lambda: AdmissionControlSettings(1, 0, 1.0), + state, + ) + + first: Final = asyncio.create_task(_call(middleware, path="/v1/messages")) + await started.wait() + rejected: Final = await _call(middleware, path="/v1/messages") + release.set() + await first + + assert rejected[0]["status"] == 503 + assert dict(rejected[0]["headers"])[b"retry-after"] == b"1" + assert json.loads(rejected[1]["body"]) == { + "type": "error", + "error": { + "type": "overloaded_error", + "message": "Worker at capacity: 1 in-flight, 0 queued requests. Retry later.", + }, + } diff --git a/tests/unit/proxy/middleware/test_request_body_decompression_middleware.py b/tests/unit/proxy/middleware/test_request_body_decompression_middleware.py index ca5774b6be0..2727857c352 100644 --- a/tests/unit/proxy/middleware/test_request_body_decompression_middleware.py +++ b/tests/unit/proxy/middleware/test_request_body_decompression_middleware.py @@ -665,3 +665,33 @@ async def test_non_http_scopes_reach_the_app_untouched(scope: Scope): [(app_scope, app_receive)] = seen assert app_scope is scope assert app_receive is receive + + +@pytest.mark.parametrize( + ("body", "max_mb", "status_code", "anthropic_type", "message"), + [ + pytest.param( + zstd.compress(_low_entropy_text(2 * int(CAP_MB * MIB))), + CAP_MB, + 413, + "request_too_large", + f"Request body is larger than {CAP_MB:g} MB when decoded", + id="decoded-past-the-cap", + ), + pytest.param(b"not zstd at all", 50.0, 400, "invalid_request_error", None, id="undecodable"), + ], +) +def test_a_messages_body_the_proxy_cannot_take_gets_the_anthropic_envelope( + body: bytes, max_mb: float, status_code: int, anthropic_type: str, message: str | None +) -> None: + response: Final = _post( + _client(max_decompressed_body_size_mb=max_mb), + ZSTD_FRAME_MAGIC + body if message is None else body, + ("zstd",), + path="/claude_code_gateway/v1/messages", + ) + + assert response.status_code == status_code + assert response.json()["type"] == "error" + assert response.json()["error"]["type"] == anthropic_type + assert message is None or response.json()["error"]["message"] == message diff --git a/tests/unit/proxy/middleware/test_request_size_limit_middleware.py b/tests/unit/proxy/middleware/test_request_size_limit_middleware.py index 3e8792e0179..0a30b4cf9ef 100644 --- a/tests/unit/proxy/middleware/test_request_size_limit_middleware.py +++ b/tests/unit/proxy/middleware/test_request_size_limit_middleware.py @@ -1,7 +1,9 @@ +from typing import Final + import pytest from starlette.responses import JSONResponse from starlette.testclient import TestClient -from starlette.types import Message +from starlette.types import Message, Receive, Scope, Send from litellm.proxy.middleware.request_size_limit_middleware import ( RequestSizeLimitMiddleware, @@ -133,3 +135,22 @@ async def test_request_size_limit_middleware_rejects_streamed_body_without_conte "more_body": False, } assert received_body_bytes == 1024 * 1024 + + +def test_an_oversized_messages_request_gets_the_anthropic_request_too_large_error() -> None: + async def app(scope: Scope, receive: Receive, send: Send) -> None: + await JSONResponse({"ok": True})(scope, receive, send) + + client: Final = TestClient( + RequestSizeLimitMiddleware(app, get_max_request_size_mb=lambda: 1, is_request_size_limit_enabled=lambda: True) + ) + + response: Final = client.post( + "/v1/messages", content=b"x" * (1024 * 1024 + 1), headers={"content-type": "application/json"} + ) + + assert response.status_code == 413 + assert response.json() == { + "type": "error", + "error": {"type": "request_too_large", "message": "Request size is too large. Max size is 1 MB"}, + } diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py new file mode 100644 index 00000000000..e7b4d087ed3 --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -0,0 +1,520 @@ +import uuid +from collections.abc import Mapping, Sequence +from datetime import datetime +from types import MappingProxyType +from typing import Final +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + OpenAIWebsocketPassthroughLoggingHandler, + openai_websocket_frame_kind, +) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id +from litellm.types.utils import CallTypes + +REALTIME_MODEL: Final = "gpt-realtime-2.1" +RESPONSES_MODEL: Final = "gpt-5.4-mini" +REASONING_MODEL: Final = "gpt-5.4" +TRANSCRIPTION_MODEL: Final = "gpt-4o-mini-transcribe" +SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0) + + +def _logging_obj() -> MagicMock: + logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id=str(uuid.uuid4())) + return logging_obj + + +def _session_created(model: str) -> dict[str, object]: + return {"type": "session.created", "session": {"id": "sess_fake", "model": model}} + + +def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict[str, object]: + return { + "type": "response.done", + "response": { + "id": "resp_fake", + "status": "completed", + "usage": { + "input_tokens": text_in + audio_in, + "output_tokens": text_out + audio_out, + "total_tokens": text_in + audio_in + text_out + audio_out, + "input_token_details": {"text_tokens": text_in, "audio_tokens": audio_in, "cached_tokens": 0}, + "output_token_details": {"text_tokens": text_out, "audio_tokens": audio_out}, + }, + }, + } + + +def _responses_completed( + model: str, + input_tokens: int, + output_tokens: int, + event: str = "response.completed", + served: Mapping[str, object] = MappingProxyType({}), +) -> dict[str, object]: + return { + "type": event, + "response": { + "id": "resp_fake", + "object": "response", + "created_at": 1761700000, + "model": model, + "status": "completed", + "output": [], + **served, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + } + + +def _handle( + messages: Sequence[Mapping[str, object]], + kwargs: Mapping[str, object] = MappingProxyType({}), + logging_obj: MagicMock | None = None, + url_route: str = "/openai/v1/realtime", +): + return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( + websocket_messages=messages, + logging_obj=logging_obj if logging_obj is not None else _logging_obj(), + url_route=url_route, + start_time=SESSION_START, + kwargs=kwargs, + ) + + +def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio_out: int) -> float: + prices: Final = litellm.model_cost[model] + return ( + text_in * prices["input_cost_per_token"] + + audio_in * prices["input_cost_per_audio_token"] + + text_out * prices["output_cost_per_token"] + + audio_out * prices["output_cost_per_audio_token"] + ) + + +def test_realtime_session_is_priced_on_the_model_the_session_reported(): + logging_obj: Final = _logging_obj() + messages: Final = [ + _session_created(REALTIME_MODEL), + {"type": "response.created", "response": {"id": "resp_fake"}}, + _realtime_response_done(text_in=33, audio_in=120, text_out=16, audio_out=240), + ] + + handled: Final = _handle(messages, logging_obj=logging_obj) + + expected_cost: Final = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240) + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == REALTIME_MODEL + assert handled["kwargs"]["custom_llm_provider"] == "openai" + assert handled["result"].model == REALTIME_MODEL + assert handled["result"].usage.total_tokens == 409 + assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost) + assert logging_obj.model_call_details["model"] == REALTIME_MODEL + + +def test_realtime_cost_covers_every_turn_of_the_session(): + one_turn: Final = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)]) + two_turns: Final = _handle( + [ + _session_created(REALTIME_MODEL), + _realtime_response_done(10, 0, 5, 0), + _realtime_response_done(10, 0, 5, 0), + ] + ) + + assert two_turns["kwargs"]["response_cost"] == pytest.approx(2 * one_turn["kwargs"]["response_cost"]) + assert two_turns["result"].usage.total_tokens == 30 + + +def test_responses_turn_is_priced_on_the_model_the_response_reported(): + handled: Final = _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)], + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == RESPONSES_MODEL + assert handled["result"].usage.prompt_tokens == 13 + assert handled["result"].usage.completion_tokens == 5 + + +def test_responses_cost_covers_every_turn_of_the_connection(): + handled: Final = _handle( + [ + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 45 + + +def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_named(): + """ + The caller picks what to ask for, and the provider picks what to run. Only + the provider's own frames may set the price of a session, so a model the + caller put in the request cannot stand in for one the provider never named. + """ + logging_obj: Final = _logging_obj() + + handled: Final = _handle( + [_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0)], + kwargs=MappingProxyType({"model": REALTIME_MODEL}), + logging_obj=logging_obj, + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REALTIME_MODEL} + assert logging_obj.model_call_details == {} + + +def test_each_responses_turn_is_priced_on_the_model_that_turn_ran(): + """ + One socket can run several models one after another, so a connection that + starts on a big model and finishes on a small one has to pay the big + model's price for the turn that used it. + """ + handled: Final = _handle( + [ + _responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + first: Final = litellm.model_cost[REASONING_MODEL] + second: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = ( + 13 * first["input_cost_per_token"] + + 5 * first["output_cost_per_token"] + + 20 * second["input_cost_per_token"] + + 7 * second["output_cost_per_token"] + ) + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 45 + + +def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model(): + handled: Final = _handle( + [ + _session_created(RESPONSES_MODEL), + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}), + ], + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + + +def test_a_session_nobody_named_a_model_for_is_left_unpriced(): + logging_obj: Final = _logging_obj() + handled: Final = _handle( + [{"type": "response.done", "response": {"usage": {"input_tokens": 10, "output_tokens": 5}}}], + kwargs=MappingProxyType({"litellm_call_id": "call-fake"}), + logging_obj=logging_obj, + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"litellm_call_id": "call-fake"} + assert logging_obj.model_call_details == {} + + +def test_a_session_that_reported_no_usage_still_records_its_model(): + logging_obj: Final = _logging_obj() + + handled: Final = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"} + assert logging_obj.model_call_details["model"] == REALTIME_MODEL + assert logging_obj.model_call_details["custom_llm_provider"] == "openai" + + +def test_a_responses_turn_without_readable_usage_is_not_billed(): + handled: Final = _handle( + [{"type": "response.completed", "response": {"model": RESPONSES_MODEL, "status": "completed"}}], + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} + + +def test_the_price_of_a_multi_model_connection_reaches_the_spend_row(): + """ + The logging layer prices whatever result it is handed, and one result can + only name one model, so a connection that ran two models has to carry the + price it already worked out or the cheaper model pays for both turns. The + row it lands in names the model the connection opened on. + """ + handled: Final = _handle( + [ + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(REASONING_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + logged_cost: Final = litellm.response_cost_calculator( + response_object=handled["result"], + model=handled["kwargs"]["model"], + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + session_prices: Final = litellm.model_cost[RESPONSES_MODEL] + priced_on_the_session_model: Final = ( + 33 * session_prices["input_cost_per_token"] + 12 * session_prices["output_cost_per_token"] + ) + assert logged_cost == pytest.approx(handled["kwargs"]["response_cost"]) + assert logged_cost > priced_on_the_session_model + assert handled["kwargs"]["model"] == RESPONSES_MODEL + + +def test_turns_that_start_in_the_same_instant_each_keep_their_own_spend_row(): + """ + Spend logs key each row by the id its result carries and skip a repeated id + without a word, so two turns that start in the same instant, on one socket + or on two, each need an id of their own or one of them is never billed. + """ + turns: Final = [ + _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)], + url_route="/openai_passthrough/v1/responses", + ) + for _ in range(2) + ] + + spend_log_ids: Final = { + get_spend_logs_id(CallTypes.pass_through.value, turn["result"].model_dump(), turn["kwargs"]) for turn in turns + } + assert len(spend_log_ids) == 2 + assert None not in spend_log_ids + + +def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): + """ + A client that closes before response.completed leaves no usage behind to + bill, but response.created already named the model, so the row says which + model answered instead of falling back to unknown. + """ + logging_obj: Final = _logging_obj() + + handled: Final = _handle( + [{"type": "response.created", "response": {"id": "resp_fake", "model": REASONING_MODEL}}], + logging_obj=logging_obj, + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"} + assert logging_obj.model_call_details["model"] == REASONING_MODEL + + +def _transcription_session(model: str) -> tuple[dict[str, object], ...]: + return ( + {"type": "session.created", "session": {"type": "transcription", "audio": {"input": {"transcription": None}}}}, + { + "type": "session.updated", + "session": {"type": "transcription", "audio": {"input": {"transcription": {"model": model}}}}, + }, + ) + + +def _transcription_completed(usage: Mapping[str, object]) -> dict[str, object]: + return { + "type": "conversation.item.input_audio_transcription.completed", + "item_id": "item_fake", + "transcript": "hello", + "usage": usage, + } + + +@pytest.mark.parametrize( + ("usage", "billed_units"), + [ + ({"type": "duration", "seconds": 3}, ((3, "input_cost_per_second"),)), + ( + { + "type": "tokens", + "input_tokens": 30, + "output_tokens": 11, + "input_token_details": {"text_tokens": 0, "audio_tokens": 30}, + }, + ((30, "input_cost_per_audio_token"), (11, "output_cost_per_token")), + ), + ], +) +def test_a_transcription_session_is_billed_for_every_clip_it_transcribed(usage, billed_units): + handled: Final = _handle( + [ + *_transcription_session(TRANSCRIPTION_MODEL), + _transcription_completed(usage), + _transcription_completed(usage), + ], + url_route="/openai_passthrough/v1/realtime", + ) + + logged_cost: Final = litellm.response_cost_calculator( + response_object=handled["result"], + model=handled["kwargs"]["model"], + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + prices: Final = litellm.model_cost[TRANSCRIPTION_MODEL] + expected_cost: Final = 2 * sum(units * prices[rate] for units, rate in billed_units) + assert expected_cost > 0 + assert handled["kwargs"]["model"] == TRANSCRIPTION_MODEL + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert logged_cost == pytest.approx(expected_cost) + + +def test_a_conversation_that_also_transcribes_the_caller_is_logged_on_its_own_model(): + handled: Final = _handle( + [ + { + "type": "session.created", + "session": { + "model": REALTIME_MODEL, + "audio": {"input": {"transcription": {"model": TRANSCRIPTION_MODEL}}}, + }, + }, + _realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0), + ] + ) + + assert handled["kwargs"]["model"] == REALTIME_MODEL + + +@pytest.mark.parametrize("event", ["response.incomplete", "response.failed"]) +def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(event): + """ + A turn cut short by max_output_tokens or a provider failure still reports + the tokens it burned, and OpenAI bills them, so the row has to as well. + """ + handled: Final = _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, event=event)], + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 18 + + +def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed(): + handled: Final = _handle( + [{"type": "response.failed", "response": {"model": RESPONSES_MODEL, "status": "failed", "usage": None}}], + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} + + +@pytest.mark.parametrize( + ("turn_model", "requested_model", "priced_as"), + [ + ("gpt-unreleased-qa-snapshot", RESPONSES_MODEL, RESPONSES_MODEL), + ("gpt-5.4-mini-qa-preview", RESPONSES_MODEL, RESPONSES_MODEL), + (RESPONSES_MODEL, REASONING_MODEL, RESPONSES_MODEL), + ], + ids=["unmapped_model", "model_the_cost_map_prices_at_zero", "priced_turn_model"], +) +def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_has_no_price( + turn_model: str, requested_model: str, priced_as: str +): + """ + The provider can answer under a model name the cost map does not price + yet. That turn still ran the model the key was allowed to ask for, so it is + billed at that model's price instead of $0, while a turn whose own model + has a price keeps it. + """ + handled: Final = _handle( + [_responses_completed(turn_model, input_tokens=13, output_tokens=5)], + kwargs=MappingProxyType({"litellm_params": {"metadata": {"model_group": requested_model}}}), + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[priced_as] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert expected_cost > 0 + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == turn_model + + +def _http_passthrough_cost(response: Mapping[str, object]) -> float: + http_response: Final = httpx.Response( + 200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses") + ) + _, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost( + RESPONSES_MODEL, http_response, _logging_obj(), "openai" + ) + return cost + + +@pytest.mark.parametrize( + "served", + [ + {"service_tier": "priority"}, + {"output": [{"type": "web_search_call", "id": "ws_fake", "status": "completed", "action": {"type": "search"}}]}, + ], + ids=["priority_tier", "hosted_web_search"], +) +def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_response(served): + """ + The provider bills a turn by the tier that served it and the hosted tools it + called, not only its tokens, so a websocket turn has to cost exactly what the + same response costs over HTTP, never the cheaper token-only price. + """ + turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served) + plain_turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5) + + handled: Final = _handle([turn], url_route="/openai_passthrough/v1/responses") + + assert handled["kwargs"]["response_cost"] == pytest.approx(_http_passthrough_cost(turn["response"])) + assert handled["kwargs"]["response_cost"] > _http_passthrough_cost(plain_turn["response"]) + + +@pytest.mark.parametrize( + ("frame_type", "kind"), + [ + ("response.completed", "ends_turn"), + ("response.incomplete", "ends_turn"), + ("response.failed", "ends_turn"), + ("response.done", "ends_turn"), + ("session.created", "opens_session"), + ("session.updated", "updates_session"), + ("conversation.item.input_audio_transcription.completed", "bills_usage"), + ("response.created", "starts_turn"), + ("response.output_text.delta", "relayed"), + ("error", "refuses_turn"), + ], +) +def test_each_frame_is_classified_by_what_it_does_to_a_turn(frame_type, kind): + assert openai_websocket_frame_kind({"type": frame_type}) == kind diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 4fe358c38a3..cd6c1ff11b5 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1,5 +1,6 @@ import asyncio import base64 +import contextvars import gzip import json import logging @@ -7,20 +8,20 @@ import os import sys import uuid import zlib -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime from io import BytesIO from types import MappingProxyType, ModuleType, SimpleNamespace from typing import Final, Optional, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch from uuid import UUID import httpx import pytest import respx -from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, UploadFile +from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, UploadFile, WebSocketDisconnect from fastapi.responses import StreamingResponse from fastapi.routing import APIRoute from fastapi.testclient import TestClient @@ -35,9 +36,13 @@ from litellm.constants import ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH, ANTHROPIC from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( @@ -46,9 +51,15 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrou from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + STORED_PROMPT_PROBE_MODEL, + WALKOUT_TURN_DRAIN_SECONDS, HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, SafeRouteAdder, + WebsocketFrameModelAllowed, + WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, _build_passthrough_failure_request_payload, _registered_pass_through_routes, _truncate_upstream_error_body, @@ -63,12 +74,15 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( set_env_variables_in_header, websocket_passthrough_request, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes @@ -7382,6 +7396,7 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) @@ -7434,6 +7449,7 @@ async def test_websocket_passthrough_propagates_active_trace_context( mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) @@ -7500,6 +7516,7 @@ async def test_websocket_passthrough_never_forwards_caller_credentials_upstream( mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) @@ -7593,10 +7610,11 @@ def _patched_websocket_passthrough_environment(upstream_ws): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) - yield + yield mock_proxy_logging async def _pending_receive(): @@ -7901,6 +7919,1487 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason) +SUCCESS_HANDLER: Final = ( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler" +) + + +KEY_MODEL_REFUSAL: Final = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." +KEY_MODEL_CLOSE_REASON: Final = "key not allowed to access model." + + +class RecordingFrameModelGate: + def __init__( + self, + refusals: Mapping[str, str] = MappingProxyType({}), + groups: Mapping[str, tuple[str, ...]] = MappingProxyType({}), + error_type: str = "key_model_access_denied", + ): + self.refusals = refusals + self.groups = groups + self.error_type = error_type + self.checked: list[str] = [] + + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: + self.checked.append(model) + refusal: Final = self.refusals.get(model) + if refusal is not None: + return WebsocketFrameModelRefused(message=refusal, error_type=self.error_type) + return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ())) + + +def _received_frame(frame: str | bytes) -> Mapping[str, str | bytes]: + if isinstance(frame, bytes): + return {"type": "websocket.receive", "bytes": frame} + return {"type": "websocket.receive", "text": frame} + + +async def _relay_client_frames( + frames: Sequence[str | bytes], + gate: WebsocketFrameModelGate | None, + user_api_key_dict: UserAPIKeyAuth | None = None, +) -> SimpleNamespace: + upstream_ws: Final = RecordingUpstreamWebSocket() + websocket: Final = _client_websocket( + AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}]) + ) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai/v1/realtime", + accept_websocket=False, + client_frame_model_gate=gate, + ) + + return SimpleNamespace(upstream=upstream_ws, client=websocket, success_handler=success_handler) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame, expected_model", + [ + ({"type": "response.create", "model": "gpt-5.4-mini"}, "gpt-5.4-mini"), + ({"type": "response.create", "response": {"model": "gpt-5.4"}}, "gpt-5.4"), + ({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}, "gpt-realtime-2.1"), + ( + {"type": "session.update", "session": {"input_audio_transcription": {"model": "gpt-4o-transcribe"}}}, + "gpt-4o-transcribe", + ), + ( + { + "type": "session.update", + "session": {"audio": {"input": {"transcription": {"model": "gpt-transcribe"}}}}, + }, + "gpt-transcribe", + ), + ], +) +async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(frame, expected_model): + gate: Final = RecordingFrameModelGate() + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked == [expected_model] + assert [call.args[0] for call in relayed.upstream.send.await_args_list] == [json.dumps(frame)] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cannot_call(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) + + relayed: Final = await _relay_client_frames([frame], gate) + + assert relayed.upstream.send.await_args_list == [] + assert json.loads(relayed.client.send_text.await_args.args[0]) == { + "type": "error", + "error": {"type": "key_model_access_denied", "message": KEY_MODEL_REFUSAL}, + } + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + relayed.upstream.close.assert_awaited() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the_refusal(): + """ + The close frame only has room for 123 bytes, so a refusal cut at that limit + reads as a truncated number. The first sentence names the reason whole, and + the error event before the close still carries every word with its type. + """ + refusal: Final = "Budget has been exceeded! Key=" + "f" * 64 + " Spend=1.23456789, Budget=1.0" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": refusal}), error_type="budget_exceeded") + + relayed: Final = await _relay_client_frames([json.dumps({"type": "response.create", "model": "gpt-5.4"})], gate) + + assert relayed.client.close.await_args_list == [call(code=1008, reason="Budget has been exceeded!")] + assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == { + "type": "budget_exceeded", + "message": refusal, + } + + +@pytest.mark.asyncio +async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps( + { + "type": "session.update", + "session": {"model": "gpt-realtime-2.1", "input_audio_transcription": {"model": "gpt-4o-transcribe"}}, + } + ) + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_stops_relaying_after_a_refused_frame(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frames: Final = [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), + json.dumps({"type": "response.create"}), + ] + + relayed: Final = await _relay_client_frames(frames, gate) + + assert relayed.upstream.send.await_args_list == [] + assert gate.checked == ["gpt-realtime-2.1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + '{"type": "input_audio_buffer.append", "audio": "AAAA"}', + '{"type": "response.create", "response": {"modalities": ["audio"]}}', + '{"type": "session.update", "session": {"model": 42}}', + 'not json at all, but it does say "model"', + '{"type": "session.update", "session": "model"}', + ], +) +async def test_websocket_passthrough_forwards_frames_that_name_no_model(frame): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == [] + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_never_bills_the_model_a_client_frame_named(): + """ + A frame is what the caller asked for, not what the provider ran. A caller + who names a cheap model in a frame the provider ignores must not move the + session onto that model's price, so only the provider's own frames decide + what a session is billed as. + """ + gate: Final = RecordingFrameModelGate() + frames: Final = [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), + json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), + ] + + relayed: Final = await _relay_client_frames(frames, gate) + + assert gate.checked == ["gpt-realtime-2.1", "gpt-realtime-2.1-mini"] + assert "model" not in relayed.success_handler.call_args.kwargs + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame, expected_model", + [ + ('{"type": "response.create", "\\u006dodel": "gpt-5.4"}', "gpt-5.4"), + ('{"type": "session.update", "session": {"\\u006d\\u006fdel": "gpt-realtime-2.1"}}', "gpt-realtime-2.1"), + ], +) +async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(frame, expected_model): + """ + JSON lets a caller spell the same key several ways, so a frame carrying + "\\u006dodel" asks for a model just as plainly as one carrying "model" and + has to face the same check. + """ + gate: Final = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL})) + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == [expected_model] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode() + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-realtime-2.1"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + + +STORED_PROMPT_FRAMES: Final = ( + {"type": "response.create", "prompt": {"id": "pmpt_fake", "version": "2"}}, + {"type": "response.create", "response": {"prompt": {"id": "pmpt_fake"}}}, + {"type": "session.update", "session": {"prompt": {"id": "pmpt_fake"}}}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES) +async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_to_some_models(frame): + """ + A stored prompt carries its own model on the provider side, so a key held to + a model list could reach any model through one. The gateway cannot see that + model, so a key that is not allowed every model cannot run a stored prompt. + """ + gate: Final = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + error_message: Final = json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"] + close_reason: Final = relayed.client.close.await_args.kwargs["reason"] + assert gate.checked == [STORED_PROMPT_PROBE_MODEL] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args.kwargs["code"] == 1008 + assert "stored prompt" in close_reason + assert close_reason.endswith(".") and error_message.startswith(close_reason) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES) +async def test_websocket_passthrough_forwards_a_stored_prompt_for_a_key_allowed_every_model(frame): + gate: Final = RecordingFrameModelGate() + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [json.dumps(frame)] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_prompt_is_refused_for_spend(): + refusal: Final = "Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0" + gate: Final = RecordingFrameModelGate( + MappingProxyType({STORED_PROMPT_PROBE_MODEL: refusal}), error_type="budget_exceeded" + ) + + relayed: Final = await _relay_client_frames([json.dumps(STORED_PROMPT_FRAMES[0])], gate) + + assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == { + "type": "budget_exceeded", + "message": refusal, + } + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_the_model_a_hosted_tool_names(): + """A hosted tool such as image generation runs its own model, so that model gets the same check as the turn's""" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-image-2": KEY_MODEL_REFUSAL})) + frame: Final = { + "type": "response.create", + "model": "gpt-5.4-mini", + "tools": [{"type": "web_search"}, {"type": "image_generation", "model": "gpt-image-2"}], + } + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked == ["gpt-5.4-mini", "gpt-image-2"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args.kwargs["code"] == 1008 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "model": "gpt-5.4-mini", "tools": [{"type": "image_generation"}]}, + {"type": "response.create", "response": {"tools": [{"type": "image_generation", "model": None}]}}, + ], +) +async def test_websocket_passthrough_refuses_an_image_tool_without_a_model_for_a_key_limited_to_some_models(frame): + gate: Final = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked[-1] == STORED_PROMPT_PROBE_MODEL + assert relayed.upstream.send.await_args_list == [] + assert "image_generation tool has to name its model" in relayed.client.close.await_args.kwargs["reason"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_a_frame_only_the_standard_json_parser_reads(): + """A lone surrogate escape is valid JSON syntax the fast parser refuses, so the frame is still checked""" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + frame: Final = '{"type": "response.create", "model": "gpt-5.4", "metadata": {"note": "\\ud800"}}' + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-5.4"] + assert relayed.upstream.send.await_args_list == [] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_forwards_binary_audio_untouched(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) + audio: Final = bytes(range(256)) * 4 + + relayed: Final = await _relay_client_frames([audio], gate) + + assert gate.checked == [] + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [audio] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_without_a_frame_gate_forwards_every_frame(): + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) + + relayed: Final = await _relay_client_frames([frame], None) + + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame] + assert "model" not in relayed.success_handler.call_args.kwargs + + +REALTIME_SESSION_CREATED: Final = json.dumps( + {"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}} +) +REALTIME_RESPONSE_DONE: Final = json.dumps( + { + "type": "response.done", + "response": {"status": "completed", "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}}, + } +) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out_of(): + """ + A realtime client closes the moment it has the answer it asked for, so the + next frame the provider sends lands on a socket that is already gone. That + is how a session ends, not a gateway failure, and the turns the caller + already used still have to reach spend, including the last frame the client + never collected. + """ + upstream_ws: Final = FakeUpstreamWebSocket( + REALTIME_SESSION_CREATED, + REALTIME_RESPONSE_DONE, + json.dumps({"type": "response.output_text.delta", "delta": "too late"}), + ) + websocket: Final = _client_websocket(_pending_receive) + websocket.send_text = AsyncMock(side_effect=[None, None, WebSocketDisconnect(code=1006)]) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai_passthrough/v1/realtime", + accept_websocket=False, + ) + + relayed_frames: Final = success_handler.call_args.kwargs["response_body"] + assert [frame["type"] for frame in relayed_frames] == [ + "session.created", + "response.done", + "response.output_text.delta", + ] + + +class ScriptedUpstreamWebSocket: + """Answers each client frame with its scripted reply""" + + def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = ()): + self._frames: asyncio.Queue[str | None] = asyncio.Queue() + for frame in opening: + self._frames.put_nowait(frame) + self._replies = iter(replies) + self._closed = False + self.send = AsyncMock(side_effect=self._reply) + self.close = AsyncMock(side_effect=self._flush_and_close) + + async def _reply(self, frame: str | bytes) -> None: + for reply in next(self._replies, ()): + self._frames.put_nowait(reply) + + async def _flush_and_close(self, *args: object, **kwargs: object) -> None: + if self._closed: + return + self._closed = True + self._frames.put_nowait(None) + while not self._frames.empty(): + await asyncio.sleep(0) + + async def recv(self, decode: bool | None = None) -> str: + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + frame: Final = await self._frames.get() + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame + + +def _turn_taking_client( + frames: Sequence[str], answered_by: str, answers_due: Sequence[int] | None = None +) -> MagicMock: + """A client that sends each frame, and then leaves, once the answers due by then reached it (by default the + answer to every earlier frame), and that Starlette refuses to send to once it is closed""" + from starlette.websockets import WebSocketState + + delivered: Final[list[str]] = [] + script: Final = iter(frames) + receives: Final = iter(range(len(frames) + 1) if answers_due is None else answers_due) + + async def receive(): + answers_due: Final = next(receives) + while sum(answered_by in text for text in delivered) < answers_due: + await asyncio.sleep(0) + frame: Final = next(script, None) + return {"type": "websocket.disconnect"} if frame is None else _received_frame(frame) + + websocket: Final = _client_websocket(receive) + + async def send_text(text: str) -> None: + if websocket.application_state == WebSocketState.DISCONNECTED: + raise RuntimeError('Cannot call "send" once a close message has been sent.') + delivered.append(text) + + websocket.send_text = AsyncMock(side_effect=send_text) + return websocket + + +async def _relay_turns( + client: MagicMock, + upstream_ws: ScriptedUpstreamWebSocket, + gate: WebsocketFrameModelGate, + endpoint: str = "/openai_passthrough/v1/responses", + target: str = "wss://api.openai.com/v1/responses", + opening_access_groups: Sequence[str] = (), + walkout_turn_drain_seconds: float = 0.05, +) -> SimpleNamespace: + with ( + _patched_websocket_passthrough_environment(upstream_ws) as proxy_logging, + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target=target, + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-key", matched_model_access_groups=list(opening_access_groups) + ), + forward_headers=False, + endpoint=endpoint, + accept_websocket=False, + client_frame_model_gate=gate, + upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=walkout_turn_drain_seconds, + ), + timeout=10, + ) + return SimpleNamespace(success_handler=success_handler, proxy_logging=proxy_logging) + + +def _responses_turn(model: str) -> list[str]: + return [_responses_turn_started(model), _responses_turn_ended(model)] + + +def _responses_turn_started(model: str) -> str: + return json.dumps({"type": "response.created", "response": {"model": model}}) + + +def _responses_turn_ended(model: str) -> str: + return ( + + json.dumps( + { + "type": "response.completed", + "response": {"model": model, "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}}, + } + ) + ) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_the_groups_that_allowed_it(): + """ + One socket can run a model from one access group and then a model from + another. Each turn has to reach spend as it finishes, under its own call, + and be charged to the group that let its model through, so a group budget + sees its share during the session and never pays for the other group's turn. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini"), _responses_turn("gpt-5.4")]) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [ + ["gpt-5.4-mini", "gpt-5.4-mini"], + ["gpt-5.4", "gpt-5.4"], + ] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + assert [row["litellm_params"]["metadata"]["user_api_key"] for row in rows] == ["hashed-key", "hashed-key"] + assert rows[0]["litellm_call_id"] != rows[1]["litellm_call_id"] + assert rows[1]["start_time"] == rows[0]["end_time"] + relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_holds_the_keys_parallel_slot_until_the_session_ends(): + """ + A key allowed one parallel request holds that slot for as long as its + socket stays open. A turn that finishes and reaches spend mid-session must + not free it, or a second session gets in while the first is still running. + """ + limiter_cache: Final = DualCache() + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(limiter_cache)) + key: Final = UserAPIKeyAuth(api_key=hash_token("sk-one-socket"), max_parallel_requests=1) + logged_rows: Final[list[asyncio.Task[None]]] = [] + + async def admit(user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> None: + await limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=limiter_cache, data=data, call_type=call_type + ) + + async def release(data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object) -> None: + await limiter.async_post_call_success_hook(data=data, user_api_key_dict=user_api_key_dict, response=response) + + async def log_row(**row: object) -> None: + await limiter.async_log_success_event( + kwargs=row, response_obj=None, start_time=row["start_time"], end_time=row["end_time"] + ) + + async def opens_another_session() -> bool: + await asyncio.gather(*logged_rows) + try: + await asyncio.create_task( + admit(key, {"litellm_call_id": str(uuid.uuid4())}, "pass_through_endpoint"), + context=contextvars.Context(), + ) + except ProxyRateLimitError: + return False + return True + + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], answered_by='"response.completed"' + ) + take_frame: Final = client.receive + while_open: Final[list[bool]] = [] + + async def receive() -> Mapping[str, str | bytes]: + message: Final = await take_frame() + if message["type"] == "websocket.disconnect": + while_open.append(await opens_another_session()) + return message + + client.receive = receive + + with ( + _patched_websocket_passthrough_environment( + ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini")]) + ) as proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker, + patch(SUCCESS_HANDLER, new_callable=AsyncMock, side_effect=log_row), + ): + proxy_logging.pre_call_hook = AsyncMock(side_effect=admit) + proxy_logging.post_call_success_hook = AsyncMock(side_effect=release) + worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine)) + ) + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target="wss://api.openai.com/v1/responses", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=key, + forward_headers=False, + endpoint="/openai_passthrough/v1/responses", + accept_websocket=False, + client_frame_model_gate=RecordingFrameModelGate(), + upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=0.05, + ), + timeout=10, + ) + after_close: Final = await opens_another_session() + + assert len(logged_rows) == 1 + assert (while_open, after_close) == ([False], True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "models, upstream_refuses, while_open", + [ + ((), False, [False]), + (("gpt-5.4-mini",), False, [False]), + (("gpt-5.4-mini", "gpt-5.4"), False, [False]), + ((), True, []), + ], + ids=["no-turn", "one-turn", "two-turns", "upstream-refused"], +) +async def test_websocket_passthrough_frees_one_legacy_parallel_slot_only_when_the_session_ends( + models: tuple[str, ...], upstream_refuses: bool, while_open: list[bool] +): + limiter: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=lambda: datetime(2026, 1, 31, 23, 59, 30) + ) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = limiter + key: Final = UserAPIKeyAuth(api_key=hash_token("sk-legacy-socket"), max_parallel_requests=2) + logged_rows: Final[list[asyncio.Task[None]]] = [] + seen_while_open: Final[list[bool]] = [] + + async def admit(user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str) -> None: + await limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=call_type + ) + + async def log_row(**row: object) -> None: + await limiter.async_log_success_event( + kwargs=row, response_obj=None, start_time=row["start_time"], end_time=row["end_time"] + ) + + async def opens_another_session() -> bool: + await asyncio.gather(*logged_rows) + try: + await admit(key, {}, "pass_through_endpoint") + except ProxyRateLimitError: + return False + return True + + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": model}) for model in models], + answered_by='"response.completed"', + ) + take_frame: Final = client.receive + + async def receive() -> Mapping[str, str | bytes]: + message: Final = await take_frame() + if message["type"] == "websocket.disconnect": + seen_while_open.append(await opens_another_session()) + return message + + client.receive = receive + await admit(key, {}, "pass_through_endpoint") + + with ExitStack() as stack: + proxy_logging: Final = stack.enter_context( + _patched_websocket_passthrough_environment( + ScriptedUpstreamWebSocket(replies=[_responses_turn(model) for model in models]) + ) + ) + worker: Final = stack.enter_context( + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") + ) + stack.enter_context(patch(SUCCESS_HANDLER, new_callable=AsyncMock, side_effect=log_row)) + if upstream_refuses: + stack.enter_context( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + side_effect=OSError("upstream refused the socket"), + ) + ) + proxy_logging.pre_call_hook = AsyncMock(side_effect=admit) + proxy_logging.arelease_responses_websocket_slot = proxy_logging_obj.arelease_responses_websocket_slot + worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine)) + ) + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target="wss://api.openai.com/v1/responses", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=key, + forward_headers=False, + endpoint="/openai_passthrough/v1/responses", + accept_websocket=False, + client_frame_model_gate=RecordingFrameModelGate(), + upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=0.05, + ), + timeout=10, + ) + after_close: Final = (await opens_another_session(), await opens_another_session()) + + assert len(logged_rows) == len(models) + assert (seen_while_open, after_close) == (while_open, (True, False)) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_asked_for_each(): + """ + A client can send its next response.create before the provider finished the + last one, and the provider then runs them one after the other. Each turn + still has to reach spend under the groups that let its own frame through, + and under the model that frame asked for, even when the provider answers + under a dated snapshot name a per-model budget was never set on. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + answers_due=(0, 0, 2), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[_responses_turn("gpt-5.4-mini-2099-01-01"), _responses_turn("gpt-5.4-2099-01-01")] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [ + ["gpt-5.4-mini-2099-01-01", "gpt-5.4-mini-2099-01-01"], + ["gpt-5.4-2099-01-01", "gpt-5.4-2099-01-01"], + ] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] + + +def _responses_error(code: str) -> str: + return json.dumps({"type": "error", "status": 400, "error": {"type": "invalid_request_error", "code": code}}) + + +def _on_lane(frame: str, stream_id: str) -> str: + return json.dumps({**json.loads(frame), "stream_id": stream_id}) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_frames_that_asked_for_them(): + """ + The provider answers a response.create it cannot run, one naming no model + or a stale previous_response_id, with an error event and keeps the socket + open. Those frames never run, so each later turn still has to reach spend + under the groups and model its own frame asked for. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "input": "hi"}), + json.dumps({"type": "response.create", "model": "gpt-5.4", "previous_response_id": "resp_gone"}), + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + answers_due=(0, 0, 0, 0, 2), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[ + [_responses_error("missing_required_parameter")], + [_responses_error("previous_response_not_found")], + _responses_turn("gpt-5.4-mini-2099-01-01"), + _responses_turn("gpt-5.4-2099-01-01"), + ] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_frame_that_asked_for_each(): + """ + A client can put each response.create on its own stream_id lane, and the + provider runs lanes side by side, so a later frame's turn can end first. + Each turn still has to reach spend on its own row, under the groups and + model its own frame asked for. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "stream_id": "slow", "model": "gpt-5.4"}), + json.dumps({"type": "response.create", "stream_id": "fast", "model": "gpt-5.4-mini"}), + ], + answered_by='"response.completed"', + answers_due=(0, 0, 2), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[ + [_on_lane(_responses_turn_started("gpt-5.4"), "slow")], + [ + *(_on_lane(frame, "fast") for frame in _responses_turn("gpt-5.4-mini")), + _on_lane(_responses_turn_ended("gpt-5.4"), "slow"), + ], + ] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[(frame["stream_id"], frame["response"]["model"]) for frame in row["response_body"]] for row in rows] == [ + [("fast", "gpt-5.4-mini"), ("fast", "gpt-5.4-mini")], + [("slow", "gpt-5.4"), ("slow", "gpt-5.4")], + ] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream_id", ["", 7]) +async def test_websocket_passthrough_refuses_a_frame_on_a_lane_the_provider_error_cannot_name(stream_id): + """ + The provider refuses a frame with an invalid stream_id through an error that + names no lane (OpenAI WebSocket mode guide, "Errors to handle", read + 2026-11-02, which also says an empty stream_id is invalid). Relayed, that + error would refuse a turn waiting on the default lane instead, and that turn + would reach spend without the groups its own frame asked for. + """ + gate: Final = RecordingFrameModelGate() + allowed: Final = json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}) + + relayed: Final = await _relay_client_frames( + [allowed, json.dumps({"type": "response.create", "stream_id": stream_id, "model": "gpt-5.4-mini"})], gate + ) + + assert [sent.args[0] for sent in relayed.upstream.send.await_args_list] == [allowed] + assert json.loads(relayed.client.send_text.await_args.args[0])["type"] == "error" + assert relayed.client.close.await_args.kwargs["code"] == 1008 + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_key_opened(): + """ + A realtime session runs the model it was opened on, so its turns stay + charged to the groups and the model the key opened it with, whatever model + a later session.update or response.create names. + """ + gate: Final = RecordingFrameModelGate(groups=MappingProxyType({"gpt-realtime-2.1-mini": ("tier-b",)})) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1-mini"}}), + json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), + ], + answered_by='"type"', + answers_due=(1, 2, 4), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [json.dumps({"type": "session.updated"})], + [json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), REALTIME_RESPONSE_DONE], + ], + ) + + relayed: Final = await _relay_turns( + client, + upstream_ws, + gate, + "/openai/v1/realtime", + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + opening_access_groups=("tier-a",), + ) + + assert gate.checked == ["gpt-realtime-2.1-mini", "gpt-realtime-2.1-mini"] + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [["tier-a"]] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-realtime-2.1"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_no_spend_row_for_a_session_refused_before_any_answer(): + """A session refused at its first frame ran nothing, so its only record is the refusal, not an extra $0 row""" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + client: Final = _turn_taking_client([json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"error"') + + relayed: Final = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate) + + assert relayed.success_handler.call_args_list == [] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_out_of(): + """ + A client can leave while a later turn is still running. That turn never + ends, but it was started and may be billed, so it still gets its own row. + """ + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.created"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[_responses_turn("gpt-5.4-mini"), [_responses_turn_started("gpt-5.4")]] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate()) + + assert [ + [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.completed"], ["response.created"]] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_each_lane_a_client_walked_out_of_to_the_frame_that_started_it(): + """ + A client can leave while turns on several stream_id lanes are still + running. Each unfinished turn still gets its own row, under the groups and + model its own frame asked for. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "stream_id": "slow", "model": "gpt-5.4"}), + json.dumps({"type": "response.create", "stream_id": "fast", "model": "gpt-5.4-mini"}), + ], + answered_by='"response.created"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[ + [_on_lane(_responses_turn_started("gpt-5.4"), "slow")], + [_on_lane(_responses_turn_started("gpt-5.4-mini"), "fast")], + ] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["stream_id"] for frame in row["response_body"]] for row in rows] == [["slow"], ["fast"]] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4", "gpt-5.4-mini"] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-b"], + ["tier-a"], + ] + + +class UpstreamFinishingTheTurnAfterTheClientLeft: + """Starts a turn for each client frame, unless it only answers once the client left, sends the rest of the session + once the client left, and stops once closed""" + + def __init__( + self, + client_left: asyncio.Event, + started: str | None, + rest_of_the_turn: Sequence[str], + opening: Sequence[str] = (), + ): + self._frames: asyncio.Queue[str | None] = asyncio.Queue() + for frame in opening: + self._frames.put_nowait(frame) + self._client_left = client_left + self._started = started + self._rest_of_the_turn = rest_of_the_turn + self._closed = False + self._finishing: list[asyncio.Task[None]] = [] + self.send = AsyncMock(side_effect=self._start_turn) + self.close = AsyncMock(side_effect=self._close) + + def _close(self, *args: object, **kwargs: object) -> None: + self._closed = True + self._frames.put_nowait(None) + + async def _start_turn(self, frame: str | bytes) -> None: + if self._started is not None: + self._frames.put_nowait(self._started) + if not self._finishing: + self._finishing.append(asyncio.create_task(self._finish_once_the_client_left())) + + async def _finish_once_the_client_left(self) -> None: + await self._client_left.wait() + for frame in self._rest_of_the_turn: + await asyncio.sleep(0) + if self._closed: + return + self._frames.put_nowait(frame) + + async def recv(self, decode: bool | None = None) -> str: + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + frame: Final = await self._frames.get() + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hang_up_read_before_the_next_send", [True, False], ids=["read_first", "found_by_a_send"]) +async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_the_provider_ends_it( + hang_up_read_before_the_next_send: bool, +): + """ + Leaving mid-answer must not make a turn free. The provider keeps running and + bills it, so the gateway keeps reading until the turn ends and logs its usage, + whether it reads the hang-up first or finds it by failing to relay a frame. + """ + client_left: Final = asyncio.Event() + client: Final = _walking_out_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4"})], + '"response.created"', + None, + client_left, + hang_up_read_before_the_next_send, + ) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, + _responses_turn_started("gpt-5.4"), + (json.dumps({"type": "response.output_text.delta", "delta": "Hi"}), _responses_turn_ended("gpt-5.4")), + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=5) + + assert [ + [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.output_text.delta", "response.completed"]] + upstream_ws.close.assert_awaited() + + +def _walking_out_client( + frames: Sequence[str], + answered_by: str, + answers_due: Sequence[int] | None, + client_left: asyncio.Event, + hang_up_read_before_the_next_send: bool = True, +) -> MagicMock: + """A turn-taking client that sets client_left as it hangs up, which the gateway either reads at once or only finds + when its next send to the client fails""" + from starlette.websockets import WebSocketDisconnect, WebSocketState + + send_failed: Final = asyncio.Event() + client: Final = _turn_taking_client(frames, answered_by, answers_due) + receive: Final = client.receive + send_text: Final = client.send_text + + async def receive_until_leaving() -> Mapping[str, object]: + message: Final = await receive() + if message["type"] == "websocket.disconnect": + client_left.set() + if not hang_up_read_before_the_next_send: + await send_failed.wait() + client.client_state = WebSocketState.DISCONNECTED + return message + + async def send_text_until_leaving(text: str) -> None: + if client_left.is_set(): + send_failed.set() + raise WebSocketDisconnect(code=1006) + await send_text(text) + + client.receive = receive_until_leaving + client.send_text = AsyncMock(side_effect=send_text_until_leaving) + return client + + +@pytest.mark.asyncio +@pytest.mark.parametrize("turns", [1, 2], ids=["one_turn", "two_turns_asked_at_once"]) +async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_before_the_provider_started_it(turns: int): + """ + A client can ask for turns and leave before the provider says they started. + The provider still runs and bills each one, so the gateway keeps reading + until the last one ends and logs each one's usage. + """ + client_left: Final = asyncio.Event() + client: Final = _walking_out_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4"})] * turns, + '"response.created"', + (0,) * (turns + 1), + client_left, + ) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, None, (_responses_turn_started("gpt-5.4"), _responses_turn_ended("gpt-5.4")) * turns + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=5) + + assert _logged_frame_types(relayed) == [["response.created", "response.completed"]] * turns + upstream_ws.close.assert_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hang_up_read_before_the_next_send", [True, False], ids=["read_first", "found_by_a_send"]) +async def test_websocket_passthrough_bills_a_realtime_turn_a_client_walked_out_of_while_the_provider_stored_its_input( + hang_up_read_before_the_next_send: bool, +): + """ + A realtime client can add an item, ask for a turn, and leave at once. The + provider confirms the item before it starts the turn, and that confirmation + must not end the session while the turn it was asked for is still owed. + """ + client_left: Final = asyncio.Event() + client: Final = _walking_out_client( + [ + json.dumps({"type": "conversation.item.create", "item": {"type": "message", "role": "user"}}), + json.dumps({"type": "response.create"}), + ], + '"session.created"', + (1, 1, 1), + client_left, + hang_up_read_before_the_next_send, + ) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, + None, + ( + json.dumps({"type": "conversation.item.created", "item": {"type": "message", "role": "user"}}), + json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), + REALTIME_RESPONSE_DONE, + ), + opening=(REALTIME_SESSION_CREATED,), + ) + + relayed: Final = await _relay_turns( + client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime", walkout_turn_drain_seconds=5 + ) + + assert _logged_frame_types(relayed) == [ + ["session.created", "conversation.item.created", "response.created", "response.done"] + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_bills_a_realtime_turn_a_client_walked_out_of_after_the_provider_started_one_itself(): + """ + A realtime provider can start a turn on its own when voice detection hears + the user stop. That turn answers no client frame, so a turn the client asks + for next and leaves before it starts is still owed and still billed. + """ + realtime_response_created: Final = json.dumps({"type": "response.created", "response": {"status": "in_progress"}}) + client_left: Final = asyncio.Event() + client: Final = _walking_out_client( + [json.dumps({"type": "response.create"})], '"response.done"', (1, 1), client_left + ) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, + None, + (realtime_response_created, REALTIME_RESPONSE_DONE), + opening=(REALTIME_SESSION_CREATED, realtime_response_created, REALTIME_RESPONSE_DONE), + ) + + relayed: Final = await _relay_turns( + client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime", walkout_turn_drain_seconds=5 + ) + + assert _logged_frame_types(relayed) == [ + ["session.created", "response.created", "response.done"], + ["session.created", "response.created", "response.done"], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "answer", + [ + _responses_turn("gpt-5.4"), + [_responses_turn_ended("gpt-5.4")], + [_responses_error("previous_response_not_found")], + ], + ids=["finished", "finished_with_no_start_frame", "refused"], +) +async def test_websocket_passthrough_lets_go_of_the_provider_at_once_when_a_client_leaves_owing_no_turn( + answer: Sequence[str], +): + """ + Once every turn a client asked for has ended or been refused, nothing more + can be billed, so a client leaving closes the provider socket right away + instead of holding it for the walkout drain. + """ + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4"})], + answered_by=json.dumps(json.loads(answer[-1])["type"]), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[answer]) + + relayed: Final = await _relay_turns( + client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=WALKOUT_TURN_DRAIN_SECONDS + ) + + upstream_ws.close.assert_awaited() + assert len(relayed.success_handler.call_args_list) == 1 + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_no_extra_row_for_what_the_provider_sends_after_the_last_turn(): + """ + A realtime provider follows each finished turn with frames that start no + turn, such as rate_limits.updated. The session ends with nothing running, + so its turn is its only row and the request count stays one per turn. + """ + client: Final = _turn_taking_client([json.dumps({"type": "response.create"})], answered_by='"rate_limits.updated"') + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[[REALTIME_RESPONSE_DONE, json.dumps({"type": "rate_limits.updated", "rate_limits": []})]], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["session.created", "response.done"]] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn(): + """ + A realtime provider names the session's model once, when it opens the + session, and never again. Each later turn logs on its own, so it has to + carry that opening frame or it would be priced as a model nobody named. + """ + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create"}), json.dumps({"type": "response.create"})], + answered_by='"response.done"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, replies=[[REALTIME_RESPONSE_DONE], [REALTIME_RESPONSE_DONE]] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["session.created", "response.done"], ["session.created", "response.done"]] + + +REALTIME_TRANSCRIPTION_COMPLETED: Final = json.dumps( + { + "type": "conversation.item.input_audio_transcription.completed", + "item_id": "item_fake", + "transcript": "hello", + "usage": {"type": "duration", "seconds": 2}, + } +) +TRANSCRIPTION_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" + + +def _transcription_session_update(frame_type: str, model: str) -> str: + return json.dumps({"type": frame_type, "session": {"audio": {"input": {"transcription": {"model": model}}}}}) + + +def _logged_frame_types(relayed: SimpleNamespace) -> list[list[str]]: + return [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_bills_a_clip_transcribed_after_the_last_answer(): + client: Final = _turn_taking_client( + [_transcription_session_update("session.update", "gpt-transcribe"), json.dumps({"type": "response.create"})], + answered_by='"type"', + answers_due=(1, 2, 4), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [_transcription_session_update("session.updated", "gpt-transcribe")], + [REALTIME_RESPONSE_DONE, REALTIME_TRANSCRIPTION_COMPLETED], + ], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert _logged_frame_types(relayed) == [ + ["session.created", "session.updated", "response.done"], + ["session.created", "session.updated", TRANSCRIPTION_COMPLETED], + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_bills_a_clip_transcribed_during_an_answer_in_that_answer_row(): + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create"})], answered_by='"type"', answers_due=(1, 4) + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [ + json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), + REALTIME_TRANSCRIPTION_COMPLETED, + REALTIME_RESPONSE_DONE, + ] + ], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert _logged_frame_types(relayed) == [ + ["session.created", "response.created", TRANSCRIPTION_COMPLETED, "response.done"] + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_carries_only_the_latest_session_update_into_later_rows(): + client: Final = _turn_taking_client( + [ + _transcription_session_update("session.update", "gpt-transcribe-a"), + json.dumps({"type": "response.create"}), + _transcription_session_update("session.update", "gpt-transcribe-b"), + json.dumps({"type": "response.create"}), + ], + answered_by='"type"', + answers_due=(1, 2, 3, 4, 5), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [_transcription_session_update("session.updated", "gpt-transcribe-a")], + [REALTIME_RESPONSE_DONE], + [_transcription_session_update("session.updated", "gpt-transcribe-b")], + [REALTIME_RESPONSE_DONE], + ], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + last_row: Final = relayed.success_handler.call_args_list[-1].kwargs["response_body"] + assert [frame for frame in last_row if frame["type"] == "session.updated"] == [ + json.loads(_transcription_session_update("session.updated", "gpt-transcribe-b")) + ] + assert len(relayed.success_handler.call_args_list) == 2 + + +@pytest.mark.asyncio +async def test_websocket_passthrough_keeps_the_session_update_a_clip_ran_under_until_its_answer_is_billed(): + client: Final = _turn_taking_client( + [ + _transcription_session_update("session.update", "gpt-transcribe-a"), + json.dumps({"type": "response.create"}), + _transcription_session_update("session.update", "gpt-transcribe-b"), + ], + answered_by='"type"', + answers_due=(1, 2, 4, 6), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [_transcription_session_update("session.updated", "gpt-transcribe-a")], + [ + json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), + REALTIME_TRANSCRIPTION_COMPLETED, + ], + [_transcription_session_update("session.updated", "gpt-transcribe-b"), REALTIME_RESPONSE_DONE], + ], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert _logged_frame_types(relayed) == [ + [ + "session.created", + "session.updated", + "response.created", + TRANSCRIPTION_COMPLETED, + "session.updated", + "response.done", + ] + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "later_frame", + [ + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + '{"type": "response.create", "model": "gpt-5.4-mini", "metadata": ' + "[" * 100_000 + "]" * 100_000 + "}", + ( + '{"type": "response.create", "model": "gpt-5.4-mini", "seed": ' + + "9" * (sys.get_int_max_str_digits() + 1) + + "}" + ), + ], + ids=["names_a_denied_model", "nests_too_deep_to_parse", "holds_a_number_too_long_to_parse"], +) +async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_later_frame_is_refused(later_frame: str): + """ + A refused frame closes the client while the provider may still be sending the + answer to an earlier, allowed frame. That answer cannot go to a closed client, + but the provider bills it once it ends, so the gateway keeps reading until it + ends and its usage reaches spend. + """ + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), later_frame], + answered_by='"response.created"', + ) + refused: Final = asyncio.Event() + close_client: Final = client.close.side_effect + + def close_client_and_note_the_refusal(*args: object, **kwargs: object) -> None: + close_client(*args, **kwargs) + refused.set() + + client.close = AsyncMock(side_effect=close_client_and_note_the_refusal) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + refused, _responses_turn_started("gpt-5.4-mini"), (_responses_turn_ended("gpt-5.4-mini"),) + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate, walkout_turn_drain_seconds=5) + + assert [json.loads(text.args[0])["type"] for text in client.send_text.await_args_list] == [ + "response.created", + "error", + ] + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.completed"]] + relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() + + DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" DEEPGRAM_INTERIM_FRAME = json.dumps( { @@ -8962,6 +10461,32 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) +async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object]: + served: Final = await _relay_client_frames( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], + gate, + UserAPIKeyAuth(api_key="hashed-key"), + ) + return served.success_handler.call_args.kwargs["litellm_params"]["metadata"] + + +@pytest.mark.asyncio +async def test_websocket_session_is_charged_to_the_group_its_frames_were_authorized_through(): + """A websocket names its model in a frame, long after the relay built its logging metadata, so the + groups the gate matched have to be picked up at the end or the group pays for nothing it allowed.""" + metadata: Final = await _billed_metadata(RecordingFrameModelGate(groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",)}))) + + assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] + assert metadata["user_api_key"] == "hashed-key" + + +@pytest.mark.asyncio +async def test_websocket_session_that_matched_no_group_is_charged_to_none(): + metadata: Final = await _billed_metadata(RecordingFrameModelGate()) + + assert not metadata[MODEL_ACCESS_GROUP_METADATA_KEY] + + _GEMINI_FLASH_TARGET: Final = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash:generateContent" _VERTEX_FLASH_TARGET: Final = ( "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" @@ -9632,6 +11157,7 @@ PROXY_STAMPED_NAMES: Final = frozenset( "secret_fields", "litellm_trusted_callback_vars", "_litellm_addressed_response_id", + "_litellm_addressed_comparison_response_id", "_litellm_strip_stream_usage", "_litellm_stream_usage_injected", "client_side_timeout", diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index cf00545f331..4934573a777 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -652,6 +652,59 @@ def test_vertex_ai_live_route_without_usage_frames_yields_no_logging_response( assert normalized["kwargs"] == {"model": _LIVE_MODEL} +REALTIME_FRAMES: Final = [ + {"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}}, + { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 33, + "output_tokens": 16, + "total_tokens": 49, + "input_token_details": {"text_tokens": 33, "audio_tokens": 0, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 16, "audio_tokens": 0}, + } + }, + }, +] + + +def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict: + logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id="call-fake") + return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=MagicMock(spec=httpx.Response), + response_body=response_body, + request_body={}, + logging_obj=logging_obj, + url_route=url_route, + result="websocket_connection_successful", + start_time=_START, + end_time=_END, + cache_hit=False, + litellm_call_id="call-fake", + ) + + +@pytest.mark.parametrize("url_route", ["/openai/v1/realtime", "/openai_passthrough/v1/realtime"]) +def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route): + normalized: Final = _normalize(url_route, REALTIME_FRAMES) + + assert normalized["kwargs"]["model"] == "gpt-realtime-2.1" + assert normalized["kwargs"]["custom_llm_provider"] == "openai" + prices: Final = litellm.model_cost["gpt-realtime-2.1"] + assert normalized["kwargs"]["response_cost"] == pytest.approx( + 33 * prices["input_cost_per_token"] + 16 * prices["output_cost_per_token"] + ) + assert normalized["standard_logging_response_object"].usage.total_tokens == 49 + + +def test_openai_http_passthrough_bodies_are_left_to_the_http_handlers(): + normalized: Final = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}}) + + assert normalized["standard_logging_response_object"] is None + assert normalized["kwargs"] == {"litellm_call_id": "call-fake"} + + @pytest.mark.asyncio @pytest.mark.parametrize("cancel_mode", ["once", "twice", "scope", "none"]) @pytest.mark.parametrize("cancel_phase", ["allocation", "lock"]) diff --git a/tests/unit/proxy/proxy_server/conftest.py b/tests/unit/proxy/proxy_server/conftest.py index 6cd9f1cd0ba..fdd251968e5 100644 --- a/tests/unit/proxy/proxy_server/conftest.py +++ b/tests/unit/proxy/proxy_server/conftest.py @@ -17,6 +17,8 @@ from unittest.mock import AsyncMock, MagicMock import pytest if TYPE_CHECKING: + import respx + from litellm import Router from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs @@ -563,3 +565,72 @@ def openai_only_router(monkeypatch: pytest.MonkeyPatch) -> Router: monkeypatch.setattr(proxy_server, "general_settings", {}) monkeypatch.setattr(proxy_server, "user_model", None) return router + + +@pytest.fixture +def failed_spend_row( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> Iterator[Callable[[], tuple[str, str]]]: + """Route tts-1, whisper-1 and omni-moderation to an OpenAI deployment whose provider answers 400, + and return a reader for the (status, call_type) of the spend row the failure wrote.""" + import litellm + from litellm import Router + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.proxy.utils import ProxyLogging + + api_base: Final = "https://failing-provider.example/v1" + respx_mock.post(url__startswith=api_base).respond( + 400, json={"error": {"message": "rejected by the provider", "type": "invalid_request_error"}} + ) + router: Final = Router( + model_list=[ + { + "model_name": model_name, + "litellm_params": { + "model": f"openai/{provider_model}", + "api_base": api_base, + "api_key": "provider-key", + "max_retries": 0, + }, + } + for model_name, provider_model in ( + ("tts-1", "tts-1"), + ("whisper-1", "whisper-1"), + ("omni-moderation", "omni-moderation-latest"), + ) + ], + num_retries=0, + ) + cache: Final = UserApiKeyCache() + spend_writer: Final = MagicMock(update_database=AsyncMock()) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger(spend_writer=lambda: spend_writer)]) + for name, value in ( + ("llm_router", router), + ("llm_model_list", router.model_list), + ("prisma_client", None), + ("general_settings", {}), + ("user_model", None), + ("proxy_config", ProxyConfig()), + ("proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)), + ("user_api_key_cache", cache), + ("version", "test-version"), + ): + monkeypatch.setattr(proxy_server, name, value) + + def last_row() -> tuple[str, str]: + db_call: Final = spend_writer.update_database.call_args.kwargs + payload: Final = get_logging_payload( + kwargs=db_call["kwargs"], + response_obj=db_call["completion_response"], + start_time=db_call["start_time"], + end_time=db_call["end_time"], + ) + return payload["status"], payload["call_type"] + + yield last_row + ProxyLogging._callback_capabilities_cache.clear() diff --git a/tests/unit/proxy/proxy_server/test_exception_handlers.py b/tests/unit/proxy/proxy_server/test_exception_handlers.py index 0f9d379afa1..341875c7cca 100644 --- a/tests/unit/proxy/proxy_server/test_exception_handlers.py +++ b/tests/unit/proxy/proxy_server/test_exception_handlers.py @@ -36,9 +36,9 @@ from litellm.proxy.proxy_server import ( from .conftest import normalize -def _make_request(parent_otel_span=None, path="/chat/completions"): +def _make_request(parent_otel_span=None, path="/chat/completions", headers=()): return Request({ - "type": "http", "method": "POST", "path": path, "headers": [], + "type": "http", "method": "POST", "path": path, "headers": list(headers), "state": {"parent_otel_span": parent_otel_span}, }) @@ -168,6 +168,182 @@ async def test_openai_exception_handler_never_fabricates_a_call_id(monkeypatch): assert "litellm_call_id" not in body["error"] +def _model_access_denied_exception() -> ProxyException: + return ProxyException( + message="Team not allowed to access model. Tried to access gpt-5.6", + type="team_model_access_denied", + param="model", + code=403, + ) + + +_MESSAGES_ROUTES: Final = ( + "/v1/messages", + "/v1/messages/count_tokens", + "/claude_code_gateway/v1/messages", + "/claude_code_gateway/v1/messages/count_tokens", +) +_ANTHROPIC_PASS_THROUGH_ROUTES: Final = ( + "/anthropic/v1/messages", + "/anthropic/v1/messages/count_tokens", + "/anthropic/v1/models", +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", (*_MESSAGES_ROUTES, *_ANTHROPIC_PASS_THROUGH_ROUTES)) +async def test_openai_exception_handler_answers_the_anthropic_envelope_on_the_messages_routes(path: str) -> None: + response: Final = await openai_exception_handler( + request=_make_request(path=path), exc=_model_access_denied_exception() + ) + body: Final = json.loads(response.body) + + assert response.status_code == 403 + assert body == { + "type": "error", + "error": { + "type": "permission_error", + "message": "Team not allowed to access model. Tried to access gpt-5.6", + "litellm_error_type": "team_model_access_denied", + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status_code,anthropic_type", + [ + (401, "authentication_error"), + (403, "permission_error"), + (404, "not_found_error"), + (429, "rate_limit_error"), + (402, "billing_error"), + (409, "conflict_error"), + (504, "timeout_error"), + (405, "invalid_request_error"), + (408, "api_error"), + (502, "api_error"), + ], +) +async def test_openai_exception_handler_maps_messages_failures_to_anthropic_error_types( + status_code: int, anthropic_type: str +) -> None: + exc: Final = ProxyException(message="denied", type="key_model_access_denied", param="model", code=status_code) + + response: Final = await openai_exception_handler(request=_make_request(path="/v1/messages"), exc=exc) + + assert response.status_code == status_code + assert json.loads(response.body)["error"]["type"] == anthropic_type + + +@pytest.mark.asyncio +async def test_openai_exception_handler_keeps_the_status_of_a_budget_failure_on_messages() -> None: + exc: Final = ProxyException(message="Budget exceeded", type="budget_exceeded", param=None, code=422) + + response: Final = await openai_exception_handler( + request=_make_request(path="/v1/messages/count_tokens"), exc=exc + ) + body: Final = json.loads(response.body) + + assert response.status_code == 422 + assert body["type"] == "error" + assert body["error"]["type"] == "invalid_request_error" + assert body["error"]["message"] == "Budget exceeded" + assert body["error"]["litellm_error_type"] == "budget_exceeded" + assert "param" not in body["error"] and "code" not in body["error"] + + +@pytest.mark.asyncio +async def test_openai_exception_handler_keeps_the_openai_envelope_off_the_messages_routes() -> None: + response: Final = await openai_exception_handler( + request=_make_request(path="/v1/chat/completions"), exc=_model_access_denied_exception() + ) + body: Final = json.loads(response.body) + + assert response.status_code == 403 + assert body == { + "error": { + "message": "Team not allowed to access model. Tried to access gpt-5.6", + "type": "team_model_access_denied", + "param": "model", + "code": "403", + } + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ("/v1/models", "/models", "/claude_code_gateway/v1/models")) +async def test_a_model_list_request_answers_errors_in_the_shape_its_listing_takes(path: str) -> None: + anthropic: Final = await openai_exception_handler( + request=_make_request(path=path, headers=[(b"anthropic-version", b"2023-06-01")]), + exc=_model_access_denied_exception(), + ) + openai: Final = await openai_exception_handler( + request=_make_request(path=path), exc=_model_access_denied_exception() + ) + + assert json.loads(anthropic.body)["error"]["type"] == "permission_error" + assert json.loads(anthropic.body)["type"] == "error" + assert json.loads(openai.body)["error"]["type"] == "team_model_access_denied" + assert "type" not in json.loads(openai.body) + + +@pytest.mark.asyncio +async def test_openai_exception_handler_keeps_the_openai_envelope_on_other_anthropic_routes() -> None: + response: Final = await openai_exception_handler( + request=_make_request(path="/v1/skills"), exc=_model_access_denied_exception() + ) + body: Final = json.loads(response.body) + + assert response.status_code == 403 + assert "type" not in body + assert body["error"]["type"] == "team_model_access_denied" + + +@pytest.mark.asyncio +async def test_openai_exception_handler_messages_route_copies_the_call_id_when_opted_in( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}) + exc: Final = ProxyException( + message="denied", + type="team_model_access_denied", + param="model", + code=403, + headers={"x-litellm-call-id": "call-8302"}, + provider_specific_fields={"guardrail": "keyword-block"}, + ) + + response: Final = await openai_exception_handler(request=_make_request(path="/v1/messages"), exc=exc) + body: Final = json.loads(response.body) + + assert response.headers["x-litellm-call-id"] == "call-8302" + assert body["error"] == { + "type": "permission_error", + "message": "denied", + "provider_specific_fields": {"guardrail": "keyword-block"}, + "litellm_error_type": "team_model_access_denied", + "litellm_call_id": "call-8302", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ("/v1/messages", "/anthropic/v1/messages")) +async def test_a_messages_request_with_no_key_answers_the_anthropic_authentication_error( + client: TestClient, monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + response: Final = client.post(path, json={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}) + body: Final = response.json() + + assert response.status_code == 401 + assert body["type"] == "error" + assert body["error"]["type"] == "authentication_error" + assert "param" not in body["error"] + + # --------------------------------------------------------------------------- # _close_dangling_otel_server_span # --------------------------------------------------------------------------- @@ -467,6 +643,22 @@ async def test_otel_unhandled_exception_handler_answers_a_db_outage_with_503_no_ assert json.loads(response.body) == _DB_OUTAGE_503_BODY +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ("/v1/messages", "/claude_code_gateway/v1/messages", "/anthropic/v1/messages")) +async def test_otel_unhandled_exception_handler_answers_the_anthropic_api_error_on_the_messages_routes( + path: str, +) -> None: + response: Final = await otel_unhandled_exception_handler( + request=_make_request(path=path), exc=RuntimeError("kaboom") + ) + + assert response.status_code == 500 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "api_error", "message": "Internal server error"}, + } + + @pytest.mark.asyncio async def test_otel_unhandled_exception_handler_reraises_proxy_exception_error(): """ProxyException / HTTPException / RequestValidationError are re-raised @@ -541,6 +733,33 @@ async def test_http_exception_on_a_plain_http_request_keeps_the_default_json_bod assert json.loads(bytes(response.body)) == {"detail": "not found"} +@pytest.mark.asyncio +async def test_a_disabled_claude_code_gateway_answers_the_anthropic_not_found_error() -> None: + request: Final = _make_request(path="/claude_code_gateway/v1/messages") + + response: Final = await otlp_http_exception_handler( + request, HTTPException(404, "Claude Code gateway is not enabled") + ) + + assert response is not None + assert response.status_code == 404 + assert json.loads(bytes(response.body)) == { + "type": "error", + "error": {"type": "not_found_error", "message": "Claude Code gateway is not enabled"}, + } + + +def test_a_get_on_messages_answers_the_anthropic_envelope_and_keeps_the_allow_header(client: TestClient) -> None: + response: Final = client.get("/v1/messages") + + assert response.status_code == 405 + assert response.headers["allow"] == "POST" + assert response.json() == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "Method Not Allowed"}, + } + + @pytest.mark.asyncio async def test_http_exception_on_a_websocket_closed_before_accept_sends_nothing_more() -> None: sent: Final[list[Message]] = [] diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index d9ff7fef9f9..712b4134c05 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -9,6 +9,7 @@ Pins covered: from __future__ import annotations import asyncio +import base64 import json import logging import os @@ -25,18 +26,24 @@ from unittest.mock import AsyncMock, MagicMock import httpx import pytest +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.asymmetric import padding, rsa +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey from prisma.errors import DataError from pydantic import JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import turn_on_json +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.proxy._types import CommonProxyErrors, ConfigGeneralSettings from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.config_resolvers.settings_store import SettingsStore from litellm.router_api import ApplyResult, RouterSettings from litellm.types.router import Deployment, LiteLLM_Params from litellm.proxy.proxy_server import ( ProxyConfig, _is_remote_module_url, + _reload_settings_store, _scrub_db_overlay_remote_module_loads, _scrub_guardrail_inner, resolve_complexity_router_plugins, @@ -70,6 +77,15 @@ def _db_model(model_id: str, model_name: str = "gpt-3.5-turbo") -> SimpleNamespa ) +def test_yaml_settings_reload_keeps_db_general_settings_until_sync() -> None: + settings: Final = SettingsStore("general_settings") + settings.apply_db_row("general_settings", {"max_file_downloads_per_minute": 2}) + + _reload_settings_store("general_settings", settings, {}) + + assert settings["max_file_downloads_per_minute"] == 2 + + @pytest.mark.asyncio async def test_proxy_config_adds_and_deletes_stale_deployment(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server @@ -102,7 +118,9 @@ def test_proxy_config_upserts_existing_deployment_without_duplicating_it(monkeyp @pytest.mark.asyncio -async def test_proxy_config_preserves_models_when_config_read_fails(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +async def test_proxy_config_preserves_models_when_config_read_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: from litellm.proxy import proxy_server router = litellm.Router(model_list=[_deployment("config-deployment").to_json(exclude_none=True)]) @@ -135,6 +153,8 @@ async def test_proxy_config_keeps_database_models_and_deletes_stale_models( assert desired == frozenset({"kept-deployment"}) assert router.get_model_ids() == ["kept-deployment"] + + from tests._master_key import MASTER_KEY @@ -150,7 +170,6 @@ async def test_proxy_config_loads_lens_store_from_yaml(tmp_path, monkeypatch) -> assert is_lens_tracing_enabled(settings["tracing"], {}) is True - @pytest.mark.asyncio @pytest.mark.parametrize("shutdown_error", [False, True]) async def test_tracing_config_automatically_exports_spend_without_a_storage_dependency( @@ -176,7 +195,9 @@ async def test_tracing_config_automatically_exports_spend_without_a_storage_depe outcome: Final = pytest.raises(RuntimeError, match="shutdown failure") if shutdown_error else nullcontext() with outcome: async with manage_tracing(enabled=True, client_factory=client): - logger: Final = next(callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter)) + logger: Final = next( + callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter) + ) await logger.async_log_success_event( {"standard_logging_object": {"id": "response-1", "response_cost": 0.25}}, None, None, None ) @@ -2655,6 +2676,213 @@ async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monke assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else []) +@pytest.mark.asyncio +async def test_load_config_refuses_allowed_ips_without_a_license(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + config_file: Final = tmp_path / "allowed_ips.yaml" + config_file.write_text('model_list: []\nlitellm_settings: {}\ngeneral_settings:\n allowed_ips: ["127.0.0.1"]\n') + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + expected_message: Final = ( + "allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment." + ) + with pytest.raises(ValueError, match=f"^{re.escape(expected_message)}$"): + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + +def _signed_test_license(private_key: RSAPrivateKey, max_users: int) -> str: + message: Final = json.dumps({"expiration_date": "9999-12-31", "user_id": "test", "max_users": max_users}).encode() + signature: Final = private_key.sign( + message, + padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH), + hashes.SHA256(), + ) + return base64.b64encode(message + b"." + signature).decode() + + +def _write_config(path: Path, general_settings: dict[str, object], litellm_settings: dict[str, object]) -> str: + path.write_text( + json.dumps({"model_list": [], "litellm_settings": litellm_settings, "general_settings": general_settings}) + ) + return str(path) + + +def _boot_proxy( + monkeypatch: pytest.MonkeyPatch, + public_key: RSAPublicKey, + env_license: str | None, + license_server_urls: list[str], + server_accepted_license: str | None = None, +) -> None: + from litellm.proxy import proxy_server + + def _license_server(request: httpx.Request) -> httpx.Response: + license_server_urls.append(str(request.url)) + return httpx.Response( + 200, json={"verify": server_accepted_license is not None and server_accepted_license in str(request.url)} + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.public_key", public_key) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.airgapped_license_data", None) + monkeypatch.setattr( + "litellm.proxy.proxy_server._license_check.http_handler", + HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_license_server))), + ) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + if env_license is None: + monkeypatch.delenv("LITELLM_LICENSE", raising=False) + else: + monkeypatch.setenv("LITELLM_LICENSE", env_license) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.license_str", env_license) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", proxy_server._license_check.is_premium()) + monkeypatch.setattr( + "litellm.proxy.proxy_server.premium_user_data", proxy_server._license_check.airgapped_license_data + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "gated_config", + [ + {"general_settings": {"allowed_ips": ["127.0.0.1"]}}, + {"general_settings": {"enforced_params": ["user"]}}, + {"litellm_settings": {"callbacks": ["openai_moderations"]}}, + ], +) +async def test_load_config_lets_a_config_license_unlock_its_enterprise_gates( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, gated_config: dict[str, dict[str, list[str]]] +): + from litellm.proxy import proxy_server + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + license_key: Final = _signed_test_license(private_key, max_users=10) + config_file: Final = _write_config( + tmp_path / "licensed.yaml", + {"litellm_license": license_key, **gated_config.get("general_settings", {})}, + dict(gated_config.get("litellm_settings", {})), + ) + _boot_proxy(monkeypatch, private_key.public_key(), env_license=None, license_server_urls=[]) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server.premium_user is True + + +@pytest.mark.asyncio +async def test_load_config_keeps_a_valid_env_license_when_the_config_license_is_invalid( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config( + tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license", "allowed_ips": ["127.0.0.1"]}, {} + ) + _boot_proxy( + monkeypatch, + private_key.public_key(), + env_license=_signed_test_license(private_key, max_users=5), + license_server_urls=[], + ) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server.premium_user is True + assert proxy_server._license_check.user_seat_limit() == 5 + + +@pytest.mark.asyncio +async def test_load_config_prefers_the_config_license_when_both_licenses_are_valid( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_license_endpoint + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config( + tmp_path / "licensed.yaml", {"litellm_license": _signed_test_license(private_key, max_users=100)}, {} + ) + _boot_proxy( + monkeypatch, + private_key.public_key(), + env_license=_signed_test_license(private_key, max_users=5), + license_server_urls=[], + ) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server._license_check.user_seat_limit() == 100 + assert (await health_license_endpoint(user_api_key_dict=UserAPIKeyAuth()))["limits"]["max_users"] == 100 + + +@pytest.mark.asyncio +async def test_load_config_still_reports_an_invalid_config_license_as_configured( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_license_endpoint + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config(tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license"}, {}) + _boot_proxy(monkeypatch, private_key.public_key(), env_license=None, license_server_urls=[]) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + license_status: Final = await health_license_endpoint(user_api_key_dict=UserAPIKeyAuth()) + assert (license_status["has_license"], license_status["license_type"]) == (True, "community") + + +@pytest.mark.asyncio +async def test_load_config_does_not_recheck_an_env_license_that_failed_at_boot( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config(tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license"}, {}) + license_server_urls: Final[list[str]] = [] + _boot_proxy(monkeypatch, private_key.public_key(), "also-not-a-valid-license", license_server_urls) + asked_at_boot: Final = len(license_server_urls) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server.premium_user is False + assert [url for url in license_server_urls[asked_at_boot:] if "also-not-a-valid-license" in url] == [] + assert any("not-a-valid-license" in url for url in license_server_urls[asked_at_boot:]) + + +@pytest.mark.asyncio +async def test_load_config_keeps_the_boot_verdict_of_a_server_checked_env_license( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config( + tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license", "allowed_ips": ["127.0.0.1"]}, {} + ) + license_server_urls: Final[list[str]] = [] + _boot_proxy( + monkeypatch, + private_key.public_key(), + "server-checked-license", + license_server_urls, + server_accepted_license="server-checked-license", + ) + asked_at_boot: Final = len(license_server_urls) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server.premium_user is True + assert proxy_server._license_check.license_str == "server-checked-license" + assert [url for url in license_server_urls[asked_at_boot:] if "server-checked-license" in url] == [] + + @pytest.mark.asyncio @pytest.mark.parametrize( "setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "claude_code_marketplace_auth_required"] @@ -2693,7 +2921,6 @@ async def test_load_config_rejects_non_boolean_deny_by_default( await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) - @pytest.mark.asyncio @pytest.mark.parametrize( "yaml_value", ["files-api-2025-04-14", "[files-api-2025-04-14, 7]"], ids=["bare-string", "non-string-entry"] @@ -2713,6 +2940,7 @@ async def test_load_config_rejects_a_malformed_vertex_drop_anthropic_betas_list( with pytest.raises(ValidationError, match="vertex_ai_pass_through_drop_anthropic_betas"): await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + @pytest.mark.asyncio async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch): """Regression: router_settings.plugins dotted-path strings must be resolved to @@ -4470,6 +4698,86 @@ async def test_ProxyConfig__update_general_settings_updates_health_check_retenti reschedule.assert_awaited_once() +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_ignores_a_stored_key_no_settings_api_manages_and_warns_once( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + pc: Final = ProxyConfig() + stored_row: Final = {"enforce_user_param": True, "store_prompts_in_spend_logs": True} + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await pc._update_general_settings(stored_row) + await pc._update_general_settings(stored_row) + from litellm.proxy import proxy_server + + assert proxy_server.general_settings["store_prompts_in_spend_logs"] is True + assert "enforce_user_param" not in proxy_server.general_settings + warnings: Final = [record.getMessage() for record in caplog.records if "enforce_user_param" in record.getMessage()] + assert warnings == [ + "general_settings.enforce_user_param is stored in the database, but the proxy does not apply that stored " + "value, so it is ignored. Set it in the config file to use it, or remove it with POST /config/field/delete." + ] + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_warns_about_an_ignored_stored_key_the_config_file_also_sets( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {"enforce_user_param": False}}) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await pc._update_general_settings({"enforce_user_param": True}) + from litellm.proxy import proxy_server + + assert proxy_server.general_settings["enforce_user_param"] is False + warnings: Final = [record.getMessage() for record in caplog.records if "enforce_user_param" in record.getMessage()] + assert warnings == [ + "general_settings.enforce_user_param is stored in the database, but the proxy does not apply that stored " + "value, so it is ignored. The config file's value applies. Remove the stored copy with POST " + "/config/field/delete." + ] + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_does_not_warn_about_a_stored_copy_of_an_admin_ui_field( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await pc._update_general_settings({"budget_rollover": True, "enforce_user_param": True}) + + warned: Final = [record.getMessage().split(" ")[0] for record in caplog.records if "is ignored" in record.getMessage()] + assert warned == ["general_settings.enforce_user_param"] + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_warns_once_each_time_the_stored_row_stops_being_an_object( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await pc._update_general_settings(None) + await pc._update_general_settings('{"store_prompts_in_spend_logs": true}') + await pc._update_general_settings('{"store_prompts_in_spend_logs": true}') + await pc._update_general_settings({"store_prompts_in_spend_logs": True}) + await pc._update_general_settings('{"store_prompts_in_spend_logs": true}') + + warnings: Final = [record.getMessage() for record in caplog.records if "JSON object" in record.getMessage()] + assert warnings == [ + "general_settings is stored in the database as something other than a JSON object, so the proxy ignores " + "every stored general setting." + ] * 2 + + def _paused_scheduler(monkeypatch): from apscheduler.schedulers.asyncio import AsyncIOScheduler @@ -4918,9 +5226,6 @@ async def test_ProxyConfig__update_general_settings_none_input_noop(): # None input returns early. result = await pc._update_general_settings(db_general_settings=None) assert result is None - # Error-style: dict() will fail on non-mapping non-None input. - with pytest.raises(TypeError): - await pc._update_general_settings(db_general_settings=12345) # type: ignore[arg-type] @pytest.mark.asyncio @@ -5080,7 +5385,9 @@ async def test_ProxyConfig__update_config_from_db_keeps_keys_the_config_file_omi @pytest.mark.asyncio @pytest.mark.parametrize("ui_settings_already_synced", [False, True]) -async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints(monkeypatch, ui_settings_already_synced): +async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints( + monkeypatch: pytest.MonkeyPatch, ui_settings_already_synced: bool +) -> None: from litellm.proxy import proxy_server pc = ProxyConfig() @@ -5106,6 +5413,32 @@ async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endp assert settings_refresh.await_count == (0 if ui_settings_already_synced else 1) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stored_row, rebuilds_routes", + [({"pass_through_endpoints": []}, True), ('{"pass_through_endpoints": []}', False), ([], False)], + ids=["object", "json-string", "array"], +) +async def test_ProxyConfig_add_deployment_rebuilds_pass_through_routes_only_from_an_object_general_settings_row( + monkeypatch: pytest.MonkeyPatch, stored_row: object, rebuilds_routes: bool +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.pass_through_endpoints import pass_through_endpoints + + pc: Final = ProxyConfig() + rebuild_routes: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock()) + monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=SimpleNamespace(param_value=stored_row))) + monkeypatch.setattr(pass_through_endpoints, "initialize_pass_through_endpoints_in_db", rebuild_routes) + monkeypatch.setattr(pc, "_should_load_db_object", lambda *, object_type: object_type == "pass_through_endpoints") + monkeypatch.setattr(pc, "get_credentials", AsyncMock()) + + await pc.add_deployment(prisma_client=MagicMock(), proxy_logging_obj=MagicMock(), ui_settings_already_synced=True) + + assert rebuild_routes.await_count == (1 if rebuilds_routes else 0) + + # --------------------------------------------------------------------------- # ProxyConfig._update_config_from_db # --------------------------------------------------------------------------- @@ -5285,9 +5618,7 @@ def revision_test_environment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv(key) -def _environment_variables_row( - values: Mapping[str, str | None] | None, revision: int | None -) -> SimpleNamespace | None: +def _environment_variables_row(values: Mapping[str, str | None] | None, revision: int | None) -> SimpleNamespace | None: if values is None: return None return SimpleNamespace( @@ -5473,9 +5804,7 @@ async def test_environment_variables_row_read_before_an_overlapping_newer_sync_i await newer_sync_finished.wait() older_sync: Final = asyncio.create_task( - _sync_from( - pc, _config_database(read_older_row, AsyncMock(return_value=older_row), wait_for_the_newer_sync) - ) + _sync_from(pc, _config_database(read_older_row, AsyncMock(return_value=older_row), wait_for_the_newer_sync)) ) await older_environment_row_read.wait() await _sync_from(pc, _rows_database(newer_row, newer_row)) diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index 59d9490f786..9fa617c0dc3 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -9,18 +9,37 @@ Pins (PR2): from __future__ import annotations +import contextlib import io +import json +import math +import wave +from collections.abc import AsyncIterator, Callable, Mapping +from contextlib import AbstractContextManager +from typing import Final from unittest.mock import AsyncMock, MagicMock import httpx import pytest +import respx from fastapi import HTTPException +from fastapi.testclient import TestClient +from openai import AsyncOpenAI +from starlette.datastructures import FormData, UploadFile +import litellm from litellm import Router +from litellm.caching.caching import Cache +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy import proxy_server from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.route_llm_request import route_request +from litellm.proxy.utils import _failure_fields_to_lift +from litellm.types.guardrails import GuardrailEventHooks, GuardrailStreamScope from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep from litellm.types.router import RouterErrors +from litellm.types.utils import TranscriptionResponse _TAG_ROUTING_DENIAL = f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model=tts-1 and tags=['x']" @@ -125,9 +144,7 @@ def patched_transcription(monkeypatch): return data monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data) - monkeypatch.setattr( - proxy_server, "check_file_size_under_limit", lambda **kwargs: True - ) + monkeypatch.setattr(proxy_server, "check_file_size_under_limit", lambda **kwargs: True) async def _form_data(request): from starlette.datastructures import FormData, UploadFile @@ -159,6 +176,47 @@ def patched_transcription_error(monkeypatch, patched_transcription): yield +@pytest.fixture +def patched_transcription_stream(monkeypatch, patched_transcription): + class _FakeEvent: + def model_dump_json(self, exclude_unset: bool = False) -> str: + return '{"type":"transcript.text.done","text":"hello world"}' + + class _FakeAsyncStream: + def __init__(self): + self.closed = False + + def __aiter__(self): + async def _events(): + yield _FakeEvent() + + return _events() + + async def aclose(self): + self.closed = True + + async def _form_data(request): + from starlette.datastructures import FormData, UploadFile + + upload = UploadFile( + filename="audio.mp3", + file=io.BytesIO(b"\x00\x01\x02"), + ) + return FormData([("file", upload), ("model", "gpt-transcribe"), ("stream", "true")]) + + stream = _FakeAsyncStream() + + async def _llm_call(): + return stream + + async def _fake_route_request(*args, **kwargs): + return _llm_call() + + monkeypatch.setattr(proxy_server, "get_form_data", _form_data) + monkeypatch.setattr(proxy_server, "route_request", _fake_route_request) + yield stream + + @pytest.mark.parametrize("path", ["/v1/audio/speech", "/audio/speech"]) def test_audio_speech_happy_path(client, auth_as, patched_speech, path): """Pins ``POST /v1/audio/speech`` and ``POST /audio/speech`` (happy).""" @@ -359,3 +417,663 @@ def test_audio_transcription_error(client, auth_as, patched_transcription_error, response = client.post(path, files=files, data=data) assert response.status_code == 500 assert len(response.content) > 0 + + +@pytest.mark.parametrize("path", ["/v1/audio/transcriptions", "/audio/transcriptions"]) +def test_audio_transcription_stream_returns_sse(client, auth_as, patched_transcription_stream, path): + files = {"file": ("audio.mp3", b"\x00\x01\x02", "audio/mpeg")} + data = {"model": "gpt-transcribe", "stream": "true"} + with auth_as(): + response = client.post(path, files=files, data=data) + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert response.text == 'data: {"type":"transcript.text.done","text":"hello world"}\n\ndata: [DONE]\n\n' + assert patched_transcription_stream.closed is True + + +@pytest.mark.usefixtures("patched_transcription_stream") +@pytest.mark.parametrize("configuration", ["default", "policy"]) +def test_streaming_transcription_keeps_streaming_with_output_guardrails( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + configuration: str, +) -> None: + guardrail: Final = CustomGuardrail( + guardrail_name="transcription-output", + event_hook=GuardrailEventHooks.post_call, + default_on=configuration == "default", + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + if configuration == "policy": + pipeline: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="transcription-output")]) + proxy_server.proxy_logging_obj.pre_call_hook.side_effect = lambda **kwargs: { + **kwargs["data"], + "metadata": {"_guardrail_pipelines": [("transcription-policy", pipeline)]}, + } + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + ) + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert '"text":"hello world"' in response.text + + +@pytest.mark.usefixtures("patched_transcription_stream") +def test_streamed_transcription_returns_the_litellm_call_and_model_headers( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + upstream: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, + content=b'data: {"type":"transcript.text.done","text":"hello world"}\n\n', + headers={"content-type": "text/event-stream"}, + ) + ) + ), + ) + + async def route_to_the_sdk( + data: Mapping[str, object], route_type: str, llm_router: object, user_model: object + ) -> object: + return litellm.atranscription( + model="openai/gpt-transcribe", + file=data["file"], + stream=True, + client=upstream, + litellm_call_id=data["litellm_call_id"], + model_info={"id": "transcription-deployment"}, + ) + + monkeypatch.setattr(proxy_server, "route_request", route_to_the_sdk) + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + headers={"x-litellm-call-id": "streamed-transcription-call"}, + ) + + assert response.status_code == 200 + assert '"text":"hello world"' in response.text + assert response.headers["x-litellm-call-id"] == "streamed-transcription-call" + assert response.headers["x-litellm-model-id"] == "transcription-deployment" + assert [name for name in response.headers if name.startswith("x-litellm-response-cost")] == [] + proxy_server.proxy_logging_obj.update_request_status.assert_called_once_with( + litellm_call_id="streamed-transcription-call", status="success" + ) + + +class _TranscriptRedactor(CustomGuardrail): + def __init__(self, stream_scope: GuardrailStreamScope = "both") -> None: + super().__init__( + guardrail_name="transcript-redactor", event_hook=GuardrailEventHooks.post_call, stream_scope=stream_scope + ) + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: object, response: object + ) -> None: + if isinstance(response, TranscriptionResponse): + response.text = "[REDACTED]" + + +def _streams_when_asked(request: httpx.Request) -> httpx.Response: + if b'name="stream"' in request.content: + return httpx.Response( + 200, + content=b'data: {"type":"transcript.text.done","text":"my pin is 1234"}\n\n', + headers={"content-type": "text/event-stream"}, + ) + return httpx.Response(200, json={"text": "my pin is 1234", "usage": {"type": "duration", "seconds": 2}}) + + +def _sse_events(body: str) -> tuple[object, ...]: + assert body.endswith("data: [DONE]\n\n") + return tuple(json.loads(frame.removeprefix("data: ")) for frame in body.removesuffix("data: [DONE]\n\n").split("\n\n") if frame) + + +@pytest.mark.usefixtures("patched_transcription_stream") +@pytest.mark.parametrize("deployment_guardrails", (["transcript-redactor"], None)) +def test_streamed_transcription_honors_a_deployment_post_call_guardrail( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + deployment_guardrails: list[str] | None, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [_TranscriptRedactor()]) + respx_mock.post("https://transcribe-upstream.example/v1/audio/transcriptions").mock(side_effect=_streams_when_asked) + deployment_params: Final = { + "model": "openai/gpt-transcribe", + "api_key": "sk-test", + "api_base": "https://transcribe-upstream.example/v1", + } + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-transcribe", + "litellm_params": ( + {**deployment_params, "guardrails": deployment_guardrails} + if deployment_guardrails + else deployment_params + ), + } + ] + ), + ) + monkeypatch.setattr(proxy_server, "route_request", route_request) + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + ) + + assert response.status_code == 200 + if deployment_guardrails: + assert response.headers["content-type"].startswith("text/event-stream") + assert _sse_events(response.text) == ( + {"type": "transcript.text.delta", "delta": "[REDACTED]"}, + {"type": "transcript.text.done", "text": "[REDACTED]", "usage": {"type": "duration", "seconds": 2}}, + ) + else: + assert response.headers["content-type"].startswith("text/event-stream") + assert '"text":"my pin is 1234"' in response.text + + +@pytest.mark.usefixtures("patched_transcription_stream") +@pytest.mark.parametrize( + "stream_scope,upstream_streamed,expected_text", + [("streaming", False, "[REDACTED]"), ("non_streaming", True, "my pin is 1234")], +) +def test_streamed_transcription_applies_a_deployment_guardrail_only_in_its_stream_scope( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + stream_scope: GuardrailStreamScope, + upstream_streamed: bool, + expected_text: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [_TranscriptRedactor(stream_scope=stream_scope)]) + upstream_route: Final = respx_mock.post("https://transcribe-upstream.example/v1/audio/transcriptions").mock( + side_effect=_streams_when_asked + ) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-transcribe", + "litellm_params": { + "model": "openai/gpt-transcribe", + "api_key": "sk-test", + "api_base": "https://transcribe-upstream.example/v1", + "guardrails": ["transcript-redactor"], + }, + } + ] + ), + ) + monkeypatch.setattr(proxy_server, "route_request", route_request) + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + ) + + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream") + assert (b'name="stream"' in upstream_route.calls.last.request.content) is upstream_streamed + assert _sse_events(response.text)[-1]["text"] == expected_text + + +async def _form_without_stream_field(request: object) -> FormData: + return FormData( + [("file", UploadFile(filename="sample.wav", file=io.BytesIO(b"audio"))), ("model", "gpt-transcribe")] + ) + + +@pytest.mark.usefixtures("patched_transcription") +@pytest.mark.parametrize("deployment_guardrails", (["transcript-redactor"], None)) +def test_a_deployment_that_streams_transcription_answers_sse_to_a_client_that_sent_no_stream_field( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + deployment_guardrails: list[str] | None, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [_TranscriptRedactor()]) + monkeypatch.setattr(proxy_server, "get_form_data", _form_without_stream_field) + respx_mock.post("https://transcribe-upstream.example/v1/audio/transcriptions").mock(side_effect=_streams_when_asked) + deployment_params: Final = { + "model": "openai/gpt-transcribe", + "api_key": "sk-test", + "api_base": "https://transcribe-upstream.example/v1", + "stream": True, + } + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-transcribe", + "litellm_params": ( + {**deployment_params, "guardrails": deployment_guardrails} + if deployment_guardrails + else deployment_params + ), + } + ] + ), + ) + monkeypatch.setattr(proxy_server, "route_request", route_request) + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe"}, + ) + + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream") + assert _sse_events(response.text)[-1]["text"] == ("[REDACTED]" if deployment_guardrails else "my pin is 1234") + header_hook_response: Final = proxy_server.proxy_logging_obj.post_call_response_headers_hook.await_args.kwargs[ + "response" + ] + if deployment_guardrails: + assert isinstance(header_hook_response, TranscriptionResponse) + assert header_hook_response.text == "[REDACTED]" + + +@pytest.mark.usefixtures("patched_transcription_stream") +def test_cached_guarded_transcription_still_streams_to_a_stream_client( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [_TranscriptRedactor()]) + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + upstream_route: Final = respx_mock.post("https://transcribe-upstream.example/v1/audio/transcriptions").mock( + side_effect=_streams_when_asked + ) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-transcribe", + "litellm_params": { + "model": "openai/gpt-transcribe", + "api_key": "sk-test", + "api_base": "https://transcribe-upstream.example/v1", + "guardrails": ["transcript-redactor"], + }, + } + ], + cache_responses=True, + ), + ) + monkeypatch.setattr(proxy_server, "route_request", route_request) + + with auth_as(): + responses: Final = [ + client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + ) + for _ in range(2) + ] + + assert upstream_route.call_count == 1 + assert [response.headers["content-type"].split(";")[0] for response in responses] == ["text/event-stream"] * 2 + assert _sse_events(responses[1].text)[-1] == { + "type": "transcript.text.done", + "text": "[REDACTED]", + "usage": {"type": "duration", "seconds": 2}, + } + + +@pytest.mark.usefixtures("patched_transcription_stream") +def test_streamed_transcription_frames_keep_the_provider_event_shape( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider_events: Final = ( + {"type": "transcript.text.delta", "delta": "hello"}, + {"type": "transcript.text.done", "text": "hello", "usage": {"type": "duration", "seconds": 2}}, + ) + upstream: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, + content="".join(f"data: {json.dumps(event)}\n\n" for event in provider_events).encode(), + headers={"content-type": "text/event-stream"}, + ) + ) + ), + ) + + async def route_to_the_sdk( + data: Mapping[str, object], route_type: str, llm_router: object, user_model: object + ) -> object: + return litellm.atranscription( + model="openai/gpt-transcribe", file=data["file"], stream=data["stream"], client=upstream + ) + + monkeypatch.setattr(proxy_server, "route_request", route_to_the_sdk) + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + ) + + assert response.status_code == 200 + assert _sse_events(response.text) == provider_events + + +class _FirstDeltaThenReset(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b'data: {"type":"transcript.text.delta","delta":"Good"}\n\n' + raise httpx.ReadError("connection reset") + + +def _silent_wav(seconds: int) -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(16_000) + wav.writeframes(b"\x00\x00" * 16_000 * seconds) + return buffer.getvalue() + + +@pytest.mark.usefixtures("patched_transcription") +def test_a_streamed_transcription_that_fails_mid_transcript_bills_the_uploaded_seconds( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + upstream: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_FirstDeltaThenReset() + ) + ) + ), + ) + + async def route_to_the_sdk( + data: Mapping[str, object], route_type: str, llm_router: object, user_model: object + ) -> object: + return litellm.atranscription(model="openai/gpt-transcribe", file=data["file"], stream=True, client=upstream) + + four_second_upload: Final = _silent_wav(4) + + async def stream_form(request: object) -> FormData: + return FormData( + [ + ("file", UploadFile(filename="sample.wav", file=io.BytesIO(four_second_upload))), + ("model", "gpt-transcribe"), + ("stream", "true"), + ] + ) + + monkeypatch.setattr(proxy_server, "get_form_data", stream_form) + monkeypatch.setattr(proxy_server, "route_request", route_to_the_sdk) + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", four_second_upload, "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + ) + + assert response.status_code == 200 + assert '"error"' in response.text + failure_request: Final = proxy_server.proxy_logging_obj.post_call_failure_hook.await_args.kwargs["request_data"] + billed_cost: Final = _failure_fields_to_lift(failure_request)["response_cost"] + assert isinstance(billed_cost, float) + assert math.isclose(billed_cost, 4 * litellm.model_cost["gpt-transcribe"]["input_cost_per_second"]) + + +async def _scribe_stream_form(request: object) -> FormData: + return FormData( + [ + ("file", UploadFile(filename="sample.wav", file=io.BytesIO(b"audio"))), + ("model", "scribe"), + ("stream", "true"), + ] + ) + + +@pytest.mark.usefixtures("patched_transcription") +def test_transcription_stream_form_field_reaches_other_providers_as_sent( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = respx_mock.post("https://elevenlabs.example/v1/speech-to-text").mock( + return_value=httpx.Response(200, json={"text": "hello world", "language_code": "en"}) + ) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "scribe", + "litellm_params": { + "model": "elevenlabs/scribe_v1", + "api_key": "sk-test", + "api_base": "https://elevenlabs.example", + }, + } + ] + ), + ) + monkeypatch.setattr(proxy_server, "route_request", route_request) + monkeypatch.setattr(proxy_server, "get_form_data", _scribe_stream_form) + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "scribe", "stream": "true"}, + ) + + assert response.status_code == 200 + assert response.json()["text"] == "hello world" + assert b'name="stream"\r\n\r\ntrue\r\n' in upstream.calls.last.request.content + + +def _streams_chat_and_transcription(request: httpx.Request) -> httpx.Response: + sse_body: Final = ( + b'data: {"id":"chat-1","object":"chat.completion.chunk","created":0,"model":"gpt-4o",' + b'"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' + if request.url.path.endswith("/chat/completions") + else b'data: {"type":"transcript.text.done","text":"hello world"}\n\n' + ) + return httpx.Response( + 200, + content=sse_body, + headers={"content-type": "text/event-stream", "x-ratelimit-remaining-requests": "99"}, + ) + + +def _litellm_header_names(response: httpx.Response) -> frozenset[str]: + return frozenset( + name + for name in response.headers + if name.startswith("x-litellm-") and not name.startswith("x-litellm-response-cost") + ) + + +def test_streamed_transcription_carries_the_headers_a_streamed_chat_completion_carries( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.route(host="stream-upstream.example").mock(side_effect=_streams_chat_and_transcription) + upstream: Final = {"api_key": "sk-test", "api_base": "https://stream-upstream.example/v1"} + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o", **upstream}}, + {"model_name": "gpt-transcribe", "litellm_params": {"model": "openai/gpt-transcribe", **upstream}}, + ] + ), + ) + + with auth_as(): + chat: Final = client.post( + "/v1/chat/completions", + json={"model": "gpt-4o", "stream": True, "messages": [{"role": "user", "content": "hi"}]}, + ) + transcription: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + ) + + assert chat.status_code == 200 + assert transcription.status_code == 200 + assert transcription.headers["content-type"].startswith("text/event-stream") + assert _litellm_header_names(chat) - _litellm_header_names(transcription) <= {"x-litellm-overhead-duration-ms"} + assert transcription.headers["x-litellm-model-name"] == "openai/gpt-transcribe" + assert transcription.headers["llm_provider-x-ratelimit-remaining-requests"] == "99" + assert [name for name in transcription.headers if name.startswith("x-litellm-response-cost")] == [] + + +class _FailsAfterTwoDeltas(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b'data: {"type":"transcript.text.delta","delta":"hello "}\n\n' + yield b'data: {"type":"transcript.text.delta","delta":"world"}\n\n' + raise httpx.ReadError("upstream connection dropped at /srv/app/stt_stream.py via stt.internal.corp") + + +@pytest.mark.usefixtures("patched_transcription_stream") +def test_streamed_transcription_error_mid_stream_ends_with_an_error_event_and_runs_the_failure_hook( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + upstream: Final = AsyncOpenAI( + api_key="sk-test", + base_url="https://upstream.example/v1", + http_client=httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, stream=_FailsAfterTwoDeltas(), headers={"content-type": "text/event-stream"} + ) + ) + ), + ) + + async def route_to_the_sdk( + data: Mapping[str, object], route_type: str, llm_router: object, user_model: object + ) -> object: + return litellm.atranscription(model="openai/gpt-transcribe", file=data["file"], stream=True, client=upstream) + + failure_hook: Final = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "post_call_failure_hook", failure_hook) + monkeypatch.setattr(proxy_server, "route_request", route_to_the_sdk) + + with auth_as(): + response: Final = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"audio", "audio/wav")}, + data={"model": "gpt-transcribe", "stream": "true"}, + ) + + data_lines: Final = tuple( + line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: ") + ) + frames: Final = tuple(json.loads(line) for line in data_lines[:-1]) + assert response.status_code == 200 + assert tuple(frame.get("delta") for frame in frames[:2]) == ("hello ", "world") + assert len(frames) == 3 + assert data_lines[-1] == "[DONE]" + assert "upstream connection dropped" in frames[2]["error"]["message"] + assert "/srv/app/stt_stream.py" not in frames[2]["error"]["message"] + assert "stt.internal.corp" not in frames[2]["error"]["message"] + assert set(frames[2]["error"]) == {"message", "type", "param", "code"} + assert (frames[2]["error"]["type"], frames[2]["error"]["param"], frames[2]["error"]["code"]) == ( + "internal_server_error", + None, + "500", + ) + failure_hook.assert_awaited_once() + assert isinstance(failure_hook.await_args.kwargs["original_exception"], litellm.APIConnectionError) + + +def _post_speech(client: TestClient) -> httpx.Response: + return client.post("/v1/audio/speech", json={"model": "tts-1", "input": "Hi", "voice": "alloy"}) + + +def _post_transcription(client: TestClient) -> httpx.Response: + return client.post( + "/v1/audio/transcriptions", + files={"file": ("audio.mp3", b"\x00\x01\x02", "audio/mpeg")}, + data={"model": "whisper-1"}, + ) + + +@pytest.mark.parametrize( + ("send", "call_type"), + [(_post_speech, "aspeech"), (_post_transcription, "atranscription")], +) +def test_a_provider_rejection_logs_the_audio_endpoint_call_type( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[object]], + failed_spend_row: Callable[[], tuple[str, str]], + send: Callable[[TestClient], httpx.Response], + call_type: str, +) -> None: + with auth_as(): + response: Final = send(client) + assert response.status_code == 400, response.text + assert failed_spend_row() == ("failure", call_type) diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index bffa4612964..6ad5ef1a4eb 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -19,6 +19,7 @@ from unittest.mock import AsyncMock, MagicMock import httpx import pytest +import respx from fastapi.testclient import TestClient import litellm @@ -576,6 +577,55 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke assert router.registry().names("gpt-4o", kind="wildcard") == frozenset({"openai/gpt-4o"}) +@pytest.mark.parametrize( + ("key_models", "expected_names", "expected_listings"), + [(["qwen-open"], ["qwen-open"], 0), (["hosted_vllm/*"], ["hosted_vllm/qwen2.5"], 1)], + ids=["key-without-the-wildcard", "key-granted-the-wildcard"], +) +def test_v1_model_info_asks_a_wildcard_server_only_for_keys_that_can_call_it( + client, auth_as, monkeypatch, request, key_models, expected_names, expected_listings +): + from litellm.proxy._types import LitellmUserRoles + + wildcard_server: Final = f"http://vllm-{request.node.callspec.id}.test:8000/v1" + other_alias_server: Final = f"http://vllm-secure-{request.node.callspec.id}.test:8000/v1" + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "qwen-open", + "litellm_params": {"model": "hosted_vllm/qwen2.5", "api_base": "http://vllm-open.test:8000/v1"}, + }, + {"model_name": "hosted_vllm/*", "litellm_params": {"model": "hosted_vllm/*", "api_base": wildcard_server}}, + { + "model_name": "vllm-secure/*", + "litellm_params": {"model": "hosted_vllm/*", "api_base": other_alias_server}, + }, + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + + with respx.mock(assert_all_called=False) as upstream: + listing: Final = upstream.get(f"{wildcard_server}/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}]}) + ) + other_alias_listing: Final = upstream.get(f"{other_alias_server}/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "qwen2.5"}]}) + ) + with auth_as(role=LitellmUserRoles.INTERNAL_USER, models=key_models): + response: Final = client.get("/v1/model/info") + + assert response.status_code == 200, response.text + assert ( + [model["model_name"] for model in response.json()["data"]], + listing.call_count, + other_alias_listing.call_count, + ) == (expected_names, expected_listings, 0) + + # --------------------------------------------------------------------------- # GET /model/info — team BYOK scoping (issue #30983) # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index 874f471f3e7..ff74cd3fa1d 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -58,6 +58,7 @@ def patched_models(monkeypatch): router.get_routable_upstream_model = MagicMock(return_value="gpt-4") router.get_configured_display_name = MagicMock(return_value=None) router.get_configured_service_tiers = MagicMock(return_value=()) + router.get_configured_max_input_tokens = MagicMock(return_value=()) router.get_model_group_info = MagicMock(return_value=None) router.get_search_tool_support = MagicMock(return_value=()) route_entry_points_to_stubs(router) diff --git a/tests/unit/proxy/proxy_server/test_routes_moderations.py b/tests/unit/proxy/proxy_server/test_routes_moderations.py index 4553a5e7cf4..47452629ebf 100644 --- a/tests/unit/proxy/proxy_server/test_routes_moderations.py +++ b/tests/unit/proxy/proxy_server/test_routes_moderations.py @@ -7,12 +7,17 @@ Pins (PR2): from __future__ import annotations +import contextlib +from typing import TYPE_CHECKING, Callable, Final from unittest.mock import AsyncMock, MagicMock import pytest from litellm.proxy import proxy_server +if TYPE_CHECKING: + from fastapi.testclient import TestClient + from .conftest import normalize # type: ignore[import-not-found] HAPPY_RESPONSE = { @@ -109,3 +114,14 @@ def test_moderation_error(client, auth_as, moderation_pipeline_raises, path): response = client.post(path, json=payload) assert response.status_code == 500 assert len(response.content) > 0 + + +def test_a_provider_rejection_logs_the_moderation_call_type( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[object]], + failed_spend_row: Callable[[], tuple[str, str]], +) -> None: + with auth_as(): + response: Final = client.post("/v1/moderations", json={"model": "omni-moderation", "input": ["one", "two"]}) + assert response.status_code == 400, response.text + assert failed_spend_row() == ("failure", "amoderation") diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index c37475bd539..80bed119e79 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -5,7 +5,6 @@ Pins covered: - ``async_assistants_data_generator`` - ``_get_client_requested_model_for_streaming`` - ``_restamp_streaming_chunk_model`` -- ``_fast_serialize_simple_model_response_stream`` - ``_serialize_streaming_chunk`` - ``_apply_streaming_chunk_hooks`` - ``_format_streaming_sse_chunk`` @@ -16,8 +15,11 @@ Pins covered: from __future__ import annotations import asyncio +import hashlib import json -from collections.abc import AsyncIterator, Callable +import logging +from collections.abc import AsyncIterator, Callable, Mapping +from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -36,7 +38,6 @@ from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.common_request_processing import create_response from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, - _fast_serialize_simple_model_response_stream, _format_fallback_metadata_sse_event, _format_streaming_sse_chunk, _get_client_requested_model_for_streaming, @@ -55,7 +56,16 @@ from litellm.types.llms.openai import ( ResponseFailedEvent, ResponsesAPIResponse, ) -from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + TextChoices, + TextCompletionResponse, + Usage, +) from tests._support.clock_loops import AutojumpClockLoop from .conftest import normalize, route_entry_points_to_stubs @@ -81,6 +91,10 @@ def _simple_chunk(model: str = "gpt-4", content: str = "hi") -> ModelResponseStr ) +def _sse_text(frame: str | bytes) -> str: + return frame.decode() if isinstance(frame, bytes) else frame + + async def _async_iter(items): for it in items: yield it @@ -520,48 +534,13 @@ def test_format_fallback_metadata_sse_event(): assert isinstance(payload["created"], int) -# --------------------------------------------------------------------------- -# _fast_serialize_simple_model_response_stream -# --------------------------------------------------------------------------- - - -def test_fast_serialize_simple_model_response_stream_returns_bytes_payload(): - chunk = _simple_chunk() - result = _fast_serialize_simple_model_response_stream(chunk) - assert isinstance(result, bytes) - payload = json.loads(result) - assert normalize(payload) == { - "id": "", - "object": "chat.completion.chunk", - "created": "", - "model": "gpt-4", - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": "hi"}, - } - ], - } - - -def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_invalid(): - """Fast path bails (returns None) when ``usage`` is populated — the slow - path is required to preserve usage fields. Returning None here is the - "I cannot handle this" sentinel, not a hard error.""" - chunk = _simple_chunk() - chunk.usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) - assert _fast_serialize_simple_model_response_stream(chunk) is None - - # --------------------------------------------------------------------------- # _serialize_streaming_chunk # --------------------------------------------------------------------------- -def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): - result = _serialize_streaming_chunk(_simple_chunk()) - assert isinstance(result, bytes) - payload = json.loads(result) +def test_serialize_streaming_chunk_simple_payload(): + payload = json.loads(_serialize_streaming_chunk(_simple_chunk())) assert normalize(payload) == { "id": "", "object": "chat.completion.chunk", @@ -576,6 +555,102 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): } +@pytest.mark.parametrize( + ("choice_extras", "serialized_extras"), + [ + ({"stop_reason": ",5", "token_ids": [20]}, {"stop_reason": ",5", "token_ids": [20]}), + ({"finish_reason": "stop", "stop_reason": None}, {"finish_reason": "stop", "stop_reason": None}), + ], + ids=["vllm-fields", "natural-stop"], +) +def test_serialize_streaming_chunk_keeps_the_vllm_choice_fields( + choice_extras: dict[str, object], serialized_extras: dict[str, object] +): + chunk: Final = ModelResponseStream( + id="chatcmpl-test", + choices=[StreamingChoices(index=0, delta=Delta(content="1,2,3,4", role="assistant"), **choice_extras)], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + + assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [ + {"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, **serialized_extras} + ] + + +def _text_completion_chunk(choice_extras: Mapping[str, object]) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", + choices=[TextChoices(finish_reason="length", index=0, text="", **choice_extras)], + created=1, + model="Qwen/Qwen3-0.6B", + ) + + +def _chat_chunk_with_fingerprint(choice_extras: Mapping[str, object]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + choices=[StreamingChoices(finish_reason="length", index=0, delta=Delta(content=None), **choice_extras)], + created=1, + model="Qwen/Qwen3-0.6B", + object="chat.completion.chunk", + system_fingerprint="fp_1", + ) + + +_VLLM_SHAPED_CHUNKS: Final = pytest.mark.parametrize( + "build_chunk", [_text_completion_chunk, _chat_chunk_with_fingerprint], ids=["text-completion", "chat-fingerprint"] +) + + +@_VLLM_SHAPED_CHUNKS +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": None}, {"stop_reason": None, "token_ids": [21]}, {"stop_reason": None, "token_ids": None}], + ids=["null-stop-reason", "null-stop-reason-with-token-ids", "null-stop-reason-and-token-ids"], +) +def test_serialize_streaming_chunk_keeps_null_vllm_choice_fields( + build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream], + vllm_choice_fields: dict[str, object], +) -> None: + chunk: Final = build_chunk(vllm_choice_fields) + plain_chunk: Final = json.loads(build_chunk({}).model_dump_json(exclude_none=True, exclude_unset=True)) + + assert json.loads(_serialize_streaming_chunk(chunk)) == { + **plain_chunk, + "choices": [{**plain_chunk["choices"][0], **vllm_choice_fields}], + } + + +@_VLLM_SHAPED_CHUNKS +@pytest.mark.parametrize("choice_extras", [{}, {"stop_reason": ",5", "token_ids": [20]}], ids=["no-vllm", "set-vllm"]) +def test_serialize_streaming_chunk_without_null_vllm_fields_matches_model_dump_json( + build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream], + choice_extras: dict[str, object], +) -> None: + chunk: Final = build_chunk(choice_extras) + + assert _serialize_streaming_chunk(chunk) == chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() + + +def test_serialize_streaming_chunk_keeps_streaming_a_null_stop_reason_chunk_with_a_value_json_cannot_encode() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-vllm", + choices=[StreamingChoices(index=0, delta=Delta(content="a"), stop_reason=None, token_ids=[21])], + created=0, + model="Qwen/Qwen3-0.6B", + object="chat.completion.chunk", + ) + setattr(chunk, "trace_chunk", hashlib.sha256(b"trace").digest()) + + payload: Final = json.loads(_serialize_streaming_chunk(chunk)) + + assert payload["choices"][0]["delta"] == {"content": "a"} + assert payload["choices"][0]["token_ids"] == [21] + assert "trace_chunk" not in payload + + def test_serialize_streaming_chunk_invalid_input_raises_attribute_error(): """The helper is typed as ``BaseModel`` — handing it a plain dict trips the attribute-access path (no ``model_dump_json``).""" @@ -700,10 +775,7 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): out.append(line) assert out[-1] == "data: [DONE]\n\n" - # First chunk is bytes (fast path) wrapped via _format_streaming_sse_chunk. - first = out[0] - assert isinstance(first, bytes) - payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert normalize(payload) == { "id": "", "object": "chat.completion.chunk", @@ -718,6 +790,205 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): } +_AZURE_CONTENT_FILTER_RESULTS: Final = { + "hate": {"filtered": False, "severity": "safe"}, + "protected_material_text": {"detected": False, "filtered": False}, +} + + +def _azure_chunk(choice_fields: dict[str, object], chunk_fields: dict[str, object]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-azure", + choices=[StreamingChoices(index=0, delta=Delta(content="Red"), **choice_fields)], + created=0, + model="gpt-4", + object="chat.completion.chunk", + **chunk_fields, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("chunk", "read_field", "sent"), + [ + pytest.param( + _azure_chunk({"content_filter_results": _AZURE_CONTENT_FILTER_RESULTS}, {}), + lambda payload: payload["choices"][0]["content_filter_results"], + _AZURE_CONTENT_FILTER_RESULTS, + id="choice-content-filter-results", + ), + pytest.param( + _azure_chunk({}, {"obfuscation": "Ab3"}), + lambda payload: payload["obfuscation"], + "Ab3", + id="chunk-obfuscation", + ), + pytest.param( + ModelResponseStream( + id="chatcmpl-or", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Red", reasoning_details=[{"type": "reasoning.text", "text": "colors"}]), + ) + ], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ), + lambda payload: payload["choices"][0]["delta"]["reasoning_details"], + [{"type": "reasoning.text", "text": "colors"}], + id="delta-reasoning-details", + ), + pytest.param( + ModelResponseStream( + id="chatcmpl-resp", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Red", reasoning_items=[{"type": "reasoning", "id": "rs_1", "summary": []}]), + ) + ], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ), + lambda payload: payload["choices"][0]["delta"]["reasoning_items"], + [{"type": "reasoning", "id": "rs_1", "summary": []}], + id="delta-reasoning-items", + ), + ], +) +async def test_async_data_generator_sse_frame_keeps_provider_fields_on_a_text_chunk( + monkeypatch, chunk: ModelResponseStream, read_field: Callable[[dict[str, object]], object], sent: object +): + _patch_logging_flags(monkeypatch) + + out = [ + line + async for line in async_data_generator( + response=_async_iter([chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4"}, + ) + ] + + assert read_field(json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))) == sent + + +def _circular_dict() -> dict[str, object]: + looped: Final[dict[str, object]] = {"keep": 1} # mutable-ok: a circular dict has to be inserted into itself + looped["self"] = looped + return looped + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unencodable", + [ + pytest.param(object(), id="unknown-type"), + pytest.param(hashlib.sha256(b"trace").digest(), id="invalid-utf8-bytes"), + ], +) +async def test_async_data_generator_streams_a_text_chunk_a_hook_tagged_with_a_value_json_cannot_encode( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, unencodable: object +) -> None: + _patch_logging_flags(monkeypatch) + chunk: Final = _azure_chunk({"content_filter_results": _AZURE_CONTENT_FILTER_RESULTS}, {"obfuscation": "Ab3"}) + setattr(chunk, "trace_chunk", unencodable) + setattr(chunk.choices[0], "trace_choice", unencodable) + setattr(chunk.choices[0].delta, "trace_delta", unencodable) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out: Final = [ + line + async for line in async_data_generator( + response=_async_iter([chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4"}, + ) + ] + + payload: Final = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) + assert payload["obfuscation"] == "Ab3" + assert payload["choices"][0]["content_filter_results"] == _AZURE_CONTENT_FILTER_RESULTS + assert payload["choices"][0]["delta"] == {"content": "Red"} + assert "trace_chunk" not in payload + assert "trace_choice" not in payload["choices"][0] + assert all(name in caplog.text for name in ("trace_chunk", "trace_choice", "trace_delta")) + + +def test_serialize_streaming_chunk_keeps_lone_surrogates_a_provider_split_across_frames( + caplog: pytest.LogCaptureFixture, +) -> None: + tool_call: Final = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Par\udc00is"}') + ) + chunk: Final = ModelResponseStream( + id="chatcmpl-split", + choices=[StreamingChoices(index=0, delta=Delta(content="Hi \ud83d", tool_calls=[tool_call]))], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delta: Final = json.loads(_serialize_streaming_chunk(chunk))["choices"][0]["delta"] + + assert caplog.text == "" + assert delta["content"] == "Hi \ud83d" + assert delta["tool_calls"][0]["id"] == "call_1" + assert delta["tool_calls"][0]["function"] == {"name": "get_weather", "arguments": '{"city": "Par\udc00is"}'} + + +def test_serialize_streaming_chunk_drops_only_the_nested_value_json_cannot_encode() -> None: + digest: Final = hashlib.sha256(b"trace").digest() + chunk: Final = ModelResponseStream( + id="chatcmpl-nested", + choices=[ + StreamingChoices(index=0, delta=Delta(content="a")), + StreamingChoices( + index=1, delta=Delta(content="b", provider_specific_fields={"citations": ["ok"], "raw": digest}) + ), + ], + created=0, + model="gpt-4", + object="chat.completion.chunk", + usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + setattr(chunk.usage, "raw", digest) + setattr(chunk, "trace_chunk", _circular_dict()) + + payload: Final = json.loads(_serialize_streaming_chunk(chunk)) + + assert payload["choices"][0]["delta"] == {"content": "a"} + assert payload["choices"][1]["delta"] == {"content": "b", "provider_specific_fields": {"citations": ["ok"]}} + assert payload["usage"] == {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} + assert payload["trace_chunk"] == {"keep": 1} + + +def test_serialize_streaming_chunk_keeps_streaming_when_unencodable_values_sit_under_non_string_keys() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-keys", + choices=[StreamingChoices(index=0, delta=Delta(content="still streaming"))], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + setattr(chunk, "by_index", {0: "kept", 1: object()}) + setattr(chunk, "mixed_keys", {"name": "x", 2: object()}) + setattr(chunk, "bytes_keys", {b"raw": object()}) + setattr(chunk, "read_only_view", MappingProxyType({"name": object()})) + + payload: Final = json.loads(_serialize_streaming_chunk(chunk)) + + assert payload["choices"][0]["delta"] == {"content": "still streaming"} + assert payload["by_index"] == {"0": "kept"} + assert "mixed_keys" not in payload + assert "bytes_keys" not in payload + assert "read_only_view" not in payload + + @pytest.mark.asyncio async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch): _patch_logging_flags(monkeypatch) @@ -739,9 +1010,7 @@ async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch) ): out.append(line) - first = out[0] - assert isinstance(first, bytes) - payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert payload["model"] == "fallback-model" @@ -764,9 +1033,7 @@ async def test_async_data_generator_uses_chunk_fallback_metadata(monkeypatch): ): out.append(line) - first = out[0] - assert isinstance(first, bytes) - payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert payload["model"] == "fallback-model" @@ -793,8 +1060,8 @@ async def test_async_data_generator_switches_model_mid_stream_on_fallback(monkey ): out.append(line) - first_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) - second_payload = json.loads(out[1].removeprefix(b"data: ").removesuffix(b"\n\n")) + first_payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) + second_payload = json.loads(_sse_text(out[1]).removeprefix("data: ").removesuffix("\n\n")) assert first_payload["model"] == "primary-model" assert second_payload["model"] == "fallback-model" @@ -830,8 +1097,7 @@ async def test_async_data_generator_emits_fallback_error_metadata_event(monkeypa ): out.append(line) - assert isinstance(out[0], bytes) - chunk_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + chunk_payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert chunk_payload["model"] == "fallback-model" assert isinstance(out[1], str) assert out[1].startswith("data: ") @@ -879,8 +1145,7 @@ async def test_async_data_generator_skips_fallback_error_event_without_opt_in( ): out.append(line) - assert isinstance(out[0], bytes) - payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert payload["model"] == "fallback-model" @@ -2434,14 +2699,11 @@ async def test_queue_request_stream_is_untouched_while_keepalives_are_unconfigur assert chunks[-1] == b"data: [DONE]\n\n" -def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier(): +def test_serialize_streaming_chunk_keeps_served_service_tier(): chunk = _simple_chunk() chunk.service_tier = "priority" - result = _fast_serialize_simple_model_response_stream(chunk) - - assert result is not None - assert json.loads(result)["service_tier"] == "priority" + assert json.loads(_serialize_streaming_chunk(chunk))["service_tier"] == "priority" @pytest.mark.asyncio diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 72755633376..29413d2b189 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -147,29 +147,20 @@ def test_get_litellm_model_cost_map_catalog_only_excludes_runtime_registered_ent ) -def test_get_litellm_model_cost_map_returns_cost_map(): - app = FastAPI() +def test_get_litellm_model_cost_map_returns_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + model_cost: Final = { + "audio-model": {"input_cost_per_second": 0.1, "mode": "audio_transcription"}, + "token-model": {"input_cost_per_token": 0.01, "mode": "chat"}, + } + monkeypatch.setattr(litellm, "model_cost", model_cost) + app: Final = FastAPI() app.include_router(router) - client = TestClient(app) + client: Final = TestClient(app) - response = client.get("/public/litellm_model_cost_map") + response: Final = client.get("/public/litellm_model_cost_map") assert response.status_code == 200 - payload = response.json() - assert isinstance(payload, dict) - assert len(payload) > 0, "Expected model cost map to contain at least one model" - - # Verify the structure contains expected keys for at least one model - # Check for a common model like gpt-4 or gpt-3.5-turbo - model_keys = list(payload.keys()) - assert len(model_keys) > 0 - - # Verify at least one model has expected cost fields - sample_model = model_keys[0] - sample_model_data = payload[sample_model] - assert isinstance(sample_model_data, dict) - # Check for common cost fields that should be present - assert "input_cost_per_token" in sample_model_data or "output_cost_per_token" in sample_model_data + assert response.json() == model_cost def test_public_ai_hub_info_is_public_by_default(monkeypatch): @@ -536,6 +527,7 @@ def _decisions_provider_entry(provider: str) -> ProviderCreateInfo: ( ("TypeSafe", "TypeSafe", LlmProviders.TYPESAFE), ("Nace", "Nace", LlmProviders.NACE), + ("PARALLEL_AI", "Parallel AI", LlmProviders.PARALLEL_AI), ("EU_JEV", "eu/jev", LlmProviders.EU_JEV), ), ) diff --git a/tests/unit/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py b/tests/unit/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py index f5c97142dde..ab8dd8567b6 100644 --- a/tests/unit/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py +++ b/tests/unit/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py @@ -4,18 +4,23 @@ Tests for LiteLLM proxy realtime WebRTC HTTP endpoints: - POST /v1/realtime/calls """ +import functools import json import time -from collections.abc import Awaitable -from typing import Protocol +from collections.abc import Awaitable, Mapping +from typing import Final, Protocol from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from starlette.routing import Match - +import litellm +from litellm.litellm_core_utils.realtime_streaming import TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE +from litellm.proxy._lazy_features import LAZY_FEATURES, attach_lazy_features, reserve_lazy_slot from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -23,8 +28,18 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.realtime_endpoints.endpoints import ( + _ALLOWED_SESSION_TYPES, + _coerce_realtime_session_type, _decode_realtime_token_payload, _encode_realtime_token_payload, + _prepare_client_secret_session, +) +from litellm.types.realtime import ( + RealtimeAudioInputConfig, + RealtimeAudioTranscriptionConfig, + RealtimeClientSecretRequest, + RealtimeSessionAudioConfig, + RealtimeSessionConfig, ) # --- Unit tests: token encode/decode helpers --- @@ -126,9 +141,7 @@ def mock_route_request_client_secrets(): future_expires_at = int(time.time()) + 3600 mock_resp = MagicMock(spec=httpx.Response) mock_resp.status_code = 200 - mock_resp.text = ( - f'{{"value":"upstream_ephemeral_key","expires_at":{future_expires_at}}}' - ) + mock_resp.text = f'{{"value":"upstream_ephemeral_key","expires_at":{future_expires_at}}}' mock_resp.content = f'{{"value":"upstream_ephemeral_key","expires_at":{future_expires_at}}}'.encode() mock_resp.headers = {} mock_resp.json.return_value = { @@ -215,9 +228,7 @@ async def test_client_secrets_success_with_mock( mock_pre_call_hook, ): """POST /v1/realtime/client_secrets returns 200 with valid auth and mocked upstream.""" - proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( - user_id="test-user", team_id="test-team" - ) + proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="test-user", team_id="test-team") try: client = TestClient(proxy_app) with ( @@ -275,13 +286,7 @@ async def test_client_secrets_transcription_rejects_disallowed_nested_model( "session": { "type": "transcription", "model": "gpt-4o-realtime-preview", - "audio": { - "input": { - "transcription": { - "model": "gpt-realtime-whisper" - } - } - }, + "audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}, }, }, ) @@ -312,12 +317,8 @@ async def test_client_secrets_transcription_routes_on_nested_model( async def _inner(): resp = MagicMock(spec=httpx.Response) resp.status_code = 200 - resp.text = ( - f'{{"value":"upstream_ephemeral_key","expires_at":{future_expires_at}}}' - ) - resp.content = ( - f'{{"value":"upstream_ephemeral_key","expires_at":{future_expires_at}}}' - ).encode() + resp.text = f'{{"value":"upstream_ephemeral_key","expires_at":{future_expires_at}}}' + resp.content = (f'{{"value":"upstream_ephemeral_key","expires_at":{future_expires_at}}}').encode() resp.headers = {} resp.json.return_value = { "value": "upstream_ephemeral_key", @@ -351,13 +352,7 @@ async def test_client_secrets_transcription_routes_on_nested_model( "session": { "type": "transcription", "model": "gpt-4o-realtime-preview", - "audio": { - "input": { - "transcription": { - "model": "gpt-realtime-whisper" - } - } - }, + "audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}, }, }, ) @@ -367,10 +362,7 @@ async def test_client_secrets_transcription_routes_on_nested_model( session = captured["data"]["session"] assert session["type"] == "transcription" assert "model" not in session - assert ( - session["audio"]["input"]["transcription"]["model"] - == "gpt-realtime-whisper" - ) + assert session["audio"]["input"]["transcription"]["model"] == "gpt-realtime-whisper" encrypted_value = response.json()["value"] decoded = _decode_realtime_token_payload( decrypt_value_helper( @@ -530,10 +522,7 @@ async def test_realtime_calls_replays_transcription_session_type( ) assert captured["session"]["type"] == "transcription" - assert ( - captured["session"]["audio"]["input"]["transcription"]["model"] - == "gpt-realtime-whisper" - ) + assert captured["session"]["audio"]["input"]["transcription"]["model"] == "gpt-realtime-whisper" # --- transcription_sessions endpoint --- @@ -587,8 +576,12 @@ def test_transcription_sessions_requires_auth(proxy_app): @pytest.mark.asyncio +@pytest.mark.parametrize("body", [ + {"input_audio_transcription": {"model": "gpt-realtime-whisper"}}, + {"audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}}, +]) async def test_transcription_sessions_rejects_disallowed_resolved_model( - proxy_app, + proxy_app, body, ): proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( user_id="test-user", @@ -605,9 +598,7 @@ async def test_transcription_sessions_rejects_disallowed_resolved_model( response = client.post( "/v1/realtime/transcription_sessions", headers={"Authorization": "Bearer sk-test-master-key"}, - json={ - "input_audio_transcription": {"model": "gpt-realtime-whisper"} - }, + json=body, ) assert response.status_code == 403 @@ -651,9 +642,7 @@ async def test_transcription_sessions_rejects_disallowed_team_model_scope( response = client.post( "/v1/realtime/transcription_sessions", headers={"Authorization": "Bearer sk-test-master-key"}, - json={ - "input_audio_transcription": {"model": "gpt-realtime-whisper"} - }, + json={"input_audio_transcription": {"model": "gpt-realtime-whisper"}}, ) assert response.status_code == 403 @@ -696,9 +685,7 @@ async def test_transcription_sessions_rejects_disallowed_project_model_scope( response = client.post( "/v1/realtime/transcription_sessions", headers={"Authorization": "Bearer sk-test-master-key"}, - json={ - "input_audio_transcription": {"model": "gpt-realtime-whisper"} - }, + json={"input_audio_transcription": {"model": "gpt-realtime-whisper"}}, ) assert response.status_code == 403 @@ -751,9 +738,7 @@ async def test_transcription_sessions_rejects_disallowed_team_member_model_scope response = client.post( "/v1/realtime/transcription_sessions", headers={"Authorization": "Bearer sk-test-master-key"}, - json={ - "input_audio_transcription": {"model": "gpt-realtime-whisper"} - }, + json={"input_audio_transcription": {"model": "gpt-realtime-whisper"}}, ) assert response.status_code == 403 @@ -947,13 +932,12 @@ async def test_transcription_sessions_encrypts_client_secret( POST /v1/realtime/transcription_sessions returns 200 and the ephemeral key under client_secret.value must be encrypted (never the raw upstream key). """ - proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( - user_id="test-user", team_id="test-team" - ) + proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="test-user", team_id="test-team") captured_route_type = {} async def _capturing_route(*args, **kwargs): captured_route_type["route_type"] = kwargs.get("route_type") + captured_route_type["session"] = kwargs["data"]["transcription_session"] return await mock_route_request_transcription_sessions(*args, **kwargs) try: @@ -976,6 +960,8 @@ async def test_transcription_sessions_encrypts_client_secret( "/v1/realtime/transcription_sessions", headers={"Authorization": "Bearer sk-test-master-key"}, json={ + "model": "gpt-realtime-whisper", + "audio": {"input": {"transcription": {"model": "gpt-live-transcribe", "language": "en"}}}, "input_audio_format": "pcm16", "input_audio_transcription": {"model": "gpt-realtime-whisper"}, }, @@ -993,16 +979,61 @@ async def test_transcription_sessions_encrypts_client_secret( assert decrypted is not None assert "upstream_ephemeral_key" in decrypted # Routed through the dedicated transcription_sessions route type. - assert ( - captured_route_type["route_type"] - == "acreate_realtime_transcription_session" - ) + assert captured_route_type["route_type"] == "acreate_realtime_transcription_session" + assert captured_route_type["session"]["input_audio_transcription"]["model"] == "gpt-realtime-whisper" + assert captured_route_type["session"]["audio"]["input"]["transcription"] == { + "model": "gpt-realtime-whisper", "language": "en", + } + finally: + proxy_app.dependency_overrides.pop(user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [ + {"model": "gpt-4o-realtime-preview", "audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}}, + {"input_audio_transcription": {"model": "gpt-4o-realtime-preview"}, "audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}}, + {"model": "gpt-4o-realtime-preview", "input_audio_transcription": {"model": "gpt-realtime-whisper"}}, +]) +async def test_transcription_sessions_route_the_allowed_model_over_a_nested_model_the_key_cannot_call( + proxy_app, + mock_route_request_transcription_sessions, + mock_add_litellm_data, + mock_pre_call_hook, + body, +): + proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user", models=["gpt-4o-realtime-preview"] + ) + route: Final = AsyncMock(side_effect=mock_route_request_transcription_sessions) + + try: + client = TestClient(proxy_app) + with ( + patch("litellm.proxy.proxy_server.route_request", route), + patch("litellm.proxy.proxy_server.add_litellm_data_to_request", side_effect=mock_add_litellm_data), + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, + ): + mock_logging.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + mock_logging.post_call_failure_hook = AsyncMock() + + response = client.post( + "/v1/realtime/transcription_sessions", + headers={"Authorization": "Bearer sk-test-master-key"}, + json=body, + ) + + assert response.status_code == 200, response.text + route.assert_awaited_once() + routed_session: Final = route.await_args.kwargs["data"]["transcription_session"] + assert { + routed_session.get("input_audio_transcription", {}).get("model"), + routed_session.get("audio", {}).get("input", {}).get("transcription", {}).get("model"), + } - {None} == {"gpt-4o-realtime-preview"} finally: proxy_app.dependency_overrides.pop(user_api_key_auth, None) def test_session_type_coerced_for_unknown_value(): - """An unrecognized session_type in the token falls back to 'realtime'.""" payload = _encode_realtime_token_payload( ephemeral_key="epk", model_id="gpt-4o", @@ -1011,12 +1042,327 @@ def test_session_type_coerced_for_unknown_value(): expires_at=None, session_type="INJECTED_TYPE", ) - # Force-deserialize and check the coercion that happens in proxy_realtime_calls. decoded = json.loads(payload) - session_type = decoded.get("session_type") or "realtime" - if session_type not in ("realtime", "transcription"): - session_type = "realtime" - assert session_type == "realtime" + assert decoded["session_type"] == "INJECTED_TYPE" + assert _coerce_realtime_session_type("INJECTED_TYPE") == "realtime" + assert _coerce_realtime_session_type(None) == "realtime" + for allowed_session_type in _ALLOWED_SESSION_TYPES: + assert _coerce_realtime_session_type(allowed_session_type) == allowed_session_type + + +@pytest.mark.asyncio +async def test_translation_client_secret_rejects_disallowed_nested_transcription_model() -> None: + req = RealtimeClientSecretRequest( + model="gpt-realtime-translate", + session=RealtimeSessionConfig( + type="translation", + model="gpt-realtime-translate", + audio=RealtimeSessionAudioConfig( + input=RealtimeAudioInputConfig( + transcription=RealtimeAudioTranscriptionConfig(model="gpt-live-transcribe"), + ) + ), + ), + ) + + with pytest.raises(Exception, match=r"gpt-live-transcribe.*not available for this API key"): + await _prepare_client_secret_session( + req=req, + user_api_key_dict=UserAPIKeyAuth(models=["gpt-realtime-translate"]), + llm_model_list=None, + llm_router=None, + forced_session_type="translation", + ) + + +@pytest.mark.asyncio +async def test_translation_client_secret_binds_authorized_nested_transcription_model() -> None: + req = RealtimeClientSecretRequest( + model="gpt-realtime-translate", + session=RealtimeSessionConfig( + type="translation", + model="gpt-realtime-translate", + audio=RealtimeSessionAudioConfig( + input=RealtimeAudioInputConfig( + transcription=RealtimeAudioTranscriptionConfig(model="gpt-realtime-whisper"), + ) + ), + ), + ) + + model, session_data, session_type = await _prepare_client_secret_session( + req=req, + user_api_key_dict=UserAPIKeyAuth(models=["gpt-realtime-translate", "gpt-realtime-whisper"]), + llm_model_list=None, + llm_router=None, + forced_session_type="translation", + ) + + assert model == "gpt-realtime-translate" + assert session_type == "translation" + assert session_data is not None + assert session_data["model"] == "gpt-realtime-translate" + assert session_data["audio"]["input"]["transcription"]["model"] == "gpt-realtime-whisper" + + +@pytest.mark.parametrize( + "path", + ( + "/openai/v1/realtime/client_secrets", + "/openai/v1/realtime/calls", + "/openai/v1/realtime/translations/client_secrets", + "/openai/v1/realtime/translations/calls", + ), +) +def test_openai_prefixed_realtime_http_paths_stay_with_the_openai_pass_through( + monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import openai_proxy_route + + monkeypatch.setenv("LITELLM_DISABLE_LAZY_ROUTES", "true") + features = tuple(feature for feature in LAZY_FEATURES if feature.name in ("llm_passthrough", "realtime")) + app = FastAPI() + reserve_lazy_slot(app, "llm_passthrough", features=features) + attach_lazy_features(app, features) + scope = {"type": "http", "method": "POST", "path": path, "root_path": "", "headers": [], "query_string": b""} + + with TestClient(app): + resolved = next(route for route in app.router.routes if route.matches(scope)[0] is Match.FULL) + + assert getattr(resolved, "endpoint", None) is openai_proxy_route + + +@pytest.mark.parametrize( + "path", + [ + "/v1/realtime/translations/client_secrets", + "/realtime/translations/client_secrets", + ], +) +def test_translation_client_secret_aliases_bind_token_family( + proxy_app, + mock_route_request_client_secrets, + mock_add_litellm_data, + mock_pre_call_hook, + path, +): + proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user", + models=["gpt-realtime-translate"], + ) + captured = {} + + async def capture_route(*args, **kwargs): + captured["route_type"] = kwargs["route_type"] + captured["data"] = kwargs["data"] + return await mock_route_request_client_secrets(*args, **kwargs) + + try: + client = TestClient(proxy_app) + with ( + patch( # test-quality-ok: endpoint test captures the proxy routing boundary + "litellm.proxy.proxy_server.route_request", side_effect=capture_route + ), + patch( # test-quality-ok: endpoint test isolates request metadata enrichment + "litellm.proxy.proxy_server.add_litellm_data_to_request", + side_effect=mock_add_litellm_data, + ), + patch( # test-quality-ok: endpoint test isolates the process-wide proxy logger + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as logging, + ): + logging.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + logging.post_call_failure_hook = AsyncMock() + response = client.post( + path, + headers={"Authorization": "Bearer sk-test-master-key"}, + json={"model": "gpt-realtime-translate"}, + ) + + assert response.status_code == 200 + assert captured["route_type"] == "acreate_realtime_translation_client_secret" + assert captured["data"]["session"] == { + "type": "translation", + "model": "gpt-realtime-translate", + } + decrypted = decrypt_value_helper( + response.json()["value"], + key="client_secret.value", + exception_type="debug", + ) + decoded = _decode_realtime_token_payload(decrypted or "") + assert decoded is not None + assert decoded["session_type"] == "translation" + finally: + proxy_app.dependency_overrides.pop(user_api_key_auth, None) + + +@pytest.mark.parametrize( + "path", + [ + "/v1/realtime/translations/calls", + "/realtime/translations/calls", + ], +) +def test_translation_calls_aliases_route_translation_session( + proxy_app, + mock_route_request_realtime_calls, + mock_add_litellm_data, + mock_pre_call_hook, + path, +): + token = encrypt_value_helper( + _encode_realtime_token_payload( + ephemeral_key="ek_test", + model_id="gpt-realtime-translate", + user_id="test-user", + team_id=None, + expires_at=int(time.time()) + 3600, + session_type="translation", + ) + ) + captured = {} + + async def capture_route(*args, **kwargs): + captured["route_type"] = kwargs["route_type"] + captured["data"] = kwargs["data"] + return await mock_route_request_realtime_calls(*args, **kwargs) + + client = TestClient(proxy_app) + with ( + patch( # test-quality-ok: endpoint test captures the proxy routing boundary + "litellm.proxy.proxy_server.route_request", side_effect=capture_route + ), + patch( # test-quality-ok: endpoint test isolates request metadata enrichment + "litellm.proxy.proxy_server.add_litellm_data_to_request", + side_effect=mock_add_litellm_data, + ), + patch( # test-quality-ok: endpoint test isolates the process-wide proxy logger + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as logging, + ): + logging.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + logging.post_call_failure_hook = AsyncMock() + response = client.post( + path, + headers={"Authorization": f"Bearer {token}"}, + content=b"v=0\r\n", + ) + + assert response.status_code == 201 + assert captured["route_type"] == "arealtime_translation_calls" + assert captured["data"]["session"] == { + "type": "translation", + "model": "gpt-realtime-translate", + } + + +@pytest.mark.parametrize( + ("disable_global_guardrails", "expected_status", "expected_body"), + [(True, 201, "v=0"), (False, 400, TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE)], +) +def test_translation_calls_honor_the_minting_keys_guardrail_opt_out( + proxy_app: FastAPI, + mock_route_request_client_secrets: object, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + monkeypatch: pytest.MonkeyPatch, + disable_global_guardrails: bool, + expected_status: int, + expected_body: str, +) -> None: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy import proxy_server + from litellm.types.guardrails import GuardrailEventHooks + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-realtime-translate", + "litellm_params": {"model": "openai/gpt-realtime-translate", "api_key": "sk-upstream"}, + } + ] + ), + ) + monkeypatch.setattr( + litellm, + "callbacks", + [CustomGuardrail(guardrail_name="global-pre-call", event_hook=GuardrailEventHooks.pre_call, default_on=True)], + ) + proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user", + models=["gpt-realtime-translate"], + metadata={"disable_global_guardrails": disable_global_guardrails}, + ) + client: Final = TestClient(proxy_app, raise_server_exceptions=False) + try: + with ( + patch( # test-quality-ok: endpoint test stands in for the upstream client_secrets call + "litellm.proxy.proxy_server.route_request", side_effect=mock_route_request_client_secrets + ), + patch( # test-quality-ok: endpoint test isolates request metadata enrichment while minting + "litellm.proxy.proxy_server.add_litellm_data_to_request", + side_effect=mock_add_litellm_data, + ), + patch( # test-quality-ok: endpoint test isolates the process-wide proxy logger while minting + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as logging, + ): + logging.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + logging.post_call_failure_hook = AsyncMock() + minted: Final = client.post( + "/v1/realtime/translations/client_secrets", + headers={"Authorization": "Bearer sk-test-master-key"}, + json={"model": "gpt-realtime-translate"}, + ) + finally: + proxy_app.dependency_overrides.pop(user_api_key_auth, None) + + with respx.mock(assert_all_called=False) as upstream: + upstream.post("https://api.openai.com/v1/realtime/translations/calls").mock( + return_value=httpx.Response(201, content=b"v=0\r\n", headers={"content-type": "application/sdp"}) + ) + answered: Final = client.post( + "/v1/realtime/translations/calls", + headers={"Authorization": f"Bearer {minted.json()['value']}"}, + content=b"v=0\r\n", + ) + + assert minted.status_code == 200 + assert answered.status_code == expected_status + assert expected_body in answered.text + + +@pytest.mark.parametrize( + "session_type,path", + [ + ("realtime", "/v1/realtime/translations/calls"), + ("translation", "/v1/realtime/calls"), + ], +) +def test_realtime_calls_reject_cross_family_token(proxy_app, session_type, path): + model = "gpt-realtime-translate" if session_type == "translation" else "gpt-realtime-2" + token = encrypt_value_helper( + _encode_realtime_token_payload( + ephemeral_key="ek_test", + model_id=model, + user_id=None, + team_id=None, + expires_at=int(time.time()) + 3600, + session_type=session_type, + ) + ) + response = TestClient(proxy_app).post( + path, + headers={"Authorization": f"Bearer {token}"}, + content=b"v=0\r\n", + ) + assert response.status_code == 401 + assert response.json()["error"] == "Token is not valid for this Realtime endpoint" @pytest.mark.asyncio @@ -1142,9 +1488,7 @@ async def test_transcription_sessions_returns_upstream_error_verbatim( return _inner() - proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( - user_id="test-user", team_id="test-team" - ) + proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="test-user", team_id="test-team") try: client = TestClient(proxy_app) with ( @@ -1184,9 +1528,7 @@ async def test_transcription_sessions_wraps_route_exception( async def _raise_http(*args, **kwargs): raise HTTPException(status_code=403, detail="Model not allowed") - proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( - user_id="test-user" - ) + proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="test-user") try: client = TestClient(proxy_app, raise_server_exceptions=False) with ( @@ -1287,3 +1629,260 @@ def test_transcription_sessions_rejection_answers_an_openai_typed_error( assert response.status_code == 400 assert (response.json()["error"]["type"], response.json()["error"]["param"]) == ("invalid_request_error", None) + + +_AZURE_REALTIME_DEPLOYMENTS: Final = ( + ("azure-realtime-alias", "azure/qa-realtime-deployment"), + ("azure-transcribe-alias", "azure/qa-whisper-deployment"), + ("azure-translate-alias", "azure/qa-translate-deployment"), +) + + +def _post_through_the_azure_deployments( + proxy_app: FastAPI, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + path: str, + body: object, + deployments: tuple[tuple[str, str], ...] = _AZURE_REALTIME_DEPLOYMENTS, +) -> tuple[httpx.Response, respx.Route]: + from litellm.proxy import proxy_server + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + model_list: Final = [ # mutable-ok: Router takes a list of deployments + { + "model_name": alias, + "litellm_params": { + "model": deployment, + "api_base": "https://realtime-qa.openai.azure.com", + "api_key": "azure-test-key", + }, + } + for alias, deployment in deployments + ] + monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=model_list)) + monkeypatch.setattr(proxy_server, "llm_model_list", model_list) + upstream: Final = respx_mock.route(host="realtime-qa.openai.azure.com").mock( + return_value=httpx.Response(200, json={"value": "ek_azure", "expires_at": int(time.time()) + 3600}) + ) + monkeypatch.setitem( + proxy_app.dependency_overrides, + user_api_key_auth, + lambda: UserAPIKeyAuth(user_id="test-user", models=[alias for alias, _ in deployments]), + ) + with ( + patch( # test-quality-ok: endpoint test isolates request metadata enrichment + "litellm.proxy.proxy_server.add_litellm_data_to_request", side_effect=mock_add_litellm_data + ), + patch( # test-quality-ok: endpoint test isolates the process-wide proxy logger + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as logging, + ): + logging.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + logging.post_call_failure_hook = AsyncMock() + response: Final = TestClient(proxy_app).post( + path, headers={"Authorization": "Bearer sk-test-master-key"}, json=body + ) + return response, upstream + + +def _upstream_session_value(upstream: respx.Route, session_path: tuple[str, ...]) -> object: + return functools.reduce( + lambda node, key: node[key], session_path, json.loads(upstream.calls.last.request.content)["session"] + ) + + +@pytest.mark.parametrize( + ("path", "body", "upstream_model_path", "expected_model"), + ( + ( + "/v1/realtime/client_secrets", + {"session": {"type": "realtime", "model": "azure-realtime-alias"}}, + ("model",), + "qa-realtime-deployment", + ), + ( + "/v1/realtime/client_secrets", + {"model": "azure-realtime-alias"}, + ("model",), + "qa-realtime-deployment", + ), + ( + "/v1/realtime/client_secrets", + { + "session": { + "type": "transcription", + "audio": {"input": {"transcription": {"model": "azure-transcribe-alias"}}}, + } + }, + ("audio", "input", "transcription", "model"), + "qa-whisper-deployment", + ), + ( + "/v1/realtime/translations/client_secrets", + {"model": "azure-translate-alias"}, + ("model",), + "qa-translate-deployment", + ), + ( + "/v1/realtime/translations/client_secrets", + { + "model": "azure-translate-alias", + "session": {"audio": {"input": {"transcription": {"model": "azure-transcribe-alias"}}}}, + }, + ("audio", "input", "transcription", "model"), + "qa-whisper-deployment", + ), + ), +) +def test_client_secrets_mint_a_secret_for_the_azure_deployment_model( + proxy_app: FastAPI, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + path: str, + body: dict[str, object], + upstream_model_path: tuple[str, ...], + expected_model: str, +) -> None: + response, upstream = _post_through_the_azure_deployments( + proxy_app, mock_add_litellm_data, mock_pre_call_hook, respx_mock, monkeypatch, path, body + ) + + assert response.status_code == 200, response.text + assert _upstream_session_value(upstream, upstream_model_path) == expected_model + + +@pytest.mark.parametrize( + ("body", "session_path", "expected"), + ( + ( + { + "session": { + "type": "realtime", + "model": "azure-realtime-alias", + "audio": {"input": {"turn_detection": None}}, + } + }, + ("audio", "input", "turn_detection"), + None, + ), + ( + { + "session": { + "type": "transcription", + "audio": {"input": {"transcription": {"model": "azure-transcribe-alias"}, "turn_detection": None}}, + } + }, + ("audio", "input", "turn_detection"), + None, + ), + ({"session": {"type": "conversation", "model": "azure-realtime-alias"}}, ("type",), "realtime"), + ( + { + "session": { + "type": "transcription", + "audio": {"input": {"transcription": {"model": "azure-transcribe-alias", "delay": "auto"}}}, + } + }, + ("audio", "input", "transcription", "delay"), + "auto", + ), + ), +) +def test_client_secrets_forward_every_session_value_the_caller_set( + proxy_app: FastAPI, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + body: Mapping[str, object], + session_path: tuple[str, ...], + expected: object, +) -> None: + response, upstream = _post_through_the_azure_deployments( + proxy_app, + mock_add_litellm_data, + mock_pre_call_hook, + respx_mock, + monkeypatch, + "/v1/realtime/client_secrets", + body, + ) + + assert response.status_code == 200, response.text + assert _upstream_session_value(upstream, session_path) == expected + + +def test_client_secrets_without_a_model_let_the_upstream_pick_its_default_through_a_wildcard_deployment( + proxy_app: FastAPI, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, upstream = _post_through_the_azure_deployments( + proxy_app, + mock_add_litellm_data, + mock_pre_call_hook, + respx_mock, + monkeypatch, + "/v1/realtime/client_secrets", + {"session": {"type": "realtime"}}, + deployments=(("*", "openai/*"),), + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content)["session"] == {"type": "realtime"} + + +def test_client_secrets_send_the_deployment_model_in_the_flat_transcription_field( + proxy_app: FastAPI, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, upstream = _post_through_the_azure_deployments( + proxy_app, + mock_add_litellm_data, + mock_pre_call_hook, + respx_mock, + monkeypatch, + "/v1/realtime/client_secrets", + {"session": {"type": "transcription", "input_audio_transcription": {"model": "azure-transcribe-alias"}}}, + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content)["session"] == { + "type": "transcription", + "input_audio_transcription": {"model": "qa-whisper-deployment"}, + } + + +@pytest.mark.parametrize( + ("path", "body"), + ( + ("/v1/realtime/client_secrets", {"session": "not-a-session"}), + ("/v1/realtime/translations/client_secrets", {"session": {"audio": {"input": "not-an-input"}}}), + ("/v1/realtime/transcription_sessions", {"audio": {"input": {"transcription": "not-a-transcription"}}}), + ), +) +def test_realtime_session_routes_answer_400_for_a_malformed_body( + proxy_app: FastAPI, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + path: str, + body: Mapping[str, object], +) -> None: + response, upstream = _post_through_the_azure_deployments( + proxy_app, mock_add_litellm_data, mock_pre_call_hook, respx_mock, monkeypatch, path, body + ) + + assert (response.status_code, response.json()["error"]["type"]) == (400, "invalid_request_error") + assert upstream.called is False diff --git a/tests/unit/proxy/rerank_endpoints/test_endpoints.py b/tests/unit/proxy/rerank_endpoints/test_endpoints.py index 52d12dd1813..c1908a97d7e 100644 --- a/tests/unit/proxy/rerank_endpoints/test_endpoints.py +++ b/tests/unit/proxy/rerank_endpoints/test_endpoints.py @@ -5,10 +5,13 @@ Tests for rerank_endpoints/endpoints.py response headers. import json import logging from collections.abc import Iterator +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import HTTPException, Request, Response +import respx +from fastapi import FastAPI, HTTPException, Request, Response +from fastapi.testclient import TestClient import litellm.proxy.common_request_processing as common_request_processing_mod import litellm.proxy.proxy_server as proxy_server_mod @@ -209,3 +212,68 @@ async def test_a_rejection_raised_before_routing_keeps_its_own_status(monkeypatc error = await _rerank_failure(rejection, raised_before_routing=True, monkeypatch=monkeypatch) assert (error.type, error.param, error.code) == ("bad_request_error", "session_id", "400") + + +def test_failed_rerank_spend_row_records_arerank( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + import litellm + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.proxy.utils import ProxyLogging + + provider: Final = respx_mock.post(url__startswith="https://rerank-provider.example").respond( + 400, json={"message": "invalid documents"} + ) + cache: Final = UserApiKeyCache() + spend_writer: Final = MagicMock(update_database=AsyncMock()) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger(spend_writer=lambda: spend_writer)]) + monkeypatch.setattr(proxy_server_mod, "general_settings", {}) + monkeypatch.setattr( + proxy_server_mod, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "rerank-model", + "litellm_params": { + "model": "cohere/rerank-v3.5", + "api_base": "https://rerank-provider.example", + "api_key": "provider-key", + }, + } + ], + num_retries=0, + ), + ) + monkeypatch.setattr(proxy_server_mod, "proxy_config", ProxyConfig()) + monkeypatch.setattr(proxy_server_mod, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server_mod, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server_mod, "user_model", None) + monkeypatch.setattr(proxy_server_mod, "version", "1.2.3") + app: Final = FastAPI() + app.include_router(rerank_router) + app.add_exception_handler(ProxyException, proxy_server_mod.openai_exception_handler) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(api_key="sk-test") + try: + with TestClient(app, raise_server_exceptions=False) as client: + response: Final = client.post( + "/v1/rerank", json={"model": "rerank-model", "query": "q", "documents": ["a", "b"]} + ) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + assert (response.status_code, provider.call_count) == (400, 1), response.text + db_call: Final = spend_writer.update_database.call_args.kwargs + payload: Final = get_logging_payload( + kwargs=db_call["kwargs"], + response_obj=db_call["completion_response"], + start_time=db_call["start_time"], + end_time=db_call["end_time"], + ) + assert (payload["status"], payload["call_type"]) == ("failure", "arerank") diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 6912049d225..51b5bf13f16 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -123,6 +123,8 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: return router.get_configured_display_name(name) if field == "service_tiers": return router.get_configured_service_tiers(name, **_given(team_id=team_id)) + if field == "max_input_tokens": + return router.get_configured_max_input_tokens(name, **_given(team_id=team_id)) if field == "listing_info": return router.get_model_listing_info(name) if field == "mode": diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index f52f40ad55d..d5771dff626 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5150,6 +5150,23 @@ def test_a_background_response_billing_event_writes_a_row_apart_from_the_read_th assert request_ids == ("resp_shared", "resp_shared_background_cost") +@pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) +def test_a_read_or_cancel_row_carries_the_id_its_callbacks_logged(call_type: str): + row: Final = get_logging_payload( + kwargs={ + "call_type": call_type, + "model": "gpt-5.6", + "litellm_call_id": "call-read", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj={"id": "resp_logged_by_callbacks", "object": "response", "status": "completed"}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert row["request_id"] == "resp_logged_by_callbacks" + + @pytest.mark.parametrize("forged_marker", [True, BackgroundResponseSettlement.SETTLES.value]) def test_a_request_body_cannot_move_its_spend_row_onto_a_background_settlement_id(forged_marker): """A failed call logs its raw request body as kwargs, so a body field named like the settlement @@ -7469,6 +7486,46 @@ def test_failure_error_information_masks_ip_addresses_in_long_text_and_at_its_tr assert sanitized["error_code"] == "500" +def test_spend_row_sanitizes_the_session_failure_a_realtime_success_row_carries( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.proxy_server import general_settings + + monkeypatch.setitem(general_settings, "store_prompts_in_spend_logs", False) + monkeypatch.setenv("MAX_STRING_LENGTH_PROMPT_IN_DB", "200") + provider_error: Final = 'stream ended {"messages": [{"role": "user", "content": "private-prompt"}]} ' + "x" * 2000 + now: Final = datetime.datetime.now(timezone.utc) + + row: Final = get_logging_payload( + kwargs={ + "model": "bedrock/amazon.nova-sonic-v1:0", + "call_type": "_arealtime", + "litellm_params": { + "metadata": { + "user_api_key": "hashed-key", + "status": "failure", + "error_information": { + "error_code": "503", + "error_class": "BedrockError", + "error_message": provider_error, + "llm_provider": "bedrock", + }, + } + }, + }, + response_obj=None, + start_time=now, + end_time=now, + ) + + stored: Final = json.loads(row["metadata"])["error_information"] + assert row["status"] == "failure" + assert "private-prompt" not in stored["error_message"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored["error_message"] + assert len(stored["error_message"]) < len(provider_error) + assert stored["error_class"] == "BedrockError" + + @pytest.mark.parametrize( "address", ["203.0.113.9", "2001:db8::9", "1:2:3:4:5:6:7:8", "1:2:3:4:5:6:7::", "::1:2:3:4:5:6:7", "fe80::1"] ) diff --git a/tests/unit/proxy/test__lazy_features.py b/tests/unit/proxy/test__lazy_features.py index a4bd47ccb74..c140641a71b 100644 --- a/tests/unit/proxy/test__lazy_features.py +++ b/tests/unit/proxy/test__lazy_features.py @@ -1,6 +1,7 @@ import sys from collections.abc import AsyncGenerator, Mapping from contextlib import asynccontextmanager +from dataclasses import replace from types import ModuleType from typing import Final @@ -443,3 +444,45 @@ def test_root_authorization_server_metadata_is_the_gateways_not_the_byok_connect assert root.json()["token_endpoint"] == "https://gateway.example.com/token" assert connector.status_code == 200 assert connector.json()["token_endpoint"] == "https://gateway.example.com/v1/mcp/oauth/token" + + +@pytest.mark.parametrize( + "load_order", + (("/realtime/unrouted", "/mcp/unrouted"), ("/mcp/unrouted", "/realtime/unrouted")), + ids=("realtime_first", "realtime_last"), +) +def test_realtime_and_mcp_features_keep_every_route_once_in_either_load_order( + monkeypatch: pytest.MonkeyPatch, load_order: tuple[str, str] +) -> None: + monkeypatch.delenv(FLAG, raising=False) + + async def mcp_mount() -> Mapping[str, str]: + return {"route": "mcp_mount"} + + mounted_app: Final = FastAPI() + mounted_app.add_api_route("/{path:path}", mcp_mount, methods=["GET"]) + mcp_module: Final = ModuleType("tests.unit.proxy.lazy_fixture_mcp_app") + mcp_module.app = mounted_app # pyright: ignore[reportAttributeAccessIssue] # fixture module built at test time + monkeypatch.setitem(sys.modules, mcp_module.__name__, mcp_module) + features: Final = tuple( + replace(feature, module_path=mcp_module.__name__) if feature.name == "mcp_app" else feature + for feature in LAZY_FEATURES + if feature.name in ("mcp_app", "realtime") + ) + + async def mcp_proxy() -> Mapping[str, str]: + return {"route": "mcp_proxy"} + + app: Final = FastAPI() + attach_lazy_features(app, features) + app.add_api_route("/mcp/proxy", mcp_proxy, methods=["GET"]) + + with TestClient(app) as client: + for path in load_order: + client.get(path) + response: Final = client.get("/mcp/proxy") + + assert loaded_lazy_modules(app) == {feature.module_path for feature in features} + assert response.json() == {"route": "mcp_proxy"} + route_ids: Final = tuple(id(route) for route in app.routes) + assert len(route_ids) == len(frozenset(route_ids)) diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 535c3da8c3a..6ce8d400daf 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -66,7 +66,7 @@ def spend_counter_state(): original_prisma_client = ps.prisma_client counter_cache = DualCache() - key_cache = DualCache() + key_cache = UserApiKeyCache() ps.spend_counter_cache = counter_cache ps.user_api_key_cache = key_cache ps.prisma_client = None diff --git a/tests/unit/proxy/test_model_dump_with_preserved_fields.py b/tests/unit/proxy/test_model_dump_with_preserved_fields.py index 31531dccd16..92c7e58126b 100644 --- a/tests/unit/proxy/test_model_dump_with_preserved_fields.py +++ b/tests/unit/proxy/test_model_dump_with_preserved_fields.py @@ -545,3 +545,31 @@ def test_provider_refusal_string_is_preserved_on_message() -> None: "object": "chat.completion", } assert model_response.choices[0].message.refusal == "I can't help with that." + + +def test_vllm_null_stop_reason_is_kept_only_on_the_choice_that_carries_it() -> None: + model_response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-vllm-natural-stop", + "created": 4, + "model": "Qwen/Qwen3-0.6B", + "object": "chat.completion", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"role": "assistant", "content": "1, 2"}, + "stop_reason": None, + }, + {"finish_reason": "stop", "index": 1, "message": {"role": "assistant", "content": "1, 2"}}, + ], + }, + model_response_object=ModelResponse(), + ) + + dumped_choices: Final = model_dump_with_preserved_fields(model_response, exclude_unset=True)["choices"] + + assert [(choice["index"], "stop_reason" in choice, choice.get("stop_reason")) for choice in dumped_choices] == [ + (0, True, None), + (1, False, None), + ] diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index 6a9cd972dc2..839ab49408f 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -1,31 +1,41 @@ -"""OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals.""" +"""OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks.""" import json -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace from typing import Final -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import httpx import pytest import respx +from fastapi import HTTPException from starlette.routing import WebSocketRoute +from starlette.websockets import WebSocket import litellm from litellm.llms.openai.workload_identity import _workload_identity_auth -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _OPENAI_WS_ALIASED_MODEL_REFUSAL, _OPENAI_WS_DISABLED_REFUSAL, - _OPENAI_WS_MODEL_RESTRICTED_REFUSAL, - _has_model_restrictions, + _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, _openai_websocket_refusal, - _proxy_model_allowlists, + _proxy_frame_model_gate, openai_websocket_proxy_route, router, ) - -Scopes = tuple[Sequence[str], ...] +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + WebsocketFrameModelAllowed, + WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, + WebsocketUpstreamFrameClassifier, +) ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True}) DISABLED_SETTINGS: Final = ( @@ -76,6 +86,8 @@ class _RelayCall: forward_headers: bool endpoint: str accept_websocket: bool + client_frame_model_gate: WebsocketFrameModelGate | None + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None class _FakeRelay: @@ -92,6 +104,8 @@ class _FakeRelay: forward_headers: bool, endpoint: str, accept_websocket: bool, + client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ) -> None: self.calls.append( _RelayCall( @@ -100,24 +114,26 @@ class _FakeRelay: forward_headers=forward_headers, endpoint=endpoint, accept_websocket=accept_websocket, + client_frame_model_gate=client_frame_model_gate, + upstream_frame_kind=upstream_frame_kind, ) ) -class _FakeModelAllowlists: - def __init__(self, scopes: Scopes) -> None: - self.scopes = scopes - self.calls: list[UserAPIKeyAuth] = [] +class _FakeFrameModelGate: + def __init__(self, refusals: Mapping[str, WebsocketFrameModelRefused] = MappingProxyType({})) -> None: + self.refusals = refusals + self.checked: list[str] = [] - async def __call__(self, valid_token: UserAPIKeyAuth, /) -> Scopes: - self.calls.append(valid_token) - return self.scopes + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: + self.checked.append(model) + return self.refusals.get(model) or WebsocketFrameModelAllowed(access_groups=()) @dataclass(frozen=True, slots=True) class _Served: relay: _FakeRelay - allowlists: _FakeModelAllowlists + frame_model_gate: _FakeFrameModelGate async def _serve( @@ -125,16 +141,16 @@ async def _serve( endpoint: str, user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], - scopes: Scopes = (), + frame_model_gate: _FakeFrameModelGate | None = None, ) -> _Served: - served = _Served(relay=_FakeRelay(), allowlists=_FakeModelAllowlists(scopes)) + served: Final = _Served(relay=_FakeRelay(), frame_model_gate=frame_model_gate or _FakeFrameModelGate()) await openai_websocket_proxy_route( websocket=websocket, endpoint=endpoint, user_api_key_dict=user_api_key_dict, general_settings=general_settings, relay=served.relay, - model_allowlists=served.allowlists, + frame_model_gate=served.frame_model_gate, ) return served @@ -143,18 +159,20 @@ async def _serve( @pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch): monkeypatch.delenv("OPENAI_API_BASE", raising=False) - websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") + websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") with patch(GET_CREDENTIALS, return_value="sk-provider"): served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) assert served.relay.calls == [ _RelayCall( - target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview", + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}), forward_headers=False, endpoint=f"/{prefix}/v1/realtime", accept_websocket=False, + client_frame_model_gate=served.frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) ] assert websocket.accepts == [None] @@ -162,11 +180,23 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m assert websocket.closed is None +@pytest.mark.asyncio +async def test_openai_websocket_reports_its_route_without_the_proxy_root_path(monkeypatch): + """Route allowlists and spend rows key on the route the proxy serves, which a root path prefix is not part of""" + monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm") + websocket: Final = _FakeWebSocket("/litellm/openai/v1/responses", "") + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served: Final = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED) + + assert [call.endpoint for call in served.relay.calls] == ["/openai/v1/responses"] + + @pytest.mark.asyncio async def test_openai_websocket_accepts_first_client_subprotocol(): websocket = _FakeWebSocket( "/openai/v1/realtime", - "model=gpt-4o-realtime-preview", + "model=gpt-realtime-2.1", subprotocols="realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1", ) @@ -239,7 +269,7 @@ async def test_openai_websocket_closes_cleanly_when_workload_identity_exchange_f @pytest.mark.asyncio async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing(): - websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") + websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1") with patch(GET_CREDENTIALS, return_value=None): served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) @@ -255,7 +285,7 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing @pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) @pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings): - websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") + websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings) @@ -265,25 +295,21 @@ async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, genera assert served.relay.calls == [] -@pytest.mark.asyncio @pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) -async def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): - refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), general_settings, _FakeModelAllowlists(())) - assert refusal is _OPENAI_WS_DISABLED_REFUSAL +def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): + assert _openai_websocket_refusal(general_settings) is _OPENAI_WS_DISABLED_REFUSAL -@pytest.mark.asyncio @pytest.mark.parametrize("value", [True, "true", "True"]) -async def test_openai_websocket_refusal_is_none_for_truthy_settings(value): - settings = MappingProxyType({"enable_openai_websocket_passthrough": value}) - assert await _openai_websocket_refusal(UserAPIKeyAuth(), settings, _FakeModelAllowlists(())) is None +def test_openai_websocket_refusal_is_none_for_truthy_settings(value): + assert _openai_websocket_refusal(MappingProxyType({"enable_openai_websocket_passthrough": value})) is None @pytest.mark.asyncio async def test_openai_websocket_refusal_echoes_requested_subprotocol(): websocket = _FakeWebSocket( "/openai_passthrough/v1/realtime", - "model=gpt-4o-realtime-preview", + "model=gpt-realtime-2.1", subprotocols="realtime, openai-beta.realtime-v1", ) @@ -294,68 +320,303 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol(): assert served.relay.calls == [] -RESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( - (("gpt-4o",),), - ((), ("gpt-4o-realtime-preview",)), - (("all-team-models",), ("gpt-4o",)), - ((), ("all-proxy-models",), ("gpt-4o",)), - ((), (), (), ("gpt-4o",)), - (("*",), (), (), (), ("gpt-4o",)), +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_models", + [["gpt-5.4-mini"], ["gpt-realtime-2.1", "gpt-5.4"], []], ) -UNRESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( - (), - ((),), - (("all-proxy-models",),), - (("*",),), - (("all-team-models",), ("all-proxy-models",)), - ((), (), (), (), ()), - (("*",), ("all-proxy-models",), ("all-team-models",), (), ()), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) -async def test_openai_websocket_rejects_model_restricted_identities(scopes): - websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") - user_api_key_dict = UserAPIKeyAuth(token="hashed-fake", user_id="user-fake", team_id="team-fake") - - served = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED, scopes) - - assert "model restrictions" in websocket.error_message() - assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason) - assert served.relay.calls == [] - assert served.allowlists.calls == [user_api_key_dict] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) -async def test_openai_websocket_disabled_refusal_skips_allowlist_lookups(scopes): - allowlists = _FakeModelAllowlists(scopes) - - refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), MappingProxyType({}), allowlists) - - assert refusal is _OPENAI_WS_DISABLED_REFUSAL - assert allowlists.calls == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", UNRESTRICTED_SCOPES) -async def test_openai_websocket_allows_unrestricted_identities(scopes): - websocket = _FakeWebSocket("/openai/v1/responses", "") +async def test_openai_websocket_connects_model_restricted_keys(key_models): + websocket: Final = _FakeWebSocket("/openai/v1/responses", "") + user_api_key_dict: Final = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake") with patch(GET_CREDENTIALS, return_value="sk-provider"): - served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED, scopes) + served: Final = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED) assert len(served.relay.calls) == 1 + assert served.relay.calls[0].client_frame_model_gate is served.frame_model_gate assert websocket.sent == [] assert websocket.closed is None -@pytest.mark.asyncio -async def test_proxy_model_allowlists_reads_the_token_scopes_without_a_database(): - token: Final = UserAPIKeyAuth(models=[], team_id="team-fake", team_models=["gpt-4o"]) - with patch("litellm.proxy.proxy_server.prisma_client", None): - scopes = await _proxy_model_allowlists()(token) +ALIAS_REFUSAL: Final = WebsocketFrameModelRefused( + message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value +) - assert tuple(tuple(scope) for scope in scopes) == ((), ("gpt-4o",)) - assert _has_model_restrictions(scopes) + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "query", + [ + "model=gpt-realtime-2.1", + "model=gpt-realtime-2.1&model=gpt-realtime-2.1-mini", + "model=gpt-realtime-2.1-mini&model=gpt-realtime-2.1", + ], + ids=["one_model", "refused_model_first", "refused_model_last"], +) +async def test_openai_websocket_refuses_an_opening_model_the_frame_gate_refuses(query): + """ + The relay sends the connection URL to OpenAI unchanged, so every model it names has to pass the + same gate as a frame, or an alias the upgrade resolves to an allowed model opens the session on the + alias name itself + """ + websocket: Final = _FakeWebSocket("/openai/v1/realtime", query) + gate: Final = _FakeFrameModelGate(MappingProxyType({"gpt-realtime-2.1": ALIAS_REFUSAL})) + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served: Final = await _serve( + websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED, gate + ) + + assert websocket.error_message() == _OPENAI_WS_ALIASED_MODEL_REFUSAL + assert websocket.closed == (1008, "This key cannot use a model alias over OpenAI websocket passthrough.") + assert served.relay.calls == [] + assert "gpt-realtime-2.1" in gate.checked + + +@pytest.mark.asyncio +async def test_openai_websocket_relays_a_session_once_every_opening_model_passes_the_frame_gate(): + websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1-mini") + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served: Final = await _serve( + websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED + ) + + assert served.frame_model_gate.checked == ["gpt-realtime-2.1-mini"] + assert len(served.relay.calls) == 1 + assert websocket.closed is None + + +class _FrameAuthorizer: + """Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names.""" + + def __init__(self, token: UserAPIKeyAuth | None = None, denial: Exception | None = None) -> None: + self.token = token or UserAPIKeyAuth() + self.denial = denial + self.authorized: list[str] = [] + + async def __call__(self, websocket: object, model: str) -> UserAPIKeyAuth: + self.authorized.append(model) + if self.denial is not None: + raise self.denial + return self.token + + +SOCKET: Final = WebSocket({"type": "websocket"}, receive=AsyncMock(), send=AsyncMock()) + + +async def _gate_verdict(authorizer: _FrameAuthorizer, model: str) -> WebsocketFrameModelVerdict: + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database + return await _proxy_frame_model_gate(authorizer)(SOCKET, model) + + +def _refusal(verdict: WebsocketFrameModelVerdict) -> WebsocketFrameModelRefused: + assert isinstance(verdict, WebsocketFrameModelRefused) + return verdict + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_runs_full_auth_for_the_model_the_frame_names(): + authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])) + + verdict: Final = await _gate_verdict(authorizer, "gpt-realtime-2.1") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + assert authorizer.authorized == ["gpt-realtime-2.1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("requested_model", ["gpt-5.4", "gpt-realtime-2.1-mini"]) +async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_model): + authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])) + + refusal: Final = _refusal(await _gate_verdict(authorizer, requested_model)) + + assert requested_model in refusal.message + assert refusal.error_type == ProxyErrorTypes.key_model_access_denied.value + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "granted_models, refusal_names", + [(["gpt-5.4-mini"], "gpt-5.4"), ([], "no model grants")], + ids=["model_outside_the_grant", "no_grants"], +) +async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(granted_models, refusal_names): + from litellm.types.agents import AgentResponse + + token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake") + token.managed_agent_policy = AgentResponse( + agent_id="agent-fake", + agent_name="Agent", + agent_card_params={}, + object_permission={"object_permission_id": "grant-fake", "models": granted_models}, + ) + + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4")) + + assert refusal_names in refusal.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "token", + [ + UserAPIKeyAuth(models=["gpt-5.4-mini"], aliases={"gpt-5.4": "gpt-5.4-mini"}), + UserAPIKeyAuth(models=["gpt-5.4-mini"], team_model_aliases={"gpt-5.4": "gpt-5.4-mini"}), + ], + ids=["key_alias", "team_alias"], +) +async def test_proxy_frame_model_gate_refuses_an_alias_that_would_reach_openai_unchanged(token): + """ + The passthrough sends each frame to OpenAI as written, so OpenAI runs the + alias name itself. An alias pointing at an allowed model must not let a key + limited to that model run the model the alias is named after. + """ + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4")) + + assert refusal == WebsocketFrameModelRefused( + message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value + ) + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_a_proxy_wide_alias_for_a_key_limited_to_its_target(monkeypatch): + monkeypatch.setattr(litellm, "model_alias_map", {"gpt-5.4": "gpt-5.4-mini"}) + + refusal: Final = _refusal( + await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=["gpt-5.4-mini"])), "gpt-5.4") + ) + + assert refusal.message == _OPENAI_WS_ALIASED_MODEL_REFUSAL + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_lets_a_key_that_may_use_any_model_send_an_alias(): + token: Final = UserAPIKeyAuth(models=[], aliases={"fast": "gpt-5.4-mini"}) + + verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "fast") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]]) +async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models): + verdict: Final = await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=unrestricted_models)), "gpt-5.4") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_hands_back_the_access_groups_that_authorized_the_model(): + """The turn a frame starts is billed to the groups that let its model through, so the gate has to report them.""" + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], matched_model_access_groups=["tier-a"]) + + verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4-mini") + + assert verdict == WebsocketFrameModelAllowed(access_groups=("tier-a",)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model", + ["gpt-5.4-mini", "gpt-5.4"], + ids=["allowed", "refused_by_the_model_check"], +) +async def test_proxy_frame_model_gate_releases_the_budget_it_reserved_for_the_check(requested_model): + """ + Full auth reserves budget for the request it authorizes, but a frame check is + not a request, so a reservation the gate leaves open would hold the key's + spend counter above its real spend and refuse later calls on every route. + """ + reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": []} + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], budget_reservation=reservation) + + await _gate_verdict(_FrameAuthorizer(token), requested_model) + + assert token.budget_reservation == {**reservation, "finalized": True} + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_does_not_count_the_sessions_own_reservation_as_key_spend(): + """ + The upgrade reserves up to the key's remaining budget for the session, and the key budget check + counts open reservations as spend, so a frame checked while that reservation is open is refused + on a key with budget left, where an HTTP call with the same key goes through. + """ + connect_reservation: Final[dict[str, object]] = {"reserved_cost": 0.0001, "entries": []} + socket: Final = WebSocket( + {"type": "websocket", "state": {"budget_reservation": connect_reservation}}, receive=AsyncMock(), send=AsyncMock() + ) + + async def key_budget_check(websocket: object, model: str) -> UserAPIKeyAuth: + if connect_reservation.get("finalized") is not True: + raise litellm.BudgetExceededError(current_cost=0.0001, max_budget=0.0001) + return UserAPIKeyAuth(models=["gpt-5.4-mini"]) + + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database + verdict: Final = await _proxy_frame_model_gate(key_budget_check)(socket, "gpt-5.4-mini") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "denial, error_type, message", + [ + ( + ProxyException( + message="Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0", + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=400, + ), + "budget_exceeded", + "Budget has been exceeded! Key=hashed", + ), + ( + litellm.BudgetExceededError( + current_cost=2.0, + max_budget=1.0, + message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0", + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, + entity_id="tier-a", + ), + "budget_exceeded", + "Model access group=tier-a", + ), + ( + ProxyException( + message="team not allowed to access model. This team can only access models=['gpt-5.4-mini']", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=401, + ), + "team_model_access_denied", + "team not allowed to access model", + ), + (HTTPException(status_code=403, detail="Invalid Authorization header format"), "invalid_request_error", ""), + ], + ids=["key_budget", "group_budget", "team_allowlist", "malformed_auth"], +) +async def test_proxy_frame_model_gate_refuses_what_full_auth_denies_with_its_reason(denial, error_type, message): + """A frame is denied for the same reasons a request is, and the caller reads which reason in the error event.""" + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(denial=denial), "gpt-5.4-mini")) + + assert refusal.error_type == error_type + assert message in refusal.message + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_out_loud_when_auth_cannot_answer(): + """ + A gate that cannot answer has to close the session with a reason the caller can read. + Letting the failure escape drops the socket with no frame at all, which the caller + cannot tell apart from the network going away. + """ + refusal: Final = _refusal( + await _gate_verdict(_FrameAuthorizer(denial=ConnectionError("connection pool is exhausted")), "gpt-5.4-mini") + ) + + assert refusal == WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error") diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 182ef833069..afcb89010fe 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -347,6 +347,241 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions assert delivered_text != "" +async def _served_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, ...]: + from litellm.caching.caching import DualCache + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = ContentFilterGuardrail( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="zebra", action=ContentFilterAction("MASK"))], + event_hook="post_call", + default_on=guardrail_default_on, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + def chunk(content: str, token_ids: list[int], finish_reason: str | None = None) -> ModelResponseStream: + streamed = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)] + ) + setattr(streamed.choices[0], "token_ids", token_ids) + return streamed + + async def fake_stream(): + yield chunk("the zebra runs", [1, 2, 3]) + yield chunk("", [4], finish_reason="stop") + + served = [ + served_chunk + async for served_chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + request_data={"model": "gpt-4o-mini", "metadata": {}}, + ) + ] + return tuple(getattr(served_chunk.choices[0], "token_ids", None) for served_chunk in served) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_drops_token_ids_that_decode_to_the_unmasked_text(monkeypatch): + served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=True) + + assert served_token_ids + assert all(token_ids is None for token_ids in served_token_ids) + + +@pytest.mark.asyncio +async def test_post_call_stream_keeps_token_ids_when_no_guardrail_runs_on_the_request(monkeypatch): + served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=False) + + assert served_token_ids == ([1, 2, 3], [4]) + + +class _ChecksEachChunkText(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return None + + +async def _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, bool]: + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + monkeypatch.setattr( + litellm, + "callbacks", + [_ChecksEachChunkText(guardrail_name="chunk-check", event_hook="post_call", default_on=guardrail_default_on)], + ) + logging_obj = Logging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="per-chunk-call", + function_id="per-chunk-call", + ) + held_back_text_chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""))]) + setattr(held_back_text_chunk.choices[0], "token_ids", [220, 19, 11]) + + served = await ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_hook( + data={"model": "m", "metadata": {}, "litellm_logging_obj": logging_obj}, + response=held_back_text_chunk, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + ) + return getattr(served.choices[0], "token_ids", None), logging_obj.drops_stream_token_ids + + +@pytest.mark.asyncio +async def test_per_chunk_stream_check_drops_token_ids_served_ahead_of_the_text_it_checks(monkeypatch): + assert await _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on=True) == (None, True) + + +@pytest.mark.asyncio +async def test_per_chunk_stream_keeps_token_ids_when_its_check_does_not_run_on_the_request(monkeypatch): + assert await _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on=False) == ([220, 19, 11], False) + + +class _ScansWholeStreamFirst(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + for chunk in [chunk async for chunk in response]: + yield chunk + + +async def _logged_stream_token_ids(monkeypatch, guardrail_default_on: bool) -> object: + import asyncio + import json + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig + + guardrail = _ScansWholeStreamFirst( + guardrail_name="whole-stream-scan", event_hook="post_call", default_on=guardrail_default_on + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + logging_obj = Logging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="token-ids-call", + function_id="token-ids-call", + ) + logged = asyncio.get_running_loop().create_future() + + async def capture_logged_response(complete_streaming_response, cache_hit): + logged.set_result(complete_streaming_response) + + logging_obj._on_deferred_stream_complete = capture_logged_response + + async def sse_lines(): + for choice in ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]}, + {"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "token_ids": [20]}, + ): + yield "data: " + json.dumps( + {"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [choice]} + ) + yield "data: [DONE]" + + stream = CustomStreamWrapper( + completion_stream=HostedVLLMChatConfig().get_model_response_iterator( + streaming_response=sse_lines(), sync_stream=False, json_mode=False + ), + model="m", + custom_llm_provider="hosted_vllm", + logging_obj=logging_obj, + ) + async for _ in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=stream, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + request_data={"model": "m", "metadata": {}, "litellm_logging_obj": logging_obj}, + ): + pass + return (await asyncio.wait_for(logged, timeout=5)).choices[0].token_ids + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_drops_token_ids_from_the_response_it_logs_and_caches(monkeypatch): + assert await _logged_stream_token_ids(monkeypatch, guardrail_default_on=True) is None + + +@pytest.mark.asyncio +async def test_post_call_stream_logs_every_token_id_when_no_guardrail_runs_on_the_request(monkeypatch): + assert await _logged_stream_token_ids(monkeypatch, guardrail_default_on=False) == [16, 11, 17, 20] + + +def _completion_chunk(text: str, token_ids: list[int], finish_reason: str | None = None) -> object: + from litellm.types.utils import TextChoices, TextCompletionResponse + + return TextCompletionResponse( + object="text_completion", + choices=[TextChoices(index=0, text=text, finish_reason=finish_reason, token_ids=token_ids)], + ) + + +def _guarded_completions_stream(monkeypatch, action: str, held_back: tuple[object, ...]): + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + + monkeypatch.setattr( + litellm, + "callbacks", + [ + ContentFilterGuardrail( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="7", action=ContentFilterAction(action))], + event_hook="post_call", + default_on=True, + ) + ], + ) + + async def upstream(): + for chunk in held_back: + yield chunk + yield _completion_chunk(", 6, 7", [21, 11, 22], finish_reason="stop") + + return ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=upstream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/completions"), + request_data={"model": "qwen3", "metadata": {}}, + ) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_blocks_completions_before_serving_token_id_only_chunks(monkeypatch): + stream = _guarded_completions_stream( + monkeypatch, "BLOCK", (_completion_chunk("", [11]), _completion_chunk("", [220])) + ) + + with pytest.raises(HTTPException): + await stream.__anext__() + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_serves_no_chunk_left_empty_by_the_token_id_drop(monkeypatch): + held_back = (_completion_chunk("", [11]), _completion_chunk("", [220])) + + served = [chunk async for chunk in _guarded_completions_stream(monkeypatch, "MASK", held_back)] + + assert [(chunk.choices[0].text, chunk.choices[0].token_ids) for chunk in served] == [ + (", 6, [KEYWORD_REDACTED]", None) + ] + assert [chunk.choices[0].token_ids for chunk in held_back] == [None, None] + + @pytest.mark.asyncio async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch): from litellm.caching.caching import DualCache diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index e4af36d99f2..a50db09128f 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -42,6 +42,7 @@ from fastapi import FastAPI, HTTPException, WebSocket # test /chat/completion request to the proxy from fastapi.testclient import TestClient +from starlette.testclient import WebSocketTestSession from litellm.integrations.custom_logger import CustomLogger from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -3661,6 +3662,91 @@ def test_live_session_pre_call_abort_releases_reservations( route_call.assert_not_awaited() +def _receive_until_closed(websocket: WebSocketTestSession) -> None: + while True: + websocket.receive_text() + + +@pytest.mark.parametrize( + ("litellm_params", "expected_upstream_url"), + ( + pytest.param( + {"model": "openai/gateway-transcribe", "api_base": "https://gateway.example/v1"}, + "wss://gateway.example/v1/realtime?model=gateway-transcribe&intent=transcription", + id="gateway", + ), + pytest.param( + {"model": "xai/grok-voice", "api_base": "https://api.x.ai/v1"}, + "wss://api.x.ai/v1/realtime?model=grok-voice&intent=transcription", + id="xai", + ), + pytest.param( + {"model": "openai/gpt-live-transcribe", "api_base": "https://api.openai.com/v1"}, + "wss://api.openai.com/v1/realtime?intent=transcription", + id="openai", + ), + ), +) +def test_realtime_transcription_websocket_keeps_the_model_upstream_except_on_openai_hosts( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, + litellm_params: Mapping[str, str], + expected_upstream_url: str, +) -> None: + from starlette.websockets import WebSocketDisconnect + + from litellm.proxy import proxy_server + + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[{"model_name": "team-transcriber", "litellm_params": {**litellm_params, "api_key": "fake"}}] + ), + ) + connect: Final = MagicMock(side_effect=RuntimeError("upstream url captured")) + monkeypatch.setattr("websockets.connect", connect) + + with ( + client_no_auth.websocket_connect( + "/v1/realtime?model=team-transcriber&intent=transcription", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket, + pytest.raises(WebSocketDisconnect), + ): + _receive_until_closed(websocket) + + assert connect.call_args.args[0] == expected_upstream_url + + +def test_openai_prefixed_translation_websocket_takes_the_openai_websocket_pass_through( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.websockets import WebSocketDisconnect + + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _OPENAI_WS_DISABLED_REFUSAL + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setitem(proxy_server.app.dependency_overrides, user_api_key_auth_websocket, lambda: UserAPIKeyAuth()) + + with client_no_auth.websocket_connect( + "/openai/v1/realtime/translations?model=gpt-realtime-translate", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + refusal: Final = websocket.receive_json() + with pytest.raises(WebSocketDisconnect) as closed: + websocket.receive_json() + + assert refusal["error"]["message"] == _OPENAI_WS_DISABLED_REFUSAL.message + assert closed.value.code == 1008 + route_call.assert_not_awaited() + + @pytest.mark.parametrize("path", ["/openai/v1/live/sessions", "/v1/live/sessions", "/live/sessions"]) def test_live_session_invalid_start_is_rejected_without_routing( path: str, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3647146eb4e..9914247bc47 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import dataclasses import importlib import json import logging @@ -10,9 +11,10 @@ import subprocess import sys import time import types -from collections.abc import Callable +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone from pathlib import Path +from types import MappingProxyType from typing import Final from unittest import mock from unittest.mock import AsyncMock, MagicMock, call, create_autospec, mock_open, patch @@ -28,6 +30,7 @@ from fastapi import FastAPI, HTTPException, Request from fastapi.encoders import jsonable_encoder from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient +from starlette.datastructures import URL import litellm import litellm.proxy.proxy_server as proxy_server_module @@ -48,6 +51,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.login_throttle import LoginThrottle from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendUpdateQueue from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler from litellm.utils import _invalidate_model_cost_lowercase_map @@ -1804,9 +1808,7 @@ class _DatabaseWithOneStoredCredential: @pytest.mark.asyncio -@pytest.mark.parametrize( - "encrypted_with, asks_to_migrate", [(PUBLICLY_KNOWN_KEY, True), ("sk-some-other-key", False)] -) +@pytest.mark.parametrize("encrypted_with, asks_to_migrate", [(PUBLICLY_KNOWN_KEY, True), ("sk-some-other-key", False)]) async def test_proxy_startup_asks_to_migrate_only_when_the_database_holds_values_under_the_unsafe_key( monkeypatch, tmp_path, encrypted_with, asks_to_migrate ): @@ -1924,9 +1926,7 @@ async def test_proxy_startup_names_the_config_file_that_set_the_unsafe_key(monke from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError from litellm.proxy.proxy_server import proxy_startup_event - config_path, announced = _boot_with_general_settings( - monkeypatch, tmp_path, {"master_key": PUBLICLY_KNOWN_KEY} - ) + config_path, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": PUBLICLY_KNOWN_KEY}) with pytest.raises(UnsafeMasterKeyError): async with proxy_startup_event(FastAPI()): @@ -3807,6 +3807,24 @@ async def test_db_reload_keeps_the_resolved_value_of_a_config_owned_env_referenc assert proxy_server_module.store_model_in_db is True +@pytest.mark.asyncio +async def test_db_reload_of_a_json_string_general_settings_row_reports_no_ignored_keys( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", MagicMock(), raising=False) + proxy_config: Final = ProxyConfig() + proxy_config.settings.load_yaml({}) + monkeypatch.setattr(proxy_server_module, "general_settings", proxy_config.settings, raising=False) + + await proxy_config._update_general_settings('{"budget_rollover": true, "master_key": "sk-stored"}') + + assert proxy_config.settings.unmanaged_db_keys() == () + assert proxy_config.settings.source("master_key") != "db" + + def test_max_ui_session_budget_default_is_one_dollar(): """LIT-4662: the dashboard session budget default is a product decision; the old 0.25 default locked admins out of auto router Test Connection and the @@ -4099,7 +4117,11 @@ async def test_load_environment_variables_litellm_license_and_edge_cases(): mock_license_check = MagicMock() mock_license_check.is_premium.return_value = True - with patch("litellm.proxy.proxy_server._license_check", mock_license_check): + with ( + patch("litellm.proxy.proxy_server._license_check", mock_license_check), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.premium_user_data", None), + ): with patch.dict(os.environ, {}, clear=False): # Call the method under test proxy_config._load_environment_variables(test_config_with_license) @@ -7295,7 +7317,7 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch): db_param_value = { "max_parallel_requests": None, "allowed_models": [], - "new_key": "new_value", + "store_prompts_in_spend_logs": True, "nested": { "key1": "updated_value1", "key3": "value3", @@ -7308,7 +7330,7 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch): assert result["max_parallel_requests"] == 10 assert result["allowed_models"] == ["gpt-3.5-turbo", "gpt-4"] - assert result["new_key"] == "new_value" + assert result["store_prompts_in_spend_logs"] is True assert result["nested"] == {"key1": "value1", "key2": "value2"} @@ -8394,10 +8416,18 @@ async def test_update_general_settings_db_pass_through_endpoint_overrides_yaml_e request.headers = {} request.query_params = {} - settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam - yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in - initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here - master_key: Final = patch("litellm.proxy.proxy_server.master_key", MASTER_KEY) # test-quality-ok: a set master key is what makes a missing Authorization header a 401 + settings: Final = patch( + "litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]} + ) # test-quality-ok: the method reads this module global; no injection seam + yaml_endpoints: Final = patch( + "litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint] + ) # test-quality-ok: module global holding the YAML endpoints the fix merges in + initialize: Final = patch( + "litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock() + ) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here + master_key: Final = patch( + "litellm.proxy.proxy_server.master_key", MASTER_KEY + ) # test-quality-ok: a set master key is what makes a missing Authorization header a 401 with settings, yaml_endpoints, initialize, master_key: await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]}) @@ -11092,7 +11122,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(): def test_realtime_websocket_route_aliases_registered(): - """Realtime sessions reach the proxy via three path aliases stacked on + """Realtime sessions reach the proxy via five path aliases stacked on `realtime_websocket_endpoint`. Dropping any of them silently 405s WebSocket upgrades because the catch-all `/openai/{endpoint:path}` HTTP passthrough only declares HTTP methods. The aliases must also be @@ -11108,7 +11138,13 @@ def test_realtime_websocket_route_aliases_registered(): websocket_paths = {route.path for route in app.routes if isinstance(route, WebSocketRoute)} openai_routes = LiteLLMRoutes.openai_routes.value - for expected in ("/openai/v1/realtime", "/v1/realtime", "/realtime"): + for expected in ( + "/openai/v1/realtime", + "/v1/realtime", + "/realtime", + "/v1/realtime/translations", + "/realtime/translations", + ): assert expected in websocket_paths, ( f"{expected!r} missing from registered WebSocket routes; the " f"realtime endpoint will 405 for clients hitting this path." @@ -11127,7 +11163,7 @@ def _lit6973_fake_realtime_ws() -> MagicMock: ws = MagicMock() ws.headers = {} ws.scope = {"headers": [], "type": "websocket"} - ws.url = "ws://testserver/v1/realtime" + ws.url = URL("ws://testserver/v1/realtime") ws.accept = AsyncMock() ws.send_text = AsyncMock() ws.close = AsyncMock() @@ -13900,7 +13936,7 @@ def _stream_usage_generator_chunks(): def test_is_injected_stream_usage_artifact(): from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact - from litellm.types.utils import ModelResponseStream, Usage + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks() assert is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True @@ -13923,6 +13959,21 @@ def test_is_injected_stream_usage_artifact(): assert is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False + held_back_token_ids_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=[16])] + ) + assert is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False + + null_token_ids_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=None)] + ) + assert is_injected_stream_usage_artifact(null_token_ids_chunk) is True + + unrelated_extra_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), text="")] + ) + assert is_injected_stream_usage_artifact(unrelated_extra_chunk) is True + async def _collect_async_data_generator_frames(request_data: dict) -> list: from litellm.proxy.proxy_server import async_data_generator @@ -14060,17 +14111,18 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog): @contextlib.contextmanager -def _window_spend_enqueue_env(cached_objects: dict): +def _window_spend_enqueue_env( + cached_objects: Mapping[str, object], prisma_client: MagicMock | None = None +) -> Iterator[WindowSpendUpdateQueue]: """Point increment_spend_counters at throwaway caches and a real WindowSpendUpdateQueue, and hand back the queue to inspect.""" import litellm.proxy.proxy_server as ps from litellm.caching.dual_cache import DualCache - from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( - WindowSpendUpdateQueue, - ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - user_api_key_cache = MagicMock() - user_api_key_cache.async_get_cache = AsyncMock(side_effect=lambda key, **_: cached_objects.get(key)) + user_api_key_cache: Final = UserApiKeyCache() + for key, value in cached_objects.items(): + user_api_key_cache.set_cache(key=key, value=value) queue = WindowSpendUpdateQueue() proxy_logging_obj = MagicMock() @@ -14084,7 +14136,7 @@ def _window_spend_enqueue_env(cached_objects: dict): ) ps.user_api_key_cache = user_api_key_cache ps.spend_counter_cache = DualCache() - ps.prisma_client = None + ps.prisma_client = prisma_client ps.proxy_logging_obj = proxy_logging_obj try: yield queue @@ -14272,6 +14324,293 @@ async def test_team_window_spend_row_carries_the_request_start_time(): assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000" +_WINDOW_RESET_AT: Final = "2026-08-17T00:00:00+00:00" + + +def _window(budget_duration: str, max_budget: float) -> dict[str, object]: + return {"budget_duration": budget_duration, "max_budget": max_budget, "reset_at": _WINDOW_RESET_AT} + + +@dataclasses.dataclass(frozen=True, slots=True) +class _StoredRow: + budget_limits: list[dict[str, object]] | None + + +def _prisma_with_rows( + keys: Mapping[str, _StoredRow] = MappingProxyType({}), + teams: Mapping[str, _StoredRow] = MappingProxyType({}), +) -> MagicMock: + """Key and team tables that answer find_unique from the given rows.""" + client: Final = MagicMock() + client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=lambda where: keys.get(where["token"])) + client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=lambda where: teams.get(where["team_id"])) + return client + + +def _charged_windows(enqueued: Sequence[Mapping[str, object]]) -> list[tuple[object, object, object]]: + return [(item["entity_id"], item["window_duration"], item["spend"]) for item in enqueued] + + +@pytest.mark.asyncio +async def test_key_window_is_charged_when_the_cached_key_expired_before_the_request_ended(): + """A stream that outlives user_api_key_cache_ttl finds no cached key when its + cost lands. Its spend still belongs to the key's budget windows.""" + import litellm.proxy.proxy_server as ps + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(keys={"hashed-token": _StoredRow([_window("30d", 100.0)])}) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + enqueued: Final = await _drain(queue) + assert await ps.user_api_key_cache.async_get_cache(key="hashed-token") is None + + assert [item["entity_type"] for item in enqueued] == ["key"] + assert _charged_windows(enqueued) == [("hashed-token", "30d", pytest.approx(0.25))] + + +@pytest.mark.asyncio +async def test_team_window_is_charged_when_the_cached_team_expired_before_the_request_ended(): + import litellm.proxy.proxy_server as ps + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(teams={"team-1": _StoredRow([_window("7d", 50.0)])}) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=1.5) + enqueued: Final = await _drain(queue) + assert await ps.user_api_key_cache.async_get_cache(key="team_id:team-1") is None + + assert [item["entity_type"] for item in enqueued] == ["team"] + assert _charged_windows(enqueued) == [("team-1", "7d", pytest.approx(1.5))] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "token, keys", + [ + pytest.param("litellm_proxy_master_key", {}, id="no-key-row"), + pytest.param("hashed-token", {"hashed-token": _StoredRow(budget_limits=None)}, id="key-row-without-windows"), + ], +) +async def test_a_token_without_budget_windows_reads_the_database_once_across_requests( + token: str, keys: Mapping[str, _StoredRow] +): + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(keys=keys) + find_unique: Final = prisma_client.db.litellm_verificationtoken.find_unique + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.25) + reads_after_first_request: Final = find_unique.await_count + for _ in range(3): + await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.25) + enqueued: Final = await _drain(queue) + + assert enqueued == [] + assert reads_after_first_request > 0 + assert find_unique.await_count == reads_after_first_request + + +@pytest.mark.asyncio +async def test_every_cache_miss_charges_the_windows_the_key_row_holds_at_that_moment(): + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(keys={"hashed-token": _StoredRow([_window("30d", 100.0)])}) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_StoredRow([_window("7d", 5.0)]) + ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) + enqueued: Final = await _drain(queue) + + assert sorted(_charged_windows(enqueued), key=lambda charge: str(charge[1])) == [ + ("hashed-token", "30d", pytest.approx(0.25)), + ("hashed-token", "7d", pytest.approx(0.5)), + ] + + +@pytest.mark.asyncio +async def test_a_failed_key_row_read_skips_that_request_and_the_next_one_reads_again(): + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows() + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=ConnectionError("database is unreachable") + ) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + after_failure: Final = await _drain(queue) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_StoredRow([_window("7d", 5.0)]) + ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) + after_recovery: Final = await _drain(queue) + + assert after_failure == [] + assert _charged_windows(after_recovery) == [("hashed-token", "7d", pytest.approx(0.5))] + + +@pytest.mark.asyncio +async def test_a_late_request_after_a_key_update_charges_windows_the_key_did_not_have_before(): + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.auth_checks import delete_cache_key_object + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(keys={"hashed-token": _StoredRow(budget_limits=None)}) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_StoredRow([_window("7d", 5.0)]) + ) + await delete_cache_key_object( + hashed_token="hashed-token", user_api_key_cache=ps.user_api_key_cache, proxy_logging_obj=None + ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) + enqueued: Final = await _drain(queue) + + assert _charged_windows(enqueued) == [("hashed-token", "7d", pytest.approx(0.5))] + + +@pytest.mark.asyncio +async def test_a_key_update_that_lands_during_the_row_read_does_not_hide_the_windows_it_adds(): + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.auth_checks import delete_cache_key_object + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows() + updated_row: Final = _StoredRow([_window("7d", 5.0)]) + + async def read_the_row_from_before_the_update(where: Mapping[str, str]) -> _StoredRow: + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=updated_row) + await delete_cache_key_object( + hashed_token=where["token"], user_api_key_cache=ps.user_api_key_cache, proxy_logging_obj=None + ) + return _StoredRow(budget_limits=None) + + prisma_client.db.litellm_verificationtoken.find_unique = read_the_row_from_before_the_update + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + ps.spend_counter_cache.set_cache(key="spend:key:hashed-token", value=0.0) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + _ = await _drain(queue) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) + after_the_update: Final = await _drain(queue) + + assert _charged_windows(after_the_update) == [("hashed-token", "7d", pytest.approx(0.5))] + + +@pytest.mark.asyncio +async def test_late_requests_from_many_keys_with_windows_do_not_push_out_a_known_no_window_key(): + import litellm.proxy.proxy_server as ps + from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY + from litellm.proxy.proxy_server import increment_spend_counters + + windowed_tokens: Final = tuple(f"windowed-{index}" for index in range(DEFAULT_MAX_SIZE_IN_MEMORY)) + rows: Final = { + "no-windows": _StoredRow(budget_limits=None), + **{token: _StoredRow([_window("7d", 5.0)]) for token in windowed_tokens}, + } + prisma_client: Final = _prisma_with_rows(keys=rows) + find_unique: Final = prisma_client.db.litellm_verificationtoken.find_unique + + def no_window_key_reads() -> int: + return sum(1 for call in find_unique.await_args_list if call.kwargs["where"]["token"] == "no-windows") + + with _window_spend_enqueue_env({}, prisma_client=prisma_client): + await increment_spend_counters(token="no-windows", team_id=None, user_id=None, response_cost=0.1) + reads_after_the_first_request: Final = no_window_key_reads() + for token in windowed_tokens: + await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.1) + ps.spend_counter_cache.set_cache(key="spend:key:no-windows", value=0.1) + await increment_spend_counters(token="no-windows", team_id=None, user_id=None, response_cost=0.1) + + assert no_window_key_reads() == reads_after_the_first_request + + +@pytest.mark.asyncio +async def test_a_slow_window_row_read_does_not_hold_back_the_key_spend_counter(): + import litellm.proxy.proxy_server as ps + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows() + read_started: Final = asyncio.Event() + finish_read: Final = asyncio.Event() + + async def slow_row_read(where: Mapping[str, str]) -> _StoredRow: + read_started.set() + await finish_read.wait() + return _StoredRow([_window("7d", 5.0)]) + + prisma_client.db.litellm_verificationtoken.find_unique = slow_row_read + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + ps.spend_counter_cache.set_cache(key="spend:key:hashed-token", value=0.0) + ps.spend_counter_cache.set_cache(key="spend:key:hashed-token:window:7d", value=1.0) + charge: Final = asyncio.create_task( + increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + ) + assert await asyncio.wait_for(read_started.wait(), timeout=5) + key_spend_while_the_row_read_waits: Final = ps.spend_counter_cache.get_cache(key="spend:key:hashed-token") + finish_read.set() + await charge + window_spend: Final = ps.spend_counter_cache.get_cache(key="spend:key:hashed-token:window:7d") + charged: Final = await _drain(queue) + + assert key_spend_while_the_row_read_waits == pytest.approx(0.25) + assert window_spend == pytest.approx(1.25) + assert _charged_windows(charged) == [("hashed-token", "7d", pytest.approx(0.25))] + + +@pytest.mark.asyncio +async def test_a_late_request_that_misses_the_updated_team_charges_windows_the_team_did_not_have_before(): + """Another worker's /team/update only evicts this worker's cached team, so its next late cost reads the row.""" + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(teams={"team-1": _StoredRow(budget_limits=None)}) + updated_row: Final = _StoredRow([_window("7d", 50.0)]) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=1.5) + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=updated_row) + await refresh_cached_team( + team_row=LiteLLM_TeamTable(team_id="team-1", budget_limits=updated_row.budget_limits), + user_api_key_cache=ps.user_api_key_cache, + proxy_logging_obj=ps.proxy_logging_obj, + ) + ps.user_api_key_cache.delete_cache(key="team_id:team-1") + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=2.0) + enqueued: Final = await _drain(queue) + + assert _charged_windows(enqueued) == [("team-1", "7d", pytest.approx(2.0))] + + +@pytest.mark.asyncio +async def test_late_costs_for_many_tokens_without_windows_leave_the_cached_team_in_place(): + from litellm.proxy.proxy_server import increment_spend_counters + + team_obj: Final = MagicMock() + team_obj.budget_limits = [_window("7d", 50.0)] + + with _window_spend_enqueue_env({"team_id:team-1": team_obj}, prisma_client=_prisma_with_rows()) as queue: + for token_number in range(250): + await increment_spend_counters( + token=f"token-{token_number}", team_id=None, user_id=None, response_cost=0.01 + ) + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=1.5) + enqueued: Final = await _drain(queue) + + assert _charged_windows(enqueued) == [("team-1", "7d", pytest.approx(1.5))] + + def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client = MagicMock() client.connect = AsyncMock(side_effect=connect_error) @@ -15829,6 +16168,166 @@ async def test_delete_config_general_settings_still_removes_a_key_the_database_o assert "max_request_size_mb" not in pc.settings +@pytest.mark.asyncio +async def test_delete_config_general_settings_removes_a_stored_key_no_settings_api_manages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + stored_row: Final = {"enforce_user_param": True, "store_prompts_in_spend_logs": True} + pc.settings.apply_db_row("general_settings", stored_row) + fake_prisma: Final = _fake_prisma_with_config(stored_row) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + + admin: Final = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + await delete_config_general_settings( + data=ConfigFieldDelete(field_name="enforce_user_param", config_type="general_settings"), + user_api_key_dict=admin, + ) + + written: Final = json.loads(fake_prisma.db.litellm_config.upsert.await_args.kwargs["data"]["update"]["param_value"]) + assert written == {"store_prompts_in_spend_logs": True} + assert pc.settings.unmanaged_db_keys() == () + + +@pytest.mark.asyncio +async def test_delete_config_general_settings_stops_serving_the_removed_pass_through_routes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import FastAPI + from starlette.routing import Route + + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import _registered_pass_through_routes + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER + + endpoint: Final = {"id": "stored-route", "path": "/stored-route", "target": "http://example.com"} + stored_row: Final = {"pass_through_endpoints": [endpoint], "store_prompts_in_spend_logs": True} + app: Final = FastAPI() + monkeypatch.setattr(proxy_server_module, "app", app) + monkeypatch.setattr(proxy_server_module, "premium_user", True) + monkeypatch.setattr(proxy_server_module, "config_passthrough_endpoints", None) + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + pc.settings.apply_db_row("general_settings", stored_row) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", _fake_prisma_with_config(stored_row)) + + def served_paths() -> frozenset[str]: + return frozenset( + route.path + for route in app.router.routes + if isinstance(route, Route) and getattr(route.endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is True + ) + + saved_registry: Final = dict(_registered_pass_through_routes) + try: + _registered_pass_through_routes.clear() + await pc.serve_pass_through_endpoints([endpoint]) + served_before: Final = served_paths() + await delete_config_general_settings( + data=ConfigFieldDelete(field_name="pass_through_endpoints", config_type="general_settings"), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + served_after: Final = served_paths() + finally: + _registered_pass_through_routes.clear() + _registered_pass_through_routes.update(saved_registry) + + assert "/stored-route" in served_before + assert "/stored-route" not in served_after + + +@pytest.mark.asyncio +async def test_delete_config_general_settings_removes_an_ignored_stored_key_the_config_file_also_sets( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {"enforce_user_param": False}}) + stored_row: Final = {"enforce_user_param": True, "store_prompts_in_spend_logs": True} + pc.settings.apply_db_row("general_settings", stored_row) + fake_prisma: Final = _fake_prisma_with_config(stored_row) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + + admin: Final = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + await delete_config_general_settings( + data=ConfigFieldDelete(field_name="enforce_user_param", config_type="general_settings"), + user_api_key_dict=admin, + ) + + written: Final = json.loads(fake_prisma.db.litellm_config.upsert.await_args.kwargs["data"]["update"]["param_value"]) + assert written == {"store_prompts_in_spend_logs": True} + assert pc.settings["enforce_user_param"] is False + assert pc.settings.unmanaged_db_keys() == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_row", ['{"max_request_size_mb": 42, "disable_custom_api_keys": true}', "not json"]) +@pytest.mark.parametrize("field_name", ["max_request_size_mb", "disable_custom_api_keys", "not_a_setting"]) +async def test_delete_config_general_settings_leaves_a_string_row_untouched( + monkeypatch: pytest.MonkeyPatch, stored_row: str, field_name: str +) -> None: + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + fake_prisma: Final = _fake_prisma_with_config(stored_row) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + + admin: Final = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as refused: + await delete_config_general_settings( + data=ConfigFieldDelete(field_name=field_name, config_type="general_settings"), + user_api_key_dict=admin, + ) + + assert (refused.value.status_code, refused.value.detail) == ( + 400, + { + "error": "general_settings is stored in the database as something other than a JSON object, so no field " + "can be removed from it" + }, + ) + fake_prisma.db.litellm_config.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_config_general_settings_refuses_a_key_neither_managed_nor_stored( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + fake_prisma: Final = _fake_prisma_with_config({"store_prompts_in_spend_logs": True}) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + + admin: Final = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as refused: + await delete_config_general_settings( + data=ConfigFieldDelete(field_name="enforce_user_param", config_type="general_settings"), + user_api_key_dict=admin, + ) + + assert refused.value.status_code == 400 + fake_prisma.db.litellm_config.upsert.assert_not_awaited() + + @pytest.mark.asyncio async def test_config_field_info_reports_the_declared_value_of_a_config_owned_secret(monkeypatch): from litellm.proxy.proxy_server import ProxyConfig, get_config_general_settings diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 68eb05e60dd..2f85d1c7516 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -25,7 +25,7 @@ import pytest import respx import tiktoken from dotenv import load_dotenv -from pydantic import BaseModel, ConfigDict, JsonValue +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter from tokenizers import Tokenizer from tokenizers.models import WordLevel from tokenizers.pre_tokenizers import Whitespace @@ -82,6 +82,7 @@ from litellm.proxy.proxy_server import ( ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from litellm.types.utils import ( ADDRESSED_RESPONSE_ID_FIELD, @@ -1741,7 +1742,10 @@ _OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5" async def _ladder( - counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True + counter: _ScriptedCounter, + settings: TokenCountSettings, + local_estimate_allowed: bool = True, + forwarded_edits: bool = False, ) -> TokenCountResponse | None: return await _provider_token_count( count_with=partial( @@ -1759,6 +1763,7 @@ async def _ladder( request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=local_estimate_allowed, + forwarded_edits=forwarded_edits, ) @@ -1920,6 +1925,23 @@ async def test_uncountable_model_without_framing_raises(): assert raised.value.code == "404" +@pytest.mark.parametrize( + "settings", + [ + TokenCountSettings(framing=None, count_model=None, tool_framing=300), + TokenCountSettings(framing=None, count_model=None, responses_framing=145), + TokenCountSettings(framing=None, count_model=None, responses_tool_framing=100), + ], +) +@pytest.mark.asyncio +async def test_uncountable_model_with_only_a_request_kind_framing_falls_through_to_the_local_estimate( + settings: TokenCountSettings, +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) + + assert isinstance(await _ladder(counter, settings), ProviderCountFailure) + + @pytest.mark.parametrize( ("count_model", "counted_as"), [ @@ -1986,6 +2008,26 @@ async def test_configured_token_count_model_error_stands_even_with_framing( assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints +@pytest.mark.parametrize(("status_code", "region_hints"), [(400, 0), (403, 1)]) +@pytest.mark.asyncio +async def test_a_pinned_token_count_model_rejecting_forwarded_edits_answers_without_the_region_hint( + status_code: int, region_hints: int +): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder( + counter, + TokenCountSettings(framing=None, count_model=replace(_OPUS_IN_US_EAST_1, configured=True)), + forwarded_edits=True, + ) + + assert raised.value.code == str(status_code) + assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints + + @pytest.mark.parametrize( ("count_model", "region_hints"), [(_OPUS_IN_US_EAST_1, frozenset({True})), (replace(_OPUS_IN_US_EAST_1, region=None), frozenset())], @@ -2745,6 +2787,480 @@ async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_ assert aliased.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + framing +@pytest.mark.asyncio +async def test_a_tool_request_counts_its_framing_in_place_of_the_plain_one_per_api(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": f"sk-unused-{name}"}, + "model_info": model_info, + } + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + deployment("plain", {}), + deployment("chat", {"token_count_framing": 7, "token_count_tool_framing": 50}), + deployment( + "per-api", + { + "token_count_framing": 7, + "token_count_tool_framing": 50, + "token_count_responses_framing": 5, + "token_count_responses_tool_framing": 2, + }, + ), + ] + ), + ) + messages: Final = [{"role": "user", "content": "What is the weather in Paris?"}] + tools: Final = [_KIMI_WEATHER] + + async def count(model: str, request_tools: list[dict[str, object]] | None, responses: bool) -> int: + response: Final = await count_request_tokens( + request=TokenCountRequest(model=model, messages=messages, tools=request_tools), + responses_request={"model": model, "input": "What is the weather in Paris?"} if responses else None, + ) + return response.total_tokens + + plain_with_tools: Final = await count("plain", tools, responses=False) + plain_without_tools: Final = await count("plain", None, responses=False) + + assert await count("plain", tools, responses=True) == plain_with_tools + assert await count("chat", tools, responses=False) == plain_with_tools + 50 + assert await count("chat", None, responses=False) == plain_without_tools + 7 + assert await count("chat", [], responses=False) == plain_without_tools + 7 + assert await count("chat", tools, responses=True) == plain_with_tools + 50 + assert await count("chat", None, responses=True) == plain_without_tools + 7 + assert await count("per-api", tools, responses=False) == plain_with_tools + 50 + assert await count("per-api", tools, responses=True) == plain_with_tools + 2 + assert await count("per-api", None, responses=True) == plain_without_tools + 5 + + +_GET_WEATHER: Final = { + "name": "get_weather", + "description": "Get the current weather for a city", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, + "required": ["city"], + }, +} +_THREE_TOOLS: Final = ( + _GET_WEATHER, + { + "name": "get_time", + "description": "Get the local time in a timezone", + "input_schema": { + "type": "object", + "properties": {"timezone": {"type": "string", "description": "IANA timezone"}}, + "required": ["timezone"], + }, + }, + { + "name": "search_docs", + "description": "Search internal documents and return the top matches", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}, "limit": {"type": "integer", "description": "Max results"}}, + "required": ["query"], + }, + }, +) +_FORM_FIELDS: Final = tuple(f"field_{index}" for index in range(10)) +_FILL_FORM: Final = { + "name": "fill_form", + "description": "Fill a form", + "input_schema": { + "type": "object", + "properties": { + field: {"type": "string", "description": f"Value for field number {index}"} + for index, field in enumerate(_FORM_FIELDS) + }, + "required": list(_FORM_FIELDS), + }, +} +_WEATHER_SYSTEM: Final = "You are a helpful assistant that answers weather questions briefly." +_AZURE_GROK_4_1: Final = "azure_ai/grok-4-1-fast-reasoning" +_FIREWORKS_DEEPSEEK_V4_1: Final = "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash" +_XAI_GROK_4_3: Final = "xai/grok-4.3" +_NATIVE_RESPONSES_BASE: Final = { + _AZURE_GROK_4_1: "https://unused.services.ai.azure.com/models", + _FIREWORKS_DEEPSEEK_V4_1: "https://unused.invalid", + _XAI_GROK_4_3: "https://unused.invalid", +} + + +def _responses_function(tool: Mapping[str, object], **strict: bool | None) -> dict[str, object]: + return { + "type": "function", + "name": tool["name"], + "description": tool["description"], + "parameters": tool["input_schema"], + **strict, + } + + +def _get_weather_copies(count: int) -> tuple[dict[str, object], ...]: + return tuple({**_GET_WEATHER, "name": f"get_weather_{index}"} for index in range(count)) + + +# usage.input_tokens plus cache_read_input_tokens that POST /v1/messages billed for these requests with the user turn +# "What is the weather in Paris?" through Azure AI Foundry grok-4-1-fast-reasoning, Fireworks deepseek-v4p1-flash and +# xAI grok-4.3 on 2026-10-29, and usage.prompt_tokens that POST /v1/chat/completions billed through DeepSeek's API on +# 2026-10-30 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "system", "tools", "billed_prompt_tokens"), + [ + (_AZURE_GROK_4_1, None, (), 163), + (_AZURE_GROK_4_1, None, (_GET_WEATHER,), 357), + (_AZURE_GROK_4_1, None, _THREE_TOOLS, 445), + (_AZURE_GROK_4_1, None, (_FILL_FORM,), 550), + (_AZURE_GROK_4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 368), + (_FIREWORKS_DEEPSEEK_V4_1, None, (), 37), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_GET_WEATHER,), 304), + (_FIREWORKS_DEEPSEEK_V4_1, None, _THREE_TOOLS, 430), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_FILL_FORM,), 550), + (_FIREWORKS_DEEPSEEK_V4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 315), + (_XAI_GROK_4_3, None, (), 199), + (_XAI_GROK_4_3, None, (_GET_WEATHER,), 283), + (_XAI_GROK_4_3, None, _THREE_TOOLS, 402), + (_XAI_GROK_4_3, None, (_FILL_FORM,), 530), + (_XAI_GROK_4_3, _WEATHER_SYSTEM, (_GET_WEATHER,), 297), + ("deepseek/deepseek-v4-flash", None, (), 37), + ("deepseek/deepseek-v4-flash", None, (_GET_WEATHER,), 304), + ("deepseek/deepseek-v4-flash", None, _THREE_TOOLS, 430), + ("deepseek/deepseek-v4-flash", _WEATHER_SYSTEM, (_GET_WEATHER,), 315), + ("deepseek/deepseek-v4-pro", None, (), 90), + ("deepseek/deepseek-v4-pro", None, (_GET_WEATHER,), 367), + ("deepseek/deepseek-v4-pro", None, _THREE_TOOLS, 493), + ("deepseek/deepseek-v4-pro", _WEATHER_SYSTEM, (_GET_WEATHER,), 378), + ], +) +async def test_local_estimate_matches_the_billed_prompt_with_the_shipped_framing( + monkeypatch: pytest.MonkeyPatch, + model: str, + system: str | None, + tools: tuple[dict[str, object], ...], + billed_prompt_tokens: int, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "billed", + "litellm_params": {"model": model, "api_key": "unused", "api_base": "https://unused.invalid"}, + } + ] + ), + ) + + counted: Final = await token_counter( + request=TokenCountRequest( + model="billed", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=list(tools) or None, + system=system, + ) + ) + + assert counted.count_source == "local_estimate" + assert abs(counted.total_tokens - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 + + +# usage.input_tokens that POST /v1/responses billed for these requests with the input "What is the weather in Paris?" +# through Azure AI Foundry grok-4-1-fast-reasoning, Fireworks deepseek-v4p1-flash and xAI grok-4.3 on 2026-10-29 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "system", "tools", "billed_prompt_tokens"), + [ + (_AZURE_GROK_4_1, None, (), 159), + (_AZURE_GROK_4_1, None, (_GET_WEATHER,), 159), + (_AZURE_GROK_4_1, None, _THREE_TOOLS, 257), + (_AZURE_GROK_4_1, None, (_FILL_FORM,), 352), + (_AZURE_GROK_4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 170), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_GET_WEATHER,), 304), + (_XAI_GROK_4_3, None, (_GET_WEATHER,), 283), + ], +) +async def test_responses_estimate_matches_the_billed_responses_prompt_with_the_shipped_framing( + monkeypatch: pytest.MonkeyPatch, + model: str, + system: str | None, + tools: tuple[dict[str, object], ...], + billed_prompt_tokens: int, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "billed", + "litellm_params": {"model": model, "api_key": "unused", "api_base": _NATIVE_RESPONSES_BASE[model]}, + } + ] + ), + ) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="billed", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_responses_function(tool) for tool in tools] or None, + system=system, + ), + responses_request={"model": "billed", "input": "What is the weather in Paris?"}, + ) + + assert counted.count_source == "local_estimate" + assert abs(counted.total_tokens - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 + + +def _single_deployment_router(litellm_params: dict[str, object], model_info: dict[str, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": {"api_key": "unused", **litellm_params}, + "model_info": model_info, + } + ] + ) + + +async def _count_with_router( + monkeypatch: pytest.MonkeyPatch, router: Router, tools: tuple[dict[str, object], ...], responses: bool +) -> int: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=list(tools) or None, + ), + responses_request={"model": "counted", "input": "What is the weather in Paris?"} if responses else None, + ) + return counted.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tools", [(), (_GET_WEATHER,), _THREE_TOOLS, (_FILL_FORM,)]) +@pytest.mark.parametrize("responses", [False, True]) +async def test_a_custom_deployment_name_counts_like_the_base_model_it_names( + monkeypatch: pytest.MonkeyPatch, tools: tuple[dict[str, object], ...], responses: bool +): + api_base: Final = _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1] + custom_name: Final = _single_deployment_router( + {"model": "azure_ai/prod-grok", "api_base": api_base}, {"base_model": _AZURE_GROK_4_1} + ) + row_name: Final = _single_deployment_router({"model": _AZURE_GROK_4_1, "api_base": api_base}, {}) + + assert await _count_with_router(monkeypatch, custom_name, tools, responses) == await _count_with_router( + monkeypatch, row_name, tools, responses + ) + + +# usage.input_tokens that POST /v1/responses billed through Azure AI Foundry grok-4-1-fast-reasoning on 2026-10-29 for +# the input "What is the weather in Paris?" with these function tools, strict left unset or sent as false or null, +# the same on the .services.ai.azure.com and .openai.azure.com hosts +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tools", "strict", "billed_prompt_tokens"), + [ + ((_GET_WEATHER,), {}, 159), + ((_GET_WEATHER,), {"strict": False}, 155), + ((_GET_WEATHER,), {"strict": None}, 155), + (_THREE_TOOLS, {}, 257), + (_get_weather_copies(5), {}, 345), + (_get_weather_copies(10), {}, 575), + (_get_weather_copies(10), {"strict": False}, 535), + (_get_weather_copies(10), {"strict": None}, 535), + ], +) +@pytest.mark.parametrize( + ("litellm_params", "model_info"), + [ + ({"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {}), + ( + {"model": "azure/grok-4-1-fast-reasoning", "api_base": "https://unused.openai.azure.com/openai/v1"}, + {"base_model": _AZURE_GROK_4_1}, + ), + ], +) +async def test_azure_responses_estimate_stays_at_or_just_over_the_billed_prompt_as_function_tools_are_added( + monkeypatch: pytest.MonkeyPatch, + tools: tuple[dict[str, object], ...], + strict: dict[str, bool | None], + billed_prompt_tokens: int, + litellm_params: dict[str, object], + model_info: dict[str, object], +): + router: Final = _single_deployment_router(litellm_params, model_info) + + counted: Final = await _count_with_router( + monkeypatch, router, tuple(_responses_function(tool, **strict) for tool in tools), responses=True + ) + + assert billed_prompt_tokens <= counted <= billed_prompt_tokens * 1.02 + + +# POST /v1/responses billed the same usage.input_tokens with strict unset and false on 2026-10-29, 283 for one tool +# through xAI grok-4.3 +@pytest.mark.asyncio +async def test_a_responses_host_that_bills_an_unset_strict_like_false_counts_them_alike( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _single_deployment_router( + {"model": _XAI_GROK_4_3, "api_base": _NATIVE_RESPONSES_BASE[_XAI_GROK_4_3]}, {} + ) + + unset: Final = await _count_with_router(monkeypatch, router, (_responses_function(_GET_WEATHER),), responses=True) + + assert unset == await _count_with_router( + monkeypatch, router, (_responses_function(_GET_WEATHER, strict=False),), responses=True + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("built_in_tool", "responses"), + [({"type": "web_search"}, True), ({"type": "web_search_20250305", "name": "web_search"}, False)], +) +async def test_a_request_whose_only_tools_are_built_in_counts_like_one_without_tools( + monkeypatch: pytest.MonkeyPatch, built_in_tool: dict[str, object], responses: bool +): + router: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + + assert await _count_with_router(monkeypatch, router, (built_in_tool,), responses) == await _count_with_router( + monkeypatch, router, (), responses + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tools", [(), (_GET_WEATHER,), _THREE_TOOLS]) +async def test_a_responses_count_the_proxy_bridges_to_chat_completions_frames_like_the_chat_count( + monkeypatch: pytest.MonkeyPatch, tools: tuple[dict[str, object], ...] +): + bridged: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": "https://unused.cognitiveservices.azure.com/models"}, {} + ) + native: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + chat_count: Final = await _count_with_router(monkeypatch, bridged, tools, responses=False) + + assert await _count_with_router(monkeypatch, bridged, tools, responses=True) == chat_count + assert await _count_with_router(monkeypatch, native, tools, responses=True) != chat_count + + +@pytest.mark.asyncio +@pytest.mark.parametrize("from_router_defaults", [False, True]) +async def test_a_responses_count_that_asks_for_the_chat_completions_bridge_frames_like_the_chat_count( + monkeypatch: pytest.MonkeyPatch, from_router_defaults: bool +): + router: Final = Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": { + "api_key": "unused", + "model": _AZURE_GROK_4_1, + "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1], + }, + "model_info": {}, + } + ], + default_litellm_params={"use_chat_completions_api": True} if from_router_defaults else None, + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_GET_WEATHER], + ), + responses_request={ + "model": "counted", + "input": "What is the weather in Paris?", + **({} if from_router_defaults else {"use_chat_completions_api": True}), + }, + ) + + assert counted.total_tokens == chat_count + + +@pytest.mark.asyncio +async def test_a_responses_count_frames_for_the_api_base_router_defaults_send_inference_to( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": { + "api_key": "unused", + "model": _AZURE_GROK_4_1, + "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1], + }, + "model_info": {}, + } + ], + default_litellm_params={"api_base": "https://unused.cognitiveservices.azure.com/models"}, + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + assert await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=True) == chat_count + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_chat_completions_api", ["maybe", {"a": 1}]) +async def test_a_responses_count_with_a_malformed_bridge_flag_reads_it_the_way_the_responses_route_does( + monkeypatch: pytest.MonkeyPatch, use_chat_completions_api: object +): + router: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_GET_WEATHER], + ), + responses_request={ + "model": "counted", + "input": "What is the weather in Paris?", + "use_chat_completions_api": use_chat_completions_api, + }, + ) + + assert counted.total_tokens == chat_count + + +@pytest.mark.asyncio +async def test_an_azure_responses_count_leaves_strict_unset_for_templates_not_measured_there( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _single_deployment_router( + {"model": "azure_ai/deepseek-v4.1-flash", "api_base": "https://unused.services.ai.azure.com/models"}, {} + ) + + async def counted(**strict: bool) -> int: + return await _count_with_router(monkeypatch, router, (_responses_function(_GET_WEATHER, **strict),), True) + + assert await counted() < await counted(strict=True) + + _KIMI_WEATHER: Final = { "type": "function", "function": { @@ -3172,15 +3688,14 @@ _READ_FILE_TOOL: Final[dict[str, JsonValue]] = { "name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, } -_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = { - "edits": [ - { - "type": "clear_tool_uses_20250919", - "trigger": {"type": "input_tokens", "value": 500}, - "keep": {"type": "tool_uses", "value": 1}, - } - ] -} +_CLEAR_TOOL_USES_EDITS: Final[list[JsonValue]] = [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } +] +_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = {"edits": _CLEAR_TOOL_USES_EDITS} _INVALID_COMPACT_EDIT: Final[dict[str, JsonValue]] = { "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] } @@ -3516,6 +4031,28 @@ async def test_gemini_messages_count_tokens_counts_the_unedited_body_when_drop_p assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("router_defaults", "request_options", "expected_status"), + (({}, {"drop_params": True}, None), ({"drop_params": True}, {}, None), ({}, {}, 400)), + ids=("request-drop-params", "router-default-drop-params", "no-drop-params"), +) +async def test_messages_count_tokens_for_a_model_with_no_deployment_drops_an_invalid_edit_only_under_drop_params( + monkeypatch: pytest.MonkeyPatch, + router_defaults: dict[str, JsonValue], + request_options: dict[str, JsonValue], + expected_status: int | None, +) -> None: + monkeypatch.setenv("GEMINI_API_KEY", "test-key") + with patch.object(litellm, "drop_params", False): + response, _ = await _count_through_the_messages_endpoint( + Router(model_list=[_gemini_deployment("gemini-agent", {})], default_litellm_params=router_defaults), + {**request_options, "model": "gemini/gemini-3.8-pro", "context_management": _INVALID_COMPACT_EDIT}, + ) + + assert (response.status_code if isinstance(response, JSONResponse) else None) == expected_status + + @pytest.mark.asyncio async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop_params() -> None: with patch.object(litellm, "drop_params", False): @@ -3529,6 +4066,576 @@ async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop assert sent_bodies == () +_MISTRAL_MODEL: Final = "mistral/mistral-large-latest" + + +def _mistral_router( + litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), + router_defaults: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Router: + return Router( + model_list=[ + { + "model_name": "mistral-agent", + "litellm_params": {"model": _MISTRAL_MODEL, "api_key": "test-key", **litellm_params}, + } + ], + default_litellm_params=dict(router_defaults), + ) + + +async def _messages_endpoint_count(router: Router, request_body: Mapping[str, JsonValue]) -> object: + import litellm.proxy.anthropic_endpoints.endpoints as anthropic_endpoints + import litellm.proxy.proxy_server as proxy_server + + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + with ( + patch.object(proxy_server, "llm_router", router), + patch.object(anthropic_endpoints, "read_request_body", AsyncMock(return_value=dict(request_body))), + ): + return await anthropic_count_tokens(mock_request, MagicMock()) + + +class _MessagesCountReply(BaseModel): + model_config = ConfigDict(frozen=True) + input_tokens: int + + +async def _input_tokens(router: Router, request_body: Mapping[str, JsonValue]) -> int: + return _MessagesCountReply.model_validate(await _messages_endpoint_count(router, request_body)).input_tokens + + +async def _messages_the_inference_polyfill_sends(model: str) -> list[dict[str, JsonValue]]: + from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request + + prepared: Final = await prepare_context_managed_request( + model=model, + messages=[dict(message) for message in _agent_turns_with_tool_results(3)], + tools=[dict(_READ_FILE_TOOL)], + system=None, + context_management_spec=dict(_CLEAR_TOOL_USES), + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + ) + assert prepared is not None + return TypeAdapter(list[dict[str, JsonValue]]).validate_python(prepared.messages) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_counts_the_body_the_inference_polyfill_sends_on_a_bridged_provider() -> None: + router: Final = _mistral_router() + unedited_body: Final[dict[str, JsonValue]] = { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + unedited: Final = await _input_tokens(router, unedited_body) + polyfilled: Final = await _input_tokens( + router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_MISTRAL_MODEL)} + ) + + assert edited < unedited + assert edited == polyfilled + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "router_defaults", "request_options"), + ( + ({"additional_drop_params": ["context_management"]}, {}, {}), + ({}, {"additional_drop_params": ["context_management"]}, {}), + ({}, {}, {"additional_drop_params": ["context_management"]}), + ), + ids=("deployment-drops", "router-default-drops", "request-drops"), +) +async def test_messages_count_tokens_on_a_bridged_provider_keeps_the_unedited_count_when_context_management_is_dropped( + deployment_params: dict[str, JsonValue], + router_defaults: dict[str, JsonValue], + request_options: dict[str, JsonValue], +) -> None: + router: Final = _mistral_router(deployment_params, router_defaults) + unedited_body: Final[dict[str, JsonValue]] = { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + **request_options, + } + + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + + assert edited == await _input_tokens(router, unedited_body) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_bridged_provider_rejects_an_invalid_edit_without_drop_params() -> None: + with patch.object(litellm, "drop_params", False): + response: Final = await _messages_endpoint_count( + _mistral_router(), + { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _INVALID_COMPACT_EDIT, + }, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_bridged_provider_counts_without_an_invalid_edit_under_drop_params() -> None: + router: Final = _mistral_router() + unedited_body: Final[dict[str, JsonValue]] = { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "drop_params": True, + } + + with patch.object(litellm, "drop_params", False): + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _INVALID_COMPACT_EDIT}) + unedited: Final = await _input_tokens(router, unedited_body) + + assert edited == unedited + + +def _turns_compacted_by_the_client() -> list[dict[str, JsonValue]]: + return [ + *_agent_turns_with_tool_results(3)[:-1], + {"role": "assistant", "content": [{"type": "compaction", "content": "The agent read files 0, 1 and 2."}]}, + {"role": "user", "content": "Summarize what you read."}, + ] + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_bridged_provider_counts_only_what_follows_a_client_compaction_block() -> None: + from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request + + router: Final = _mistral_router() + sent: Final = await prepare_context_managed_request( + model=_MISTRAL_MODEL, + messages=[dict(message) for message in _turns_compacted_by_the_client()], + tools=[dict(_READ_FILE_TOOL)], + system=None, + context_management_spec=None, + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + ) + assert sent is not None + + compacted: Final = await _input_tokens( + router, {"model": "mistral-agent", "messages": _turns_compacted_by_the_client(), "tools": [_READ_FILE_TOOL]} + ) + uncompacted: Final = await _input_tokens( + router, {"model": "mistral-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL]} + ) + sent_body_count: Final = await _input_tokens( + router, + { + "model": "mistral-agent", + "messages": TypeAdapter(list[dict[str, JsonValue]]).validate_python(sent.messages), + "system": TypeAdapter(JsonValue).validate_python(sent.system), + "tools": [_READ_FILE_TOOL], + }, + ) + + assert compacted < uncompacted + assert compacted == sent_body_count + + +@pytest.mark.asyncio +async def test_messages_count_tokens_applies_the_edits_a_responses_api_bridge_polyfills_for_a_non_openai_provider() -> ( + None +): + perplexity_model: Final = "perplexity/perplexity/sonar" + router: Final = Router( + model_list=[{"model_name": "sonar-agent", "litellm_params": {"model": perplexity_model, "api_key": "test-key"}}] + ) + unedited_body: Final[dict[str, JsonValue]] = { + "model": "sonar-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + unedited: Final = await _input_tokens(router, unedited_body) + polyfilled: Final = await _input_tokens( + router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(perplexity_model)} + ) + + assert edited < unedited + assert edited == polyfilled + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_responses_api_bridge_counts_the_unedited_body_when_an_edit_fails() -> None: + from litellm.llms.anthropic.messages.responses_adapters import handler as responses_adapters_handler + + router: Final = Router( + model_list=[ + { + "model_name": "sonar-agent", + "litellm_params": {"model": "perplexity/perplexity/sonar", "api_key": "test-key"}, + } + ] + ) + unedited_body: Final[dict[str, JsonValue]] = { + "model": "sonar-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + + with patch.object( + responses_adapters_handler, "apply_context_management", AsyncMock(side_effect=RuntimeError("editor bug")) + ): + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + + assert edited == await _input_tokens(router, unedited_body) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_responses_api_bridge_rejects_an_invalid_edit_without_drop_params() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "sonar-agent", + "litellm_params": {"model": "perplexity/perplexity/sonar", "api_key": "test-key"}, + } + ] + ) + invalid_clear_edit: Final[dict[str, JsonValue]] = { + "edits": [{"type": "clear_tool_uses_20250919", "clear_at_least": "all of them"}] + } + + with patch.object(litellm, "drop_params", False): + response: Final = await _messages_endpoint_count( + router, + { + "model": "sonar-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": invalid_clear_edit, + }, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + message: Final = json.loads(bytes(response.body))["error"]["message"] + assert "context_management.edits.0.clear_tool_uses_20250919.clear_at_least" in message + + +class _CountTokensPeer: + def __init__(self, reply: Mapping[str, JsonValue]) -> None: + self._reply: Final = reply + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + return httpx.Response(200, json=dict(self._reply), request=request) + + +class _BedrockInvokeModel(BaseModel): + model_config = ConfigDict(frozen=True) + body: str + + +class _BedrockCountTokensInput(BaseModel): + model_config = ConfigDict(frozen=True) + invoke_model: _BedrockInvokeModel = Field(alias="invokeModel") + + +class _BedrockCountTokensBody(BaseModel): + model_config = ConfigDict(frozen=True) + input: _BedrockCountTokensInput + + +def _bedrock_counted_request(request_body: bytes) -> str: + invoke_model: Final = _BedrockCountTokensBody.model_validate_json(request_body).input.invoke_model + return base64.b64decode(invoke_model.body).decode() + + +async def _count_on_a_bridged_provider_count_api( + request_body: Mapping[str, JsonValue], + monkeypatch: pytest.MonkeyPatch, + bedrock_model: str = "bedrock/us.amazon.nova-2-lite-v1:0", + peer_reply: Mapping[str, JsonValue] = MappingProxyType({"inputTokens": 57}), + model_info: Mapping[str, JsonValue] = MappingProxyType({}), +) -> tuple[object, tuple[httpx.Request, ...]]: + from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_handler + from litellm.llms.bedrock.count_tokens import mantle_handler as bedrock_mantle_count_tokens_handler + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + peer: Final = _CountTokensPeer(peer_reply) + router: Final = Router( + model_list=[ + { + "model_name": "nova-agent", + "litellm_params": { + "model": bedrock_model, + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + }, + "model_info": dict(model_info), + } + ] + ) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + with ( + patch.object(bedrock_count_tokens_handler, "get_async_httpx_client", return_value=client), + patch.object(bedrock_mantle_count_tokens_handler, "get_async_httpx_client", return_value=client), + ): + response: Final = await _messages_endpoint_count(router, {**request_body, "model": "nova-agent"}) + return response, peer.requests + + +@pytest.mark.asyncio +async def test_messages_count_tokens_sends_the_edited_messages_to_a_bridged_provider_count_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent_bodies = await _count_on_a_bridged_provider_count_api( + { + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + }, + monkeypatch, + ) + + assert response == {"input_tokens": 57, "litellm_count_source": "provider"} + assert len(sent_bodies) == 1 + assert _bedrock_counted_request(sent_bodies[0].content).count(_GEMINI_TOOL_OUTPUT) == 1 + + +@pytest.mark.asyncio +async def test_messages_count_tokens_sends_a_body_with_only_empty_text_to_the_bridged_provider_count_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + empty_turn: Final[dict[str, JsonValue]] = {"role": "user", "content": [{"type": "text", "text": ""}]} + + response, sent_bodies = await _count_on_a_bridged_provider_count_api({"messages": [empty_turn]}, monkeypatch) + + assert response == {"input_tokens": 57, "litellm_count_source": "provider"} + assert len(sent_bodies) == 1 + + +_BEDROCK_CLAUDE_MODEL: Final = "bedrock/us.anthropic.claude-sonnet-5-5" +_CLEAR_THINKING_EDIT: Final[dict[str, JsonValue]] = {"type": "clear_thinking_20251015"} +_COUNTED_ON_MANTLE: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"token_count_model": {"model": "us.anthropic.claude-sonnet-5-5", "litellm_provider": "bedrock_mantle"}} +) +_COUNTED_ON_RUNTIME: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"token_count_model": {"model": "us.anthropic.claude-sonnet-5-5", "litellm_provider": "bedrock"}} +) + + +def _bedrock_claude_count_request(edits: list[JsonValue]) -> Mapping[str, JsonValue]: + return MappingProxyType( + { + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": {"edits": edits}, + } + ) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_beta_to_the_mantle_count_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent = await _count_on_a_bridged_provider_count_api( + _bedrock_claude_count_request([*_CLEAR_TOOL_USES_EDITS, _CLEAR_THINKING_EDIT]), + monkeypatch, + _BEDROCK_CLAUDE_MODEL, + MappingProxyType({"input_tokens": 57, "context_management": {"original_input_tokens": 4100}}), + _COUNTED_ON_MANTLE, + ) + + assert isinstance(response, dict) + assert (response["input_tokens"], response["context_management"]) == (57, {"original_input_tokens": 4100}) + assert len(sent) == 1 + assert json.loads(sent[0].content)["context_management"] == _CLEAR_TOOL_USES + assert sent[0].headers["anthropic-beta"] == ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value + + +@pytest.mark.asyncio +async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_beta_to_the_runtime_count_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent = await _count_on_a_bridged_provider_count_api( + _bedrock_claude_count_request([*_CLEAR_TOOL_USES_EDITS, _CLEAR_THINKING_EDIT]), + monkeypatch, + _BEDROCK_CLAUDE_MODEL, + model_info=_COUNTED_ON_RUNTIME, + ) + + assert isinstance(response, dict) + assert response["input_tokens"] == 57 + assert len(sent) == 1 + invoke_body: Final = json.loads(_bedrock_counted_request(sent[0].content)) + assert (invoke_body["context_management"], invoke_body["anthropic_beta"]) == ( + _CLEAR_TOOL_USES, + [ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value], + ) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_sends_no_context_management_to_bedrock_when_no_edit_is_bedrock_supported( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent = await _count_on_a_bridged_provider_count_api( + _bedrock_claude_count_request([_CLEAR_THINKING_EDIT]), + monkeypatch, + _BEDROCK_CLAUDE_MODEL, + MappingProxyType({"input_tokens": 57}), + _COUNTED_ON_MANTLE, + ) + + assert isinstance(response, dict) + assert response["input_tokens"] == 57 + assert len(sent) == 1 + assert "context_management" not in json.loads(sent[0].content) + assert "anthropic-beta" not in sent[0].headers + + +@pytest.mark.asyncio +async def test_messages_count_tokens_rejects_a_non_string_tool_result_text_before_a_bridged_provider_applies_edits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent_bodies = await _count_on_a_bridged_provider_count_api( + {"messages": _build_log_messages(5), "context_management": _CLEAR_TOOL_USES}, monkeypatch + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert "tool_result toolu_log content.1.text" in json.loads(bytes(response.body))["error"]["message"] + assert sent_bodies == () + + +@pytest.mark.asyncio +async def test_messages_count_tokens_applies_a_responses_api_deployment_edits_once() -> None: + peer: Final = _CountTokensPeer({"object": "response.input_tokens", "input_tokens": 71}) + router: Final = Router( + model_list=[{"model_name": "gpt-agent", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "test-key"}}] + ) + + with patch.object( + openai_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ): + response: Final = await _messages_endpoint_count( + router, + { + "model": "gpt-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + }, + ) + + assert response == {"input_tokens": 71, "litellm_count_source": "provider"} + assert len(peer.requests) == 1 + assert peer.requests[0].content.decode().count(_GEMINI_TOOL_OUTPUT) == 1 + + +_DATABRICKS_CLAUDE_MODEL: Final = "databricks/databricks-claude-sonnet-4-6" + + +def _databricks_claude_router() -> Router: + return Router( + model_list=[ + { + "model_name": "databricks-claude-agent", + "litellm_params": { + "model": _DATABRICKS_CLAUDE_MODEL, + "api_key": "test-key", + "api_base": "https://example.cloud.databricks.com/serving-endpoints", + }, + } + ] + ) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_counts_the_edited_body_on_a_native_messages_deployment_without_a_count_api() -> ( + None +): + router: Final = _databricks_claude_router() + unedited_body: Final[dict[str, JsonValue]] = { + "model": "databricks-claude-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + unedited: Final = await _input_tokens(router, unedited_body) + polyfilled: Final = await _input_tokens( + router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_DATABRICKS_CLAUDE_MODEL)} + ) + + assert edited < unedited + assert edited == polyfilled + + +@pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", [False, True]) +async def test_messages_count_tokens_on_a_native_messages_deployment_without_a_count_api_rejects_an_invalid_edit( + drop_params: bool, +) -> None: + with patch.object(litellm, "drop_params", drop_params): + response: Final = await _messages_endpoint_count( + _databricks_claude_router(), + { + "model": "databricks-claude-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _INVALID_COMPACT_EDIT, + }, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_messages_count_tokens_leaves_the_edits_to_a_native_messages_deployment() -> None: + peer: Final = _CountTokensPeer({"input_tokens": 63}) + router: Final = Router( + model_list=[ + { + "model_name": "claude-agent", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "test-key"}, + } + ] + ) + + with patch.object( + anthropic_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ): + response: Final = await _messages_endpoint_count( + router, + { + "model": "claude-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + }, + ) + + assert response == {"input_tokens": 63, "litellm_count_source": "provider"} + assert len(peer.requests) == 1 + assert peer.requests[0].content.decode().count(_GEMINI_TOOL_OUTPUT) == 3 + + @pytest.mark.asyncio @pytest.mark.parametrize( "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index fd98bdcb418..db0db97cd88 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -818,6 +818,165 @@ class TestPostCallFailureHookLiftsCallTypeAndStartTime: assert payload["model"] == "MCP: deepwiki-ask_question" assert payload["endTime"] - payload["startTime"] >= real_datetime.timedelta(seconds=2) + @pytest.mark.asyncio + async def test_failed_speech_spend_row_records_the_endpoint_call_type_over_the_body( + self, monkeypatch: pytest.MonkeyPatch + ): + from types import SimpleNamespace + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + rejected: Final = litellm.BadRequestError( + message="Voice name nonexistent is not supported", model="gemini-tts", llm_provider="gemini" + ) + request_data: Final = { + "model": "gemini-tts", + "input": "Hello there", + "voice": "nonexistent", + "call_type": "acompletion", + } + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + spend_writer: Final = SimpleNamespace(update_database=AsyncMock()) + monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger(spend_writer=lambda: spend_writer)]) + try: + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=rejected, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route="/v1/audio/speech"), + dispatched_call_type="aspeech", + ) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + db_call: Final = spend_writer.update_database.call_args.kwargs + payload: Final = get_logging_payload( + kwargs=db_call["kwargs"], + response_obj=db_call["completion_response"], + start_time=db_call["start_time"], + end_time=db_call["end_time"], + ) + assert payload["call_type"] == "aspeech" + + @pytest.mark.asyncio + async def test_failed_mcp_list_tools_row_inside_a_responses_request_keeps_list_mcp_tools( + self, monkeypatch: pytest.MonkeyPatch + ): + from types import SimpleNamespace + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + request_data: Final = { + "model": "MCP: list_tools", + "call_type": "list_mcp_tools", + "input": [{"role": "system", "content": "list_tools"}], + } + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + spend_writer: Final = SimpleNamespace(update_database=AsyncMock()) + monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger(spend_writer=lambda: spend_writer)]) + try: + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=RuntimeError("MCP server unreachable"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route="/v1/responses"), + route="/mcp/list_tools", + ) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + db_call: Final = spend_writer.update_database.call_args.kwargs + payload: Final = get_logging_payload( + kwargs=db_call["kwargs"], + response_obj=db_call["completion_response"], + start_time=db_call["start_time"], + end_time=db_call["end_time"], + ) + assert payload["call_type"] == "list_mcp_tools" + + @pytest.mark.asyncio + async def test_failure_without_an_endpoint_call_type_does_not_guess_one_from_a_multi_method_path( + self, monkeypatch: pytest.MonkeyPatch + ): + from types import SimpleNamespace + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + spend_writer: Final = SimpleNamespace(update_database=AsyncMock()) + monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger(spend_writer=lambda: spend_writer)]) + try: + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data={}, + original_exception=RuntimeError("vector store lookup failed"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route="/v1/vector_stores"), + ) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + db_call: Final = spend_writer.update_database.call_args.kwargs + payload: Final = get_logging_payload( + kwargs=db_call["kwargs"], + response_obj=db_call["completion_response"], + start_time=db_call["start_time"], + end_time=db_call["end_time"], + ) + assert payload["call_type"] == "" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("request_data", "route", "endpoint_call_type"), + [ + ({"model": "gemini-tts", "input": "Hello there", "voice": "Kore"}, "/v1/audio/speech", "aspeech"), + ({"model": "omni-moderation", "input": ["one", "two"]}, "/v1/moderations", "amoderation"), + ({"model": "gpt-transcribe"}, "/v1/audio/transcriptions", "atranscription"), + ], + ) + async def test_request_rejected_by_the_proxy_records_the_endpoint_call_type( + self, + monkeypatch: pytest.MonkeyPatch, + request_data: dict[str, object], + route: str, + endpoint_call_type: str, + ): + from types import SimpleNamespace + + from fastapi import HTTPException + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + spend_writer: Final = SimpleNamespace(update_database=AsyncMock()) + monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger(spend_writer=lambda: spend_writer)]) + try: + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data=dict(request_data), + original_exception=HTTPException(status_code=429, detail="Rate limit exceeded"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route=route), + dispatched_call_type=endpoint_call_type, + ) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + db_call: Final = spend_writer.update_database.call_args.kwargs + payload: Final = get_logging_payload( + kwargs=db_call["kwargs"], + response_obj=db_call["completion_response"], + start_time=db_call["start_time"], + end_time=db_call["end_time"], + ) + assert payload["call_type"] == endpoint_call_type + class TestPostCallFailureHookEstimatesDispatchedInputTokens: """A non-stream request that failed after dispatch (timeout, provider @@ -3576,6 +3735,44 @@ def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) +@pytest.mark.asyncio +async def test_a_billed_responses_route_failure_captures_usage_from_the_newest_compaction_item() -> None: + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation + from litellm.proxy.db.accounting_ledger import UsageWrite + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings + from litellm.proxy.utils import _failure_fields_to_lift + + async def persist(_: tuple[UsageWrite, ...]) -> None: + return None + + def captured_prompt_tokens(request_input: list[dict[str, str]]) -> int: + logging_obj: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params()) + logging_obj.model_call_details["call_type"] = "aresponses" + logging_obj.model_call_details["messages"] = request_input + lifted: Final = _failure_fields_to_lift( + {"litellm_logging_obj": logging_obj}, provider_billed=True, route_call_type="aresponses" + ) + snapshot: Final = lifted["_accounting_usage_snapshot"] + assert isinstance(snapshot, UsageCaptureObservation) and snapshot.usage_json is not None + return json.loads(snapshot.usage_json)["prompt_tokens"] + + runtime: Final = UsageCaptureRuntime() + runtime.start(persist, UsageCaptureSettings()) + try: + sent: Final = captured_prompt_tokens( + [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + since_compaction: Final = captured_prompt_tokens( + [COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + finally: + await runtime.stop(1) + + assert sent == since_compaction + + @pytest.mark.parametrize( "current_attempt_api_call_start_time", [None, 1700000001.0], diff --git a/tests/unit/proxy/test_response_model_sanitization.py b/tests/unit/proxy/test_response_model_sanitization.py index f52736453c1..b3ad53221eb 100644 --- a/tests/unit/proxy/test_response_model_sanitization.py +++ b/tests/unit/proxy/test_response_model_sanitization.py @@ -98,42 +98,13 @@ def test_restamp_streaming_chunk_skips_matching_model(): assert model_mismatch_logged is False -def test_fast_serialize_simple_streaming_chunk_matches_model_dump_json(): +def test_serialize_streaming_chunk_omits_a_missing_model_instead_of_null(): from litellm.proxy.proxy_server import _serialize_streaming_chunk - chunk = _make_model_response_stream_chunk("client-model") - - assert json.loads(_serialize_streaming_chunk(chunk)) == json.loads( - chunk.model_dump_json(exclude_none=True, exclude_unset=True) - ) - - -def test_fast_serialize_returns_none_when_model_field_is_missing(): - """ - The fast path must mirror ``model_dump_json(exclude_none=True)``: when - ``chunk.model`` is ``None`` the slow path omits the field entirely. - Emitting ``"model": null`` would diverge and trip strict OpenAI- - compatible clients that reject ``null`` for optional string fields. - Falling back to ``None`` lets the canonical serializer handle the edge. - """ - from litellm.proxy.proxy_server import ( - _fast_serialize_simple_model_response_stream, - _serialize_streaming_chunk, - ) - chunk = _make_model_response_stream_chunk("client-model") chunk.model = None # type: ignore[assignment] - assert _fast_serialize_simple_model_response_stream(chunk) is None - - # Going through the public ``_serialize_streaming_chunk`` should still - # produce a serialized result via the slow-path fallback, and it must - # not contain ``"model": null``. - serialized = _serialize_streaming_chunk(chunk) - payload_str = serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized - assert '"model": null' not in payload_str - assert '"model":null' not in payload_str - assert json.loads(payload_str) == json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True)) + assert "model" not in json.loads(_serialize_streaming_chunk(chunk)) def test_proxy_chat_completion_does_not_return_provider_prefixed_model(tmp_path, monkeypatch): diff --git a/tests/unit/proxy/test_zero_cost_model_budget_bypass.py b/tests/unit/proxy/test_zero_cost_model_budget_bypass.py index c2931470482..ce1dfce1646 100644 --- a/tests/unit/proxy/test_zero_cost_model_budget_bypass.py +++ b/tests/unit/proxy/test_zero_cost_model_budget_bypass.py @@ -32,6 +32,29 @@ from litellm.router import Router from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +@pytest.mark.parametrize( + ("model", "expected_free"), + [("parallel_ai/decision", False), ("fireworks_ai/accounts/fireworks/routers/rune", True)], +) +def test_query_price_keeps_budgets_only_where_it_is_billed(model: str, expected_free: bool) -> None: + router = Router( + model_list=[ + { + "model_name": "query-priced", + "litellm_params": { + "model": model, + "api_key": "test-key", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "input_cost_per_query": 0.025, + }, + } + ] + ) + + assert is_model_cost_zero("query-priced", router) is expected_free + + @pytest.fixture def mock_router_with_zero_cost_model(): """Create a mock router with a zero-cost model.""" diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index ea39d5d59e2..e66050877df 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -689,16 +689,11 @@ class TestProxySettingEndpoints: # Verify settings were updated settings = data["settings"] assert settings["google_client_id"] == new_sso_settings["google_client_id"] - assert ( - settings["google_client_secret"] == new_sso_settings["google_client_secret"] - ) + assert settings["google_client_secret"] == "new_****************cret" assert ( settings["microsoft_client_id"] == new_sso_settings["microsoft_client_id"] ) - assert ( - settings["microsoft_client_secret"] - == new_sso_settings["microsoft_client_secret"] - ) + assert settings["microsoft_client_secret"] == "new_*******************cret" assert settings["proxy_base_url"] == new_sso_settings["proxy_base_url"] assert settings["user_email"] == new_sso_settings["user_email"] @@ -722,6 +717,163 @@ class TestProxySettingEndpoints: create_sso_settings = json.loads(create_data["sso_settings"]) assert create_sso_settings["google_client_id"] == "new_google_client_id" + def _patch_sso_settings(self, monkeypatch, stored, body, environment_variables=None): + """PATCH /update/sso_settings against a stored row; returns (response, saved row).""" + import json + from unittest.mock import AsyncMock, MagicMock + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + mock_prisma = MagicMock() + stored_record = None + if stored is not None: + stored_record = MagicMock() + stored_record.sso_settings = stored + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock( + return_value=stored_record + ) + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + env_var_entry = None + if environment_variables is not None: + env_var_entry = MagicMock() + env_var_entry.param_value = environment_variables + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=env_var_entry) + mock_prisma.db.litellm_config.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr( + proxy_config, "encrypt_env_variables", lambda environment_variables: environment_variables + ) + monkeypatch.setattr(proxy_config, "decrypt_db_variables", lambda variables: variables) + + response = client.patch("/update/sso_settings", json=body) + upsert_kwargs = mock_prisma.db.litellm_ssoconfig.upsert.call_args.kwargs + return response, json.loads(upsert_kwargs["data"]["update"]["sso_settings"]) + + _STORED_GENERIC_SSO = { + "generic_client_id": "stored-client-id", + "generic_client_secret": "stored-client-secret", + "generic_authorization_endpoint": "https://idp.example.com/authorize", + "user_email": "admin@example.com", + "role_mappings": { + "provider": "generic", + "group_claim": "groups", + "default_role": "internal_user", + "roles": {"proxy_admin": ["admins"]}, + }, + } + + def test_update_sso_settings_omitted_fields_keep_stored_values( + self, mock_proxy_config, mock_auth, monkeypatch + ): + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "stored-client-secret") + new_role_mappings = { + "provider": "generic", + "group_claim": "roles", + "default_role": "internal_user", + "roles": {"proxy_admin": ["platform-admins"]}, + } + + response, saved = self._patch_sso_settings( + monkeypatch, self._STORED_GENERIC_SSO, {"role_mappings": new_role_mappings} + ) + + assert response.status_code == 200 + assert saved["generic_client_secret"] == "stored-client-secret" + assert saved["generic_client_id"] == "stored-client-id" + assert saved["user_email"] == "admin@example.com" + assert saved["role_mappings"]["group_claim"] == "roles" + assert os.environ["GENERIC_CLIENT_SECRET"] == "stored-client-secret" + + def test_update_sso_settings_null_clears_only_that_field( + self, mock_proxy_config, mock_auth, monkeypatch + ): + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "stored-client-secret") + monkeypatch.setenv("GENERIC_CLIENT_ID", "stored-client-id") + + response, saved = self._patch_sso_settings( + monkeypatch, + self._STORED_GENERIC_SSO, + {"generic_client_secret": None, "role_mappings": None}, + ) + + assert response.status_code == 200 + assert saved["generic_client_secret"] is None + assert saved["role_mappings"] is None + assert saved["generic_client_id"] == "stored-client-id" + assert "GENERIC_CLIENT_SECRET" not in os.environ + assert os.environ["GENERIC_CLIENT_ID"] == "stored-client-id" + + def test_update_sso_settings_new_value_replaces_stored_value( + self, mock_proxy_config, mock_auth, monkeypatch + ): + response, saved = self._patch_sso_settings( + monkeypatch, self._STORED_GENERIC_SSO, {"generic_client_secret": "rotated-secret"} + ) + + assert response.status_code == 200 + assert saved["generic_client_secret"] == "rotated-secret" + assert saved["generic_client_id"] == "stored-client-id" + assert os.environ["GENERIC_CLIENT_SECRET"] == "rotated-secret" + + def test_update_sso_settings_omitted_field_leaves_env_only_value_alone( + self, mock_proxy_config, mock_auth, monkeypatch + ): + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "env-only-secret") + + response, saved = self._patch_sso_settings( + monkeypatch, None, {"user_email": "admin@example.com"} + ) + + assert response.status_code == 200 + assert saved["user_email"] == "admin@example.com" + assert saved["generic_client_secret"] is None + assert os.environ["GENERIC_CLIENT_SECRET"] == "env-only-secret" + + def test_update_sso_settings_response_masks_stored_secrets( + self, mock_proxy_config, mock_auth, monkeypatch + ): + response, saved = self._patch_sso_settings( + monkeypatch, self._STORED_GENERIC_SSO, {} + ) + + assert response.status_code == 200 + assert saved["generic_client_secret"] == "stored-client-secret" + returned = response.json()["settings"] + assert returned["generic_client_id"] == "stored-client-id" + assert returned["generic_client_secret"] != "stored-client-secret" + assert "stored-client-secret" not in response.text + + def test_update_sso_settings_moves_omitted_env_row_values_into_sso_row( + self, mock_proxy_config, mock_auth, monkeypatch + ): + import json + + environment_variables = { + "MICROSOFT_TENANT": "env-row-tenant", + "MICROSOFT_CLIENT_ID": "env-row-client-id", + "GENERIC_CLIENT_ID": "env-row-generic-id", + "LANGFUSE_PUBLIC_KEY": "unrelated", + } + + response, saved = self._patch_sso_settings( + monkeypatch, + self._STORED_GENERIC_SSO, + {"user_email": "new@example.com", "microsoft_client_id": None}, + environment_variables=environment_variables, + ) + + from litellm.proxy.proxy_server import prisma_client + + assert response.status_code == 200 + assert saved["microsoft_tenant"] == "env-row-tenant" + assert saved["microsoft_client_id"] is None + assert saved["generic_client_id"] == "stored-client-id" + assert saved["user_email"] == "new@example.com" + env_row_update = prisma_client.db.litellm_config.update.call_args.kwargs + assert json.loads(env_row_update["data"]["param_value"]) == {"LANGFUSE_PUBLIC_KEY": "unrelated"} + def test_update_sso_settings_maps_saml_fields_to_env_vars( self, mock_proxy_config, mock_auth, monkeypatch ): @@ -804,7 +956,7 @@ class TestProxySettingEndpoints: mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock( + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( side_effect=ValueError("cleanup failed") ) mock_prisma.db.litellm_config.update = AsyncMock() diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index 46777b32319..ad3036560a3 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -1,28 +1,40 @@ import asyncio +import functools import json import time from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime, timezone from pathlib import Path from types import TracebackType -from typing import Final, Protocol +from typing import ClassVar, Final, Literal, Protocol from unittest.mock import AsyncMock, MagicMock, patch import httpx import openai import pytest import respx -from websockets.exceptions import ConnectionClosedError +from openai import AsyncOpenAI +from websockets.exceptions import ConnectionClosed, ConnectionClosedError from websockets.frames import Close import litellm +from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.realtime_streaming import TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE, RealTimeStreaming +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.openai.workload_identity import _workload_identity_auth from litellm.models.credentials import CredentialItem +from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.litellm_pre_call_utils import _apply_resolved_guardrails_to_metadata from litellm.realtime_api import main as realtime_main -from litellm.realtime_api.main import _with_resolved_session_model +from litellm.realtime_api.main import with_resolved_session_model from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyMatchContext +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep +from litellm.types.realtime import RealtimeQueryParams @pytest.fixture @@ -48,19 +60,19 @@ class _CallCapture: def test_resolves_top_level_session_model(): - resolved = _with_resolved_session_model({"model": "alias/gpt-realtime"}, "gpt-realtime") + resolved = with_resolved_session_model({"model": "alias/gpt-realtime"}, "gpt-realtime") assert resolved == {"model": "gpt-realtime"} def test_session_without_model_is_returned_unchanged(): session = {"type": "realtime", "audio": {"input": {}}} - assert _with_resolved_session_model(session, "gpt-realtime") == session + assert with_resolved_session_model(session, "gpt-realtime") == session def test_does_not_clobber_flat_transcription_model(): """The nested transcription model is a different model than the realtime conversation model and must not be overwritten with the routing model.""" - resolved = _with_resolved_session_model( + resolved = with_resolved_session_model( {"model": "gpt-4o-realtime-preview", "input_audio_transcription": {"model": "whisper-1"}}, "gpt-4o-realtime-preview", ) @@ -68,7 +80,7 @@ def test_does_not_clobber_flat_transcription_model(): def test_does_not_clobber_nested_audio_transcription_model(): - resolved = _with_resolved_session_model( + resolved = with_resolved_session_model( { "model": "gpt-4o-realtime-preview", "audio": {"input": {"transcription": {"model": "whisper-1"}}}, @@ -80,7 +92,7 @@ def test_does_not_clobber_nested_audio_transcription_model(): def test_original_session_is_not_mutated(): session = {"model": "alias/gpt-realtime"} - _with_resolved_session_model(session, "gpt-realtime") + with_resolved_session_model(session, "gpt-realtime") assert session == {"model": "alias/gpt-realtime"} @@ -388,6 +400,52 @@ async def test_arealtime_live_session_rejects_streaming_only_pre_call_guardrail( assert live_handler.called is False +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_metadata", + ( + {"user_api_key_metadata": {"disable_global_guardrails": True}}, + {"user_api_key_team_metadata": {"opted_out_global_guardrails": ["default-on-blocker"]}}, + ), +) +async def test_arealtime_key_metadata_cannot_opt_out_of_a_default_on_realtime_guardrail( + monkeypatch: pytest.MonkeyPatch, key_metadata: dict[str, object] +) -> None: + class FakeLiveHandler: + def __init__(self) -> None: + self.called = False + + async def async_live_session(self, **kwargs: object) -> None: + self.called = True + + live_handler: Final = FakeLiveHandler() + monkeypatch.setattr(realtime_main, "openai_live_sessions", live_handler) + monkeypatch.setattr( + litellm, + "callbacks", + [ + CustomGuardrail( + guardrail_name="default-on-blocker", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + ], + ) + + with pytest.raises(ValueError, match="Guardrails are not supported on OpenAI Live sessions"): + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-live-1", + websocket=MagicMock(), + api_key="fake-key", + api_base="https://live-gateway.example/v1", + litellm_logging_obj=FakeLogging(), + metadata=key_metadata, + live_session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + ) + + assert live_handler.called is False + + @pytest.mark.asyncio async def test_arealtime_live_session_rejects_non_openai_provider(monkeypatch: pytest.MonkeyPatch) -> None: def mock_get_llm_provider( @@ -986,6 +1044,114 @@ async def test_realtime_http_routes_send_a_bare_model_name_to_its_explicit_provi assert gateway.url == f"https://realtime-gateway.example/v1/realtime/{path}" +class _GatewayThatRecordsTheBody: + body: object = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.body = json.loads(request.content) + return httpx.Response(200, json={"client_secret": {"value": "ek_test"}}) + + +@pytest.mark.asyncio +async def test_transcription_session_sends_the_deployment_model_upstream_in_every_transcription_shape() -> None: + gateway: Final = _GatewayThatRecordsTheBody() + await realtime_main.acreate_realtime_transcription_session.__wrapped__( + model="openai/gpt-transcribe", + transcription_session={ + "audio": {"input": {"transcription": {"model": "team-transcriber", "language": "en"}}}, + "input_audio_transcription": {"model": "team-transcriber"}, + "input_audio_format": "pcm16", + }, + api_base="https://realtime-gateway.example/v1", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + ) + assert gateway.body == { + "audio": {"input": {"transcription": {"model": "gpt-transcribe", "language": "en"}}}, + "input_audio_transcription": {"model": "gpt-transcribe"}, + "input_audio_format": "pcm16", + } + + +def _create_realtime_session_with_an_explicit_null(route: str, client: AsyncHTTPHandler) -> Awaitable[httpx.Response]: + if route == "client_secrets": + return realtime_main.acreate_realtime_client_secret.__wrapped__( + session={ + "type": "realtime", + "model": "gpt-realtime", + "instructions": None, + "tools": None, + "audio": {"input": {"turn_detection": None}}, + }, + custom_llm_provider="openai", + api_base="https://realtime-gateway.example/v1", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + client=client, + ) + if route == "translation_client_secrets": + return realtime_main.acreate_realtime_translation_client_secret.__wrapped__( + model="gpt-realtime-translate", + session={"include": None, "audio": {"input": {"noise_reduction": None}}}, + custom_llm_provider="openai", + api_base="https://realtime-gateway.example/v1", + api_key="fake-key", + litellm_logging_obj=_realtime_http_logging(route), + client=client, + ) + return realtime_main.acreate_realtime_transcription_session.__wrapped__( + model="openai/gpt-transcribe", + transcription_session={"include": None, "audio": {"input": {"turn_detection": None}}}, + api_base="https://realtime-gateway.example/v1", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + client=client, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "body_path"), + ( + ("client_secrets", ("session", "audio", "input", "turn_detection")), + ("translation_client_secrets", ("session", "audio", "input", "noise_reduction")), + ("transcription_sessions", ("audio", "input", "turn_detection")), + ), +) +async def test_realtime_session_requests_forward_an_explicit_null_upstream( + route: str, body_path: tuple[str, ...] +) -> None: + gateway: Final = _GatewayThatRecordsTheBody() + + await _create_realtime_session_with_an_explicit_null( + route, AsyncHTTPHandler(transport=httpx.MockTransport(gateway)) + ) + + parent: Final = functools.reduce(lambda node, key: node[key], body_path[:-1], gateway.body) + assert body_path[-1] in parent + assert parent[body_path[-1]] is None + session_fields: Final = gateway.body["session"] if body_path[0] == "session" else gateway.body + assert [key for key, value in session_fields.items() if value is None] == [] + + +@pytest.mark.asyncio +async def test_transcription_client_secret_sends_the_deployment_model_as_the_session_model() -> None: + gateway: Final = _GatewayThatRecordsTheSession() + + await realtime_main.acreate_realtime_client_secret.__wrapped__( + session={"type": "transcription", "model": "openai/gpt-live-transcribe"}, + api_key="sk-test", + api_base="https://upstream.example", + litellm_logging_obj=_realtime_http_logging("client_secrets"), + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + ) + + assert gateway.session is not None + assert gateway.session["model"] == "gpt-live-transcribe" + assert gateway.session["audio"] == {"input": {"transcription": {"model": "gpt-live-transcribe"}}} + + class _ClientWebSocketWithHeaders: def __init__(self, headers: tuple[tuple[bytes, bytes], ...]) -> None: self.scope: Final = {"headers": headers} @@ -995,11 +1161,51 @@ _GA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=()) _BETA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=((b"openai-beta", b"realtime=v1"),)) +_GA_DEFAULT_SESSIONS: Final = ( + pytest.param("gpt-realtime-2.1", "realtime", None, id="ga-only-model"), + pytest.param("gpt-realtime-2", "realtime", None, id="ga-only-alias"), + pytest.param("gpt-realtime-2-2026-05-06", "realtime", None, id="ga-only-snapshot"), + # Azure answers beta sessions on these with "The Realtime preview API is no longer available" + # (https://aka.ms/aoai/realtime-ga), checked live 2026-10-28 + pytest.param("gpt-realtime", "realtime", None, id="beta-retired-gpt-realtime"), + pytest.param("gpt-realtime-mini", "realtime", None, id="beta-retired-gpt-realtime-mini"), + pytest.param("gpt-realtime-1.5", "realtime", None, id="beta-retired-gpt-realtime-1.5"), + pytest.param("my-transcribe-deployment", "realtime", {"intent": "transcription"}, id="transcription-intent"), + pytest.param("my-translate-deployment", "translation", None, id="translation-mode"), +) + + +@pytest.mark.parametrize(("model", "realtime_mode", "query_params"), _GA_DEFAULT_SESSIONS) +def test_azure_realtime_protocol_defaults_to_ga_when_none_is_configured( + local_model_cost_map, model: str, realtime_mode: str, query_params: RealtimeQueryParams | None +) -> None: + from litellm.llms.azure.realtime.handler import azure_realtime_protocol_for_client + + protocol: Final = azure_realtime_protocol_for_client( + None, model=model, realtime_mode=realtime_mode, query_params=query_params, websocket=_BETA_CLIENT + ) + + assert protocol.name == "GA" + + +@pytest.mark.parametrize(("model", "realtime_mode", "query_params"), _GA_DEFAULT_SESSIONS) +def test_azure_realtime_protocol_configured_by_the_operator_wins( + local_model_cost_map, model: str, realtime_mode: str, query_params: RealtimeQueryParams | None +) -> None: + from litellm.llms.azure.realtime.handler import azure_realtime_protocol_for_client + + protocol: Final = azure_realtime_protocol_for_client( + "beta", model=model, realtime_mode=realtime_mode, query_params=query_params, websocket=_GA_CLIENT + ) + + assert protocol.name == "beta" + + async def _azure_backend_url_dialed_for(websocket: _ClientWebSocketWithHeaders, **kwargs: object) -> str | None: connect: Final = _ConnectThatStopsAfterCapturingTheUrl() with patch("websockets.connect", connect): await realtime_main._arealtime.__wrapped__( - model="azure/gpt-realtime", + model="azure/my-realtime-deployment", websocket=websocket, api_base="https://my-endpoint.openai.azure.com", api_key="fake-key", @@ -1014,7 +1220,7 @@ async def test_arealtime_azure_ga_client_without_beta_header_dials_the_ga_upstre monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) assert ( await _azure_backend_url_dialed_for(_GA_CLIENT) - == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-realtime" + == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=my-realtime-deployment" ) @@ -1022,7 +1228,7 @@ async def test_arealtime_azure_ga_client_without_beta_header_dials_the_ga_upstre async def test_arealtime_azure_beta_header_client_keeps_the_beta_upstream(monkeypatch): monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) assert await _azure_backend_url_dialed_for(_BETA_CLIENT) == ( - "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=my-realtime-deployment" ) @@ -1030,7 +1236,7 @@ async def test_arealtime_azure_beta_header_client_keeps_the_beta_upstream(monkey async def test_arealtime_azure_explicit_beta_protocol_wins_over_a_ga_client(monkeypatch): monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) assert await _azure_backend_url_dialed_for(_GA_CLIENT, realtime_protocol="beta") == ( - "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=my-realtime-deployment" ) @@ -1038,7 +1244,197 @@ async def test_arealtime_azure_explicit_beta_protocol_wins_over_a_ga_client(monk async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypatch): monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta") assert await _azure_backend_url_dialed_for(_GA_CLIENT) == ( - "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=my-realtime-deployment" + ) + + +class _BetaClientThatStaysConnected: + def __init__(self) -> None: + self.scope: Final = {"headers": ((b"openai-beta", b"realtime=v1"),)} + self.sent: tuple[str, ...] = () + + async def receive_text(self) -> str: + return await asyncio.get_running_loop().create_future() + + async def send_text(self, data: str) -> None: + self.sent = (*self.sent, data) + + async def close(self, code: int = 1000, reason: str = "") -> None: + return None + + +class _ConnectToABackendThatSendsOneFrame: + def __init__(self, frame: Mapping[str, object]) -> None: + self.url: str | None = None + self._backend: Final = MagicMock(recv=AsyncMock(side_effect=[json.dumps(frame), ConnectionClosed(None, None)])) + + def __call__( + self, url: str, *, additional_headers: Mapping[str, str], max_size: int, ssl: object + ) -> "_ConnectToABackendThatSendsOneFrame": + self.url = url + return self + + async def __aenter__(self) -> MagicMock: + return self._backend + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + return None + + +def _realtime_websocket_logging(*failure_callbacks: CustomLogger) -> Logging: + return Logging( + model="gpt-realtime", + messages=[], + stream=False, + call_type="_arealtime", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="test-realtime-websocket", + function_id="test-realtime-websocket", + dynamic_async_failure_callbacks=list(failure_callbacks), + ) + + +_GA_SESSION_CREATED: Final = { + "type": "session.created", + "event_id": "event_1", + "session": { + "type": "realtime", + "id": "sess_1", + "model": "my-realtime-deployment", + "output_modalities": ["text"], + "audio": {"output": {"voice": "marin"}}, + }, +} + + +async def _azure_session_created_seen_by_a_beta_client( + session_kwargs: Mapping[str, object], +) -> tuple[str | None, object]: + client: Final = _BetaClientThatStaysConnected() + connect: Final = _ConnectToABackendThatSendsOneFrame(_GA_SESSION_CREATED) + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="azure/my-realtime-deployment", + websocket=client, + api_base="https://my-endpoint.openai.azure.com", + api_key="fake-key", + litellm_logging_obj=_realtime_websocket_logging(), + **session_kwargs, + ) + return connect.url, json.loads(client.sent[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("session_kwargs", "expected_url"), + ( + pytest.param( + {"query_params": {"intent": "transcription"}}, + "wss://my-endpoint.openai.azure.com/openai/v1/realtime?intent=transcription", + id="transcription-intent", + ), + ), +) +async def test_arealtime_azure_beta_client_on_a_ga_session_it_did_not_need_bridging_for_keeps_ga_session_events( + monkeypatch: pytest.MonkeyPatch, session_kwargs: Mapping[str, object], expected_url: str +) -> None: + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + + assert await _azure_session_created_seen_by_a_beta_client(session_kwargs) == (expected_url, _GA_SESSION_CREATED) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("env_protocol", "session_kwargs"), + ( + pytest.param(None, {"realtime_protocol": "GA"}, id="operator-set-ga"), + pytest.param("GA", {}, id="env-set-ga"), + ), +) +async def test_arealtime_azure_bridges_a_beta_client_to_ga_when_the_operator_pins_ga( + monkeypatch: pytest.MonkeyPatch, env_protocol: str | None, session_kwargs: Mapping[str, object] +) -> None: + if env_protocol is None: + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + else: + monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", env_protocol) + + assert await _azure_session_created_seen_by_a_beta_client(session_kwargs) == ( + "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=my-realtime-deployment", + _GA_SESSION_CREATED_BRIDGED_TO_BETA, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "base_model_kwargs", + ( + pytest.param({"base_model": "azure/gpt-realtime-2.1"}, id="base-model-param"), + pytest.param({"metadata": {"model_info": {"base_model": "azure/gpt-realtime-2.1"}}}, id="model-info"), + ), +) +async def test_arealtime_azure_bridges_a_beta_client_to_ga_when_the_deployment_base_model_is_ga_only( + local_model_cost_map, monkeypatch: pytest.MonkeyPatch, base_model_kwargs: Mapping[str, object] +) -> None: + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + + assert await _azure_session_created_seen_by_a_beta_client(base_model_kwargs) == ( + "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=my-realtime-deployment", + _GA_SESSION_CREATED_BRIDGED_TO_BETA, + ) + + +_GA_SESSION_CREATED_BRIDGED_TO_BETA: Final = { + "type": "session.created", + "event_id": "event_1", + "session": {"id": "sess_1", "model": "my-realtime-deployment", "modalities": ["text"], "voice": "marin"}, +} + +_ROUTER_REQUEST_SHAPES: Final = ( + pytest.param({}, id="sdk-router-call"), + pytest.param({"metadata": {"user_api_key_hash": "hashed-key"}}, id="proxy-request-metadata"), +) + + +def _router_with_an_azure_deployment_aliasing(deployment_model: str, base_model: str) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "team-alias", + "litellm_params": { + "model": deployment_model, + "api_base": "https://my-endpoint.openai.azure.com", + "api_key": "fake-key", + }, + "model_info": {"base_model": base_model}, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", _ROUTER_REQUEST_SHAPES) +async def test_router_realtime_bridges_a_beta_client_to_ga_when_model_info_base_model_is_ga_only( + local_model_cost_map, monkeypatch: pytest.MonkeyPatch, request_kwargs: Mapping[str, object] +) -> None: + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + router: Final = _router_with_an_azure_deployment_aliasing( + "azure/my-realtime-deployment", "azure/gpt-realtime-2.1-mini" + ) + client: Final = _BetaClientThatStaysConnected() + connect: Final = _ConnectToABackendThatSendsOneFrame(_GA_SESSION_CREATED) + + with patch("websockets.connect", connect): + await router._arealtime(model="team-alias", websocket=client, **request_kwargs) + + assert (connect.url, json.loads(client.sent[0])) == ( + "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=my-realtime-deployment", + _GA_SESSION_CREATED_BRIDGED_TO_BETA, ) @@ -1130,6 +1526,30 @@ async def test_arealtime_openai_forwards_the_intent_query_param_to_the_upstream_ assert connect.url == "wss://api.openai.com/v1/realtime?model=gpt-realtime&intent=chat" +@pytest.mark.asyncio +@pytest.mark.parametrize("intent", (None, "transcription")) +async def test_arealtime_openai_connects_to_api_base_rather_than_the_callers_client_base_url( + intent: str | None, +) -> None: + connect: Final = _ConnectThatStopsAfterCapturingTheUrl() + query_params: Final[RealtimeQueryParams] = ( + {"model": "openai/gateway-realtime", "intent": intent} if intent else {"model": "openai/gateway-realtime"} + ) + with patch("websockets.connect", connect), patch("websockets.asyncio.client.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="openai/gateway-realtime", + websocket=_ClosableGaClientWebSocket(), + api_key="fake-key", + api_base="https://configured-gateway.example/", + client=AsyncOpenAI(api_key="caller-key", base_url="https://caller-client.example/v1"), + query_params=query_params, + litellm_logging_obj=FakeLogging(), + ) + + expected_query: Final = "model=gateway-realtime&intent=transcription" if intent else "model=gateway-realtime" + assert connect.url == f"wss://configured-gateway.example/v1/realtime?{expected_query}" + + @pytest.mark.parametrize( ("model", "api_base", "client_query_params", "expected_backend_url"), [ @@ -1416,26 +1836,82 @@ class _GatewayThatRecordsTheAuthorization: @pytest.mark.asyncio @pytest.mark.parametrize( - "route", - [realtime_main.acreate_realtime_client_secret, realtime_main.acreate_realtime_transcription_session], - ids=["client_secrets", "transcription_sessions"], + ("route", "model"), + [ + (realtime_main.acreate_realtime_client_secret, "gpt-realtime"), + (realtime_main.acreate_realtime_transcription_session, "gpt-realtime"), + (realtime_main.acreate_realtime_translation_client_secret, "gpt-realtime-translate"), + ], + ids=["client_secrets", "transcription_sessions", "translation_client_secrets"], ) @respx.mock async def test_openai_realtime_http_routes_send_the_federated_bearer_for_a_keyless_deployment( - federated_deployment: dict[str, str], route: Callable[..., Awaitable[object]] + federated_deployment: dict[str, str], route: Callable[..., Awaitable[object]], model: str ) -> None: _mock_token_exchange() gateway: Final = _GatewayThatRecordsTheAuthorization() await route.__wrapped__( - model="gpt-realtime", + model=model, custom_llm_provider="openai", - litellm_logging_obj=FakeLogging(), + litellm_logging_obj=Logging( + model=model, + messages=[], + stream=False, + call_type="acreate_realtime_client_secret", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="test-federated-realtime", + function_id="test-federated-realtime", + ), client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), **federated_deployment, ) assert gateway.authorization == "Bearer federated-bearer" +@pytest.mark.asyncio +@pytest.mark.parametrize("flow", ("client_secret", "calls")) +async def test_translation_http_routes_preserve_explicit_azure_provider(flow: str) -> None: + expected_path: Final = ( + "/openai/v1/realtime/translations/" + + { + "client_secret": "client_secrets", + "calls": "calls", + }[flow] + ) + + def send_response(request: httpx.Request) -> httpx.Response: + assert request.url.host == "azure-deployment.example" + assert request.url.path == expected_path + return httpx.Response(200, json={"value": "ek-test", "routed_path": request.url.path}) + + params: Final[Mapping[str, object]] = { + "model": "translation-deployment", + "custom_llm_provider": "azure", + "api_base": "https://azure-deployment.example/openai/v1/", + "api_key": "test-key", + "client": AsyncHTTPHandler(transport=httpx.MockTransport(send_response)), + "litellm_logging_obj": Logging( + model="translation-deployment", + messages=[], + stream=False, + call_type="acreate_realtime_translation_client_secret", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="test-translation-routing", + function_id="test-translation-routing", + ), + } + response: Final = await ( + realtime_main.acreate_realtime_translation_client_secret.__wrapped__(**params) + if flow == "client_secret" + else realtime_main.arealtime_translation_calls.__wrapped__( + openai_ephemeral_key="ek-test", sdp_body=b"v=0", **params + ) + ) + + assert response.status_code == 200 + assert response.json() == {"value": "ek-test", "routed_path": expected_path} + + @pytest.mark.asyncio @respx.mock async def test_openai_realtime_health_check_sends_the_federated_bearer_for_a_keyless_deployment( @@ -1453,6 +1929,532 @@ async def test_openai_realtime_health_check_sends_the_federated_bearer_for_a_key assert connect.kwargs["additional_headers"] == {"Authorization": "Bearer federated-bearer"} +def _realtime_http_route_call( + route: str, + logging_obj: Logging, + custom_llm_provider: str, + api_base: str, + client: AsyncHTTPHandler | None = None, + metadata: Mapping[str, object] | None = None, +) -> Awaitable[httpx.Response]: + if route == "client_secrets": + return realtime_main.acreate_realtime_client_secret.__wrapped__( + model="gpt-realtime", + custom_llm_provider=custom_llm_provider, + api_key="sk-test", + api_base=api_base, + client=client, + litellm_logging_obj=logging_obj, + ) + if route == "translation_client_secrets": + return realtime_main.acreate_realtime_translation_client_secret.__wrapped__( + model="gpt-realtime-translate", + custom_llm_provider=custom_llm_provider, + api_key="sk-test", + api_base=api_base, + client=client, + litellm_logging_obj=logging_obj, + metadata=metadata, + ) + if route == "calls": + return realtime_main.arealtime_calls.__wrapped__( + openai_ephemeral_key="ek_test", + sdp_body=b"v=0", + model="gpt-realtime", + custom_llm_provider=custom_llm_provider, + api_base=api_base, + client=client, + litellm_logging_obj=logging_obj, + ) + return realtime_main.arealtime_translation_calls.__wrapped__( + openai_ephemeral_key="ek_test", + sdp_body=b"v=0", + model="gpt-realtime-translate", + custom_llm_provider=custom_llm_provider, + api_base=api_base, + client=client, + litellm_logging_obj=logging_obj, + metadata=metadata, + ) + + +def _realtime_http_logging(route: str) -> Logging: + return Logging( + model="gpt-realtime", + messages=[], + stream=False, + call_type="acreate_realtime_client_secret", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id=f"test-realtime-{route}", + function_id=f"test-realtime-{route}", + ) + + +_OPENAI_REALTIME_HTTP_ROUTE_PATHS: Final = { + "client_secrets": "/v1/realtime/client_secrets", + "translation_client_secrets": "/v1/realtime/translations/client_secrets", + "calls": "/v1/realtime/calls", + "translation_calls": "/v1/realtime/translations/calls", +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", tuple(_OPENAI_REALTIME_HTTP_ROUTE_PATHS)) +async def test_openai_realtime_http_routes_surface_the_upstream_error_body_verbatim(route: str) -> None: + upstream_error: Final = { + "error": { + "message": "Invalid model", + "type": "invalid_request_error", + "param": "session.model", + "code": "invalid_value", + } + } + + with pytest.raises(BaseLLMException) as raised: + await _realtime_http_route_call( + route, + _realtime_http_logging(route), + custom_llm_provider="openai", + api_base="https://upstream.example", + client=AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(400, json=upstream_error)) + ), + ) + + assert raised.value.status_code == 400 + assert json.loads(raised.value.message) == upstream_error + + +class _GatewayThatRecordsTheSession: + session: Mapping[str, object] | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.session = json.loads(request.content)["session"] + return httpx.Response(200, json={"value": "ek_test"}) + + +@pytest.mark.asyncio +async def test_client_secret_keeps_the_realtime_session_model_when_input_transcription_names_another() -> None: + gateway: Final = _GatewayThatRecordsTheSession() + logging_obj: Final = _realtime_http_logging("client_secrets") + + await realtime_main.acreate_realtime_client_secret.__wrapped__( + model="gpt-realtime", + session={ + "type": "realtime", + "model": "gpt-realtime", + "audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}, + }, + custom_llm_provider="openai", + api_key="sk-test", + api_base="https://upstream.example", + litellm_logging_obj=logging_obj, + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + ) + + assert gateway.session is not None + assert gateway.session["model"] == "gpt-realtime" + assert gateway.session["audio"] == {"input": {"transcription": {"model": "gpt-realtime-whisper"}}} + assert logging_obj.model == "gpt-realtime" + + +@pytest.mark.asyncio +async def test_client_secret_session_model_wins_over_a_provider_qualified_top_level_model( + local_model_cost_map: None, +) -> None: + gateway: Final = _GatewayThatRecordsTheSession() + logging_obj: Final = _realtime_http_logging("client_secrets") + + await realtime_main.acreate_realtime_client_secret.__wrapped__( + model="openai/gpt-realtime-2.1", + session={"type": "realtime", "model": "gpt-realtime-2.1-mini"}, + api_key="sk-test", + api_base="https://upstream.example", + litellm_logging_obj=logging_obj, + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + ) + + assert gateway.session is not None + assert gateway.session["model"] == "gpt-realtime-2.1-mini" + assert logging_obj.model == "gpt-realtime-2.1-mini" + + +class _UpstreamThatMustNotBeCalled: + called: bool = False + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.called = True + return httpx.Response(200, json={"value": "ek_test"}) + + +class _RecordingClientWebSocket: + def __init__(self) -> None: + self.scope: Final = {"headers": ()} + self.sent: tuple[str, ...] = () + self.close_code: int | None = None + + async def send_text(self, data: str) -> None: + self.sent = (*self.sent, data) + + async def close(self, code: int = 1000, reason: str = "") -> None: + self.close_code = code + + +def _block_every_request_with_a_guardrail( + monkeypatch: pytest.MonkeyPatch, event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call +) -> None: + monkeypatch.setattr( + litellm, + "callbacks", + [CustomGuardrail(guardrail_name="translation-blocker", event_hook=event_hook, default_on=True)], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ("translation_client_secrets", "translation_calls")) +@pytest.mark.parametrize("event_hook", (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call)) +async def test_translation_http_routes_refuse_a_guarded_request_before_any_upstream_call( + route: str, + monkeypatch: pytest.MonkeyPatch, + event_hook: GuardrailEventHooks, +) -> None: + _block_every_request_with_a_guardrail(monkeypatch, event_hook) + upstream: Final = _UpstreamThatMustNotBeCalled() + + with pytest.raises(litellm.BadRequestError, match=TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE) as raised: + await _realtime_http_route_call( + route, + _realtime_http_logging(route), + custom_llm_provider="openai", + api_base="https://upstream.example", + client=AsyncHTTPHandler(transport=httpx.MockTransport(upstream)), + ) + + assert raised.value.status_code == 400 + assert upstream.called is False + + +_DEFAULT_ON_GUARDRAIL_OPT_OUTS: Final = ( + pytest.param({"user_api_key_metadata": {"disable_global_guardrails": True}}, id="key-disables-global"), + pytest.param( + {"user_api_key_team_metadata": {"opted_out_global_guardrails": ["translation-blocker"]}}, id="team-opts-out" + ), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "path"), + (("translation_client_secrets", "client_secrets"), ("translation_calls", "calls")), +) +@pytest.mark.parametrize("key_metadata", _DEFAULT_ON_GUARDRAIL_OPT_OUTS) +@pytest.mark.parametrize("event_hook", (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call)) +async def test_translation_http_routes_serve_a_key_opted_out_of_the_default_on_guardrail( + route: str, + path: str, + key_metadata: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + event_hook: GuardrailEventHooks, +) -> None: + _block_every_request_with_a_guardrail(monkeypatch, event_hook) + gateway: Final = _GatewayThatRecordsTheUrl() + + await _realtime_http_route_call( + route, + _realtime_http_logging(route), + custom_llm_provider="openai", + api_base="https://upstream.example", + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + metadata=key_metadata, + ) + + assert gateway.url == f"https://upstream.example/v1/realtime/translations/{path}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ("translation_client_secrets", "translation_calls")) +async def test_translation_http_routes_refuse_a_provider_other_than_openai_and_azure_before_any_upstream_call( + route: str, +) -> None: + upstream: Final = _UpstreamThatMustNotBeCalled() + + with pytest.raises(litellm.BadRequestError, match="supports only the azure, openai providers, got xai") as raised: + await _realtime_http_route_call( + route, + _realtime_http_logging(route), + custom_llm_provider="xai", + api_base="https://upstream.example", + client=AsyncHTTPHandler(transport=httpx.MockTransport(upstream)), + ) + + assert raised.value.status_code == 400 + assert upstream.called is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_hook", (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call)) +async def test_translation_websocket_closes_1008_for_a_guarded_request_before_the_upstream_connect( + monkeypatch: pytest.MonkeyPatch, + event_hook: GuardrailEventHooks, +) -> None: + _block_every_request_with_a_guardrail(monkeypatch, event_hook) + websocket: Final = _RecordingClientWebSocket() + connect: Final = _CapturingConnect() + + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime-translate", + websocket=websocket, + api_key="fake-key", + realtime_mode="translation", + litellm_logging_obj=_realtime_websocket_logging(), + ) + + assert connect.url is None + assert websocket.close_code == 1008 + assert [json.loads(event)["error"]["message"] for event in websocket.sent] == [ + TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key_metadata", _DEFAULT_ON_GUARDRAIL_OPT_OUTS) +@pytest.mark.parametrize("event_hook", (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call)) +async def test_translation_websocket_opens_for_a_key_opted_out_of_the_default_on_guardrail( + key_metadata: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + event_hook: GuardrailEventHooks, +) -> None: + _block_every_request_with_a_guardrail(monkeypatch, event_hook) + websocket: Final = _RecordingClientWebSocket() + connect: Final = _CapturingConnect() + + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime-translate", + websocket=websocket, + api_key="fake-key", + realtime_mode="translation", + litellm_logging_obj=FakeLogging(), + metadata=key_metadata, + ) + + assert connect.url == "wss://api.openai.com/v1/realtime/translations?model=gpt-realtime-translate" + assert all(TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE not in event for event in websocket.sent) + + +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: tuple[object, ...] = () + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.failures = (*self.failures, kwargs.get("exception")) + + +@pytest.mark.asyncio +async def test_translation_websocket_refusal_logs_the_session_as_a_failure(monkeypatch: pytest.MonkeyPatch) -> None: + _block_every_request_with_a_guardrail(monkeypatch) + recorder: Final = _FailureRecorder() + logging_obj: Final = _realtime_websocket_logging(recorder) + + with patch("websockets.connect", _CapturingConnect()): + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime-translate", + websocket=_RecordingClientWebSocket(), + api_key="fake-key", + realtime_mode="translation", + litellm_logging_obj=logging_obj, + ) + + assert [(type(failure), getattr(failure, "status_code", None)) for failure in recorder.failures] == [ + (litellm.BadRequestError, 400) + ] + assert TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE in str(recorder.failures[0]) + assert logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] is True + + +class _NativeLifecycleGuardrail(CustomGuardrail): + use_native_lifecycle_hooks: ClassVar[bool] = True + + +@pytest.mark.asyncio +async def test_translation_websocket_opens_under_a_native_pre_call_guardrail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm, + "callbacks", + [ + _NativeLifecycleGuardrail( + guardrail_name="runner-skipped", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + ], + ) + websocket: Final = _RecordingClientWebSocket() + connect: Final = _CapturingConnect() + + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime-translate", + websocket=websocket, + api_key="fake-key", + realtime_mode="translation", + litellm_logging_obj=FakeLogging(), + ) + + assert connect.url == "wss://api.openai.com/v1/realtime/translations?model=gpt-realtime-translate" + assert all(TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE not in event for event in websocket.sent) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ("translation_client_secrets", "translation_calls", "websocket")) +@pytest.mark.parametrize( + ("key_metadata", "blocked"), + ( + pytest.param({"user_api_key_metadata": {"disable_global_guardrails": True}}, False, id="key-disables-global"), + pytest.param({}, True, id="no-opt-out"), + ), +) +async def test_translation_lets_only_an_opted_out_key_past_a_legacy_string_callback_guardrail( + route: str, + key_metadata: Mapping[str, object], + blocked: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + initialize_callbacks_on_proxy( + value=["lakera_prompt_injection"], premium_user=False, config_file_path="", litellm_settings={} + ) + if route == "websocket": + websocket: Final = _RecordingClientWebSocket() + connect: Final = _CapturingConnect() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime-translate", + websocket=websocket, + api_key="fake-key", + realtime_mode="translation", + litellm_logging_obj=_realtime_websocket_logging(), + metadata=key_metadata, + ) + if blocked: + assert connect.url is None + assert websocket.close_code == 1008 + assert TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE in websocket.sent[0] + else: + assert connect.url == "wss://api.openai.com/v1/realtime/translations?model=gpt-realtime-translate" + assert all(TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE not in event for event in websocket.sent) + return + gateway: Final = _GatewayThatRecordsTheUrl() + call: Final = _realtime_http_route_call( + route, + _realtime_http_logging(route), + custom_llm_provider="openai", + api_base="https://upstream.example", + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + metadata=key_metadata, + ) + if blocked: + with pytest.raises(litellm.BadRequestError, match=TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE): + await call + assert gateway.url is None + else: + await call + assert gateway.url is not None + + +@pytest.mark.asyncio +async def test_translation_websocket_closes_1008_for_a_provider_other_than_openai_and_azure() -> None: + websocket: Final = _RecordingClientWebSocket() + connect: Final = _CapturingConnect() + + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="xai/grok-voice", + websocket=websocket, + api_key="fake-key", + realtime_mode="translation", + litellm_logging_obj=_realtime_websocket_logging(), + ) + + assert connect.url is None + assert websocket.close_code == 1008 + assert [json.loads(event)["error"]["message"] for event in websocket.sent] == [ + "Realtime translation supports only the azure, azure_ai, openai providers, got xai" + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ("translation_client_secrets", "translation_calls", "websocket")) +@pytest.mark.parametrize("mode", ("pre_call", "post_call")) +@pytest.mark.parametrize("metadata_bucket", ("metadata", "litellm_metadata")) +@pytest.mark.parametrize("selected,condition_matches", ((True, True), (True, False), (False, True))) +async def test_translation_rejects_only_applicable_policy_pipelines( + route: str, + mode: Literal["pre_call", "post_call"], + metadata_bucket: str, + selected: bool, + condition_matches: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + request_data: Final[dict[str, object]] = {"metadata": {}, "litellm_metadata": {}} + _apply_resolved_guardrails_to_metadata( + data=request_data, + metadata_variable_name=metadata_bucket, + context=PolicyMatchContext(model="translation-deployment"), + policy_names=["translation-policy"] if selected else [], + policies={ + "translation-policy": Policy( + condition=PolicyCondition(model="translation-deployment" if condition_matches else "unrelated-model"), + pipeline=GuardrailPipeline(mode=mode, steps=[PipelineStep(guardrail="speech-policy", on_fail="block")]), + ) + }, + ) + metadata: Final = realtime_main._build_litellm_metadata(request_data) + blocked: Final = selected and condition_matches + if route == "websocket": + websocket: Final = _RecordingClientWebSocket() + connect: Final = _CapturingConnect() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime-translate", + websocket=websocket, + api_key="fake-key", + realtime_mode="translation", + litellm_logging_obj=_realtime_websocket_logging(), + litellm_metadata=metadata, + ) + if blocked: + assert connect.url is None + assert websocket.close_code == 1008 + assert TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE in websocket.sent[0] + else: + assert connect.url == "wss://api.openai.com/v1/realtime/translations?model=gpt-realtime-translate" + return + upstream: Final = _UpstreamThatMustNotBeCalled() + call: Final = _realtime_http_route_call( + route, + _realtime_http_logging(route), + custom_llm_provider="openai", + api_base="https://upstream.example", + client=AsyncHTTPHandler(transport=httpx.MockTransport(upstream)), + metadata=metadata, + ) + if blocked: + with pytest.raises(litellm.BadRequestError, match=TRANSLATION_GUARDRAILS_UNSUPPORTED_MESSAGE) as raised: + await call + assert raised.value.status_code == 400 + assert upstream.called is False + else: + response: Final = await call + assert response.status_code == 200 + assert upstream.called is True + + @pytest.mark.asyncio async def test_arealtime_routes_soniox_with_the_deployment_options(monkeypatch): from litellm.llms.soniox.realtime.transformation import SonioxRealtimeConfig diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b4cd606374a..012926c6cb3 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6564,6 +6564,31 @@ def test_an_application_inference_profile_gets_no_item_marks_under_drop_params() assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 0 + +def test_a_bedrock_deployment_keeps_item_marks_by_its_own_flag_not_a_siblings_on_the_same_model( + local_model_cost_map: None, +): + from litellm import Router + + Router( + model_list=[ + { + "model_name": "opted-in", + "litellm_params": {"model": "bedrock/moonshotai.kimi-k3", "aws_region_name": "us-east-1"}, + "model_info": {"id": "opted-in-id", "supports_prompt_cache_breakpoint": True}, + } + ] + ) + marked_system: Final = {"role": "system", "content": "House rules", "cache_control": _EPHEMERAL} + + def bridged_system(model_info: Mapping[str, object]) -> Mapping[str, object]: + return _bridged_request(marked_system, model="bedrock/moonshotai.kimi-k3", model_info=model_info)[1] + + assert ( + bridged_system({"id": "flagless-id"}), + bridged_system({"id": "opted-in-id", "supports_prompt_cache_breakpoint": True}), + ) == ({"role": "system", "content": "House rules"}, marked_system) + def test_the_shared_input_transform_keeps_every_item_mark_for_a_request_it_sends(): messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=cast(ResponseInputParam, list(_marked_turns(5))), # cast-ok: cache_control is not an input TypedDict key @@ -7192,6 +7217,29 @@ async def test_bridge_partial_response_on_disconnect_bills_only_the_input_since_ ) == await _input_tokens_billed_on_disconnect(since_compaction) +def test_a_hosted_mcp_tool_sent_to_responses_on_bedrock_converse_answers_400(monkeypatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.responses( + model="bedrock/converse/us.openai.gpt-oss-120b-1:0", + input="What does the deepwiki MCP server say about litellm?", + tools=[ + { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", + } + ], + aws_access_key_id="AKIAHOSTEDMCP", + aws_secret_access_key="hosted-mcp-secret", + aws_region_name="us-east-1", + ) + assert raised.value.status_code == 400 + assert "`deepwiki`" in str(raised.value) + assert "litellm_proxy/mcp/" in str(raised.value) + + _ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-5-5" _WEATHER_ASK: Final = {"role": "user", "content": "Weather in Paris?"} diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index f1d841a8e5e..9a8fa19e2e4 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1083,6 +1083,128 @@ async def test_previous_response_history_starts_at_a_compaction_item_in_a_stored ] +def _stored_response_turn( + request_id: str, call_type: str, status: str, output: list[dict[str, object]], prompt: str | None = None +) -> dict[str, object]: + return { + "request_id": request_id, + "call_type": call_type, + "session_id": "session-with-reads", + "proxy_server_request": json.dumps({} if prompt is None else {"input": [{"role": "user", "content": prompt}]}), + "response": {"id": request_id, "object": "response", "status": status, "output": output}, + } + + +def _identified_output_message(item_id: str, text: str) -> dict[str, object]: + return {**_output_message(text), "id": item_id} + + +def _replayed_roles_and_texts(session: ChatCompletionSession) -> list[tuple[object, object]]: + return [ + (role, _text_of(content)) for role, content in (_role_and_content(message) for message in session["messages"]) + ] + + +@pytest.mark.asyncio +async def test_reads_and_cancels_of_a_response_the_session_already_holds_replay_no_turns(): + answer: Final = [_identified_output_message("msg_answer", "OK")] + spend_logs: Final = ( + _stored_response_turn("resp_create", "aresponses", "completed", answer, prompt="Reply with OK"), + _stored_response_turn("resp_read", "aget_responses", "completed", answer), + _stored_response_turn("resp_streamed_read", "aget_responses", "completed", answer), + _stored_response_turn("resp_cancel", "acancel_responses", "completed", answer), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_create") + + assert _replayed_roles_and_texts(session) == [("user", "Reply with OK"), ("assistant", "OK")] + + +@pytest.mark.asyncio +async def test_a_polled_background_response_replays_its_finished_output_once(): + answer: Final = [_identified_output_message("msg_answer", "One, two, three")] + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to three"), + _stored_response_turn("resp_poll_running", "aget_responses", "in_progress", []), + _stored_response_turn("resp_poll_done", "aget_responses", "completed", answer), + _stored_response_turn("resp_poll_again", "aget_responses", "completed", answer), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to three"), ("assistant", "One, two, three")] + + +@pytest.mark.asyncio +async def test_a_background_response_that_grows_between_polls_replays_each_output_item_once(): + first: Final = _identified_output_message("msg_first", "One") + second: Final = _identified_output_message("msg_second", "Two") + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to two"), + _stored_response_turn("resp_poll_partial", "aget_responses", "in_progress", [first]), + _stored_response_turn("resp_poll_done", "aget_responses", "completed", [first, second]), + _stored_response_turn("resp_poll_again", "aget_responses", "completed", [first, second]), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to two"), ("assistant", "One"), ("assistant", "Two")] + + +@pytest.mark.asyncio +async def test_an_item_a_poll_caught_half_written_replays_its_finished_text_once(): + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to three"), + _stored_response_turn( + "resp_poll_partial", "aget_responses", "in_progress", [_identified_output_message("msg_answer", "One, t")] + ), + _stored_response_turn( + "resp_poll_done", + "aget_responses", + "completed", + [_identified_output_message("msg_answer", "One, two, three"), _identified_output_message("msg_more", "Done")], + ), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [ + ("user", "Count to three"), + ("assistant", "One, two, three"), + ("assistant", "Done"), + ] + + +@pytest.mark.asyncio +async def test_a_polled_output_item_without_an_id_still_replays_in_order(): + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to two"), + _stored_response_turn( + "resp_poll_done", + "aget_responses", + "completed", + [_output_message("One"), _identified_output_message("msg_two", "Two")], + ), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to two"), ("assistant", "One"), ("assistant", "Two")] + + +@pytest.mark.asyncio +async def test_a_read_that_ends_before_its_streamed_create_replays_the_answer_once_after_the_prompt(): + answer: Final = [_identified_output_message("msg_answer", "OK")] + spend_logs: Final = ( + _stored_response_turn("resp_read_mid_stream", "aget_responses", "completed", answer), + _stored_response_turn("resp_create", "aresponses", "completed", answer, prompt="Say OK"), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_create") + + assert _replayed_roles_and_texts(session) == [("user", "Say OK"), ("assistant", "OK")] + + def _tool_call_ids(message: object) -> tuple[str, ...]: tool_calls: Final = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) return tuple(call["id"] if isinstance(call, dict) else call.id for call in tool_calls or ()) diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 684f5a89e80..87858741908 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -477,7 +477,6 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): assert evt is not None assert evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED assert evt.output_index == 0 - assert hasattr(evt, "__dict__") and "sequence_number" in evt.__dict__ # Collect all remaining delta events from the pending queue by creating empty chunks delta_events = [] @@ -511,17 +510,11 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): assert len(evt.delta) <= 10 assert evt.item_id == "fc_call_test" assert evt.output_index == 0 - assert hasattr(evt, "__dict__") and "sequence_number" in evt.__dict__ # Verify all deltas concatenated equal the original arguments concatenated = "".join(evt.delta for evt in delta_events) assert concatenated == large_arguments - # Verify sequence numbers are increasing - sequence_numbers = [evt.__dict__["sequence_number"] for evt in delta_events] - assert sequence_numbers == sorted(sequence_numbers) - assert len(set(sequence_numbers)) == len(sequence_numbers) # All unique - def test_tool_call_delta_without_id_uses_index_mapping(): iterator = LiteLLMCompletionStreamingIterator( @@ -1990,6 +1983,65 @@ def test_the_partial_response_billed_on_disconnect_carries_the_logging_objects_b ), "the bill is priced on the partial's own usage" +def _serialized_sequence_numbers(events: Iterable[BaseLiteLLMOpenAIResponseObject]) -> list[int]: + payloads: Final = tuple(json.loads(event.model_dump_json()) for event in events) + assert all("sequence_number" in payload for payload in payloads), [ + payload["type"] for payload in payloads if "sequence_number" not in payload + ] + return [payload["sequence_number"] for payload in payloads] + + +def _long_text_chunks() -> tuple[ModelResponseStream, ...]: + return (*(_chunk(f"word{i} ") for i in range(300)), _chunk("end", finish_reason="stop")) + + +def _text_then_tool_call_chunks() -> tuple[ModelResponseStream, ...]: + return (_text_opening_a_tool_call_chunk("Checking. "), _tool_arguments_chunk('{"x": 1}', finish_reason="tool_calls")) + + +@pytest.mark.parametrize( + "stream", + [ + lambda: (_chunk("Hello"), _chunk("!", finish_reason="stop")), + _long_text_chunks, + _text_then_tool_call_chunks, + _reasoning_then_text_chunks, + _reasoning_and_text_in_one_chunk, + _reasoning_citation_then_tool_call_finishing_chunks, + lambda: (_tool_call_chunk(), _chunk("", finish_reason="tool_calls")), + lambda: (_refusal_chunk("I can't "), _refusal_chunk("help."), _chunk("", finish_reason="stop")), + lambda: (_chunk("Sure. "), _refusal_chunk("I can't help."), _chunk("", finish_reason="stop")), + ], + ids=[ + "text", + "long_text", + "text_then_tool", + "reasoning_then_text", + "reasoning_and_text", + "citation_tool", + "tool", + "refusal", + "text_then_refusal", + ], +) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_every_event_serializes_its_position_in_the_stream_from_zero(stream: StreamChunks, sync_mode: bool): + iterator: Final = _build_iterator(stream()) + events: Final = list(iterator) if sync_mode else [event async for event in iterator] + + assert _serialized_sequence_numbers(events) == list(range(len(events))) + + +def test_bridged_events_answer_dict_reads_of_their_number(): + events: Final = list(_build_iterator((_chunk("Hello"), _chunk("!", finish_reason="stop")))) + positions: Final = list(range(len(events))) + + assert [event["sequence_number"] for event in events] == positions + assert [event.get("sequence_number") for event in events] == positions + assert all("sequence_number" in event for event in events) + + def _refusal_chunk(refusal: str, finish_reason: str | None = None) -> ModelResponseStream: return ModelResponseStream( id=CHAT_COMPLETION_ID, diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index f98d7ca188e..5072f3d3108 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -2867,3 +2867,54 @@ async def test_websocket_turns_logging_survives_garbage_collection(terminal_even await _run_websocket_turn(logger, terminal_event) assert await _logging_tasks_left_after_garbage_collection(logger) == 1 + + +@pytest.mark.parametrize( + "output", + [ + [ + { + "type": "message", + "id": "msg_seq", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello world", "annotations": []}], + } + ], + [ + { + "type": "reasoning", + "id": "rs_seq", + "summary": [{"type": "summary_text", "text": "thinking about it"}], + }, + { + "type": "function_call", + "id": "fc_seq", + "call_id": "call_seq", + "name": "lookup", + "arguments": '{"id":1}', + }, + ], + ], + ids=["message", "reasoning_and_tool"], +) +def test_synthetic_stream_serializes_every_event_position_from_zero(output: list[dict[str, object]]): + transformed: Final = ResponsesAPIResponse.model_validate( + { + "id": "resp_seq", + "created_at": 1741476542, + "status": "completed", + "model": "claude-sonnet-5-5", + "object": "response", + "output": output, + } + ) + + events: Final = streaming_module.build_synthetic_response_events( + transformed=transformed, + logging_obj=None, + chunk_size=5, + ) + + payloads: Final = tuple(json.loads(event.model_dump_json()) for event in events) + assert [payload.get("sequence_number") for payload in payloads] == list(range(len(payloads))) diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index ef7b9e485d4..d0e2b381253 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -359,6 +359,52 @@ class TestComplexityRouterInit: "tier-probability:MEDIUM=0.900000", ) + @pytest.mark.asyncio + async def test_jev_choice_without_confidence_routes_and_reports_no_confidence(self, mock_router_instance): + client = _StaticJevClient( + JevSystemOneResponse( + answers={ + "tier": JevChoiceAnswer( + type="choice", + choice="MEDIUM", + probabilities={"SIMPLE": 0.1, "MEDIUM": 0.9}, + ) + }, + ) + ) + router = ComplexityRouter( + "test-router", + mock_router_instance, + { + "classifier_type": "jev", + "jev_classifier_config": {"api_key": "test", "timeout_ms": 100}, + "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}, + }, + derive_savings_baseline=False, + jev_client=client, + ) + + outcome = await router.aclassify("Explain this") + result = await router.async_pre_routing_hook( + model="test-router", + request_kwargs={}, + messages=[{"role": "user", "content": "Explain this"}], + ) + + assert outcome.tier == ComplexityTier.MEDIUM + assert outcome.cause == "jev_classifier" + assert outcome.jev_verdict is not None + assert outcome.jev_verdict.confidence is None + assert outcome.signals == ( + "jev-classifier:MEDIUM", + "tier-probability:SIMPLE=0.100000", + "tier-probability:MEDIUM=0.900000", + ) + assert result is not None + assert result.routing_decision is not None + assert result.routing_decision["classifier_probabilities"] == {"SIMPLE": 0.1, "MEDIUM": 0.9} + assert "classifier_confidence" not in result.routing_decision + @pytest.mark.asyncio async def test_jev_pre_routing_hook_exposes_routing_decision_provenance( self, mock_router_instance, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index d184bf84035..a4d279e32e6 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -15,17 +15,21 @@ The mechanism works without any cache and supports two encoding strategies: encrypted_content back to their original forms before sending to the upstream provider. """ +import datetime import time +from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, patch import pytest import litellm +from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_api import RoutingContext from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem @pytest.fixture(autouse=True) @@ -1100,14 +1104,7 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance(): plain_key = EncryptedContentAffinityCheck._encryption_boundary_key(plain_params) assert pydantic_key is not None - assert ( - pydantic_key - == plain_key - == ( - "https://mateo-resource.openai.azure.com", - "fake-azure-resource-key-a", - ) - ) + assert pydantic_key == plain_key def test_boundary_key_resolves_missing_values_from_named_credential(): @@ -1131,14 +1128,35 @@ def test_boundary_key_resolves_missing_values_from_named_credential(): ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"}) - assert boundary == ("https://account-a.example.com", "credential-key-a") + assert boundary is not None + assert boundary == EncryptedContentAffinityCheck._encryption_boundary_key( + {"api_base": "https://account-a.example.com", "api_key": "credential-key-a"} + ) + assert boundary != EncryptedContentAffinityCheck._encryption_boundary_key( + {"api_base": "https://account-a.example.com", "api_key": "credential-key-b"} + ) -def test_boundary_key_matches_named_credential_precedence(): +@pytest.mark.parametrize( + "deployment_values", + [ + pytest.param({"api_base": "https://deployment.example.com", "api_key": "deployment-key"}, id="set values"), + pytest.param({"api_base": "", "api_key": ""}, id="empty values"), + ], +) +@pytest.mark.usefixtures("no_process_wide_api_base") +def test_boundary_key_resolves_named_credential_values_the_way_the_request_does( + monkeypatch: pytest.MonkeyPatch, deployment_values: dict[str, str] +): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) + from litellm.utils import with_named_credential + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + deployment: Final = {"model": "openai/gpt-5.5", **deployment_values, "litellm_credential_name": "account-a"} with patch.object( litellm, "credential_list", @@ -1153,46 +1171,14 @@ def test_boundary_key_matches_named_credential_precedence(): ) ], ): - boundary = EncryptedContentAffinityCheck._encryption_boundary_key( - { - "api_base": "https://deployment.example.com", - "api_key": "deployment-key", - "litellm_credential_name": "account-a", - } - ) + boundary: Final = EncryptedContentAffinityCheck._encryption_boundary_key(deployment) + sent_params: Final = with_named_credential(deployment) - assert boundary == ("https://credential.example.com", "credential-key-a") - - -def test_boundary_key_resolves_credential_when_explicit_values_are_empty(): - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, + assert boundary is not None + assert boundary == EncryptedContentAffinityCheck._encryption_boundary_key( + {key: value for key, value in sent_params.items() if key != "litellm_credential_name"} ) - with patch.object( - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ): - boundary = EncryptedContentAffinityCheck._encryption_boundary_key( - { - "api_base": "", - "api_key": "", - "litellm_credential_name": "account-a", - } - ) - - assert boundary == ("https://credential.example.com", "credential-key-a") - def test_boundary_fallback_matches_deployments_with_same_named_credential_values(): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( @@ -2728,3 +2714,604 @@ async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_th assert [deployment["model_info"]["id"] for deployment in result] == ["region-a"] assert reasoning_input[0] == original_reasoning + + +def _default_endpoint_deployment(model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]: + return {"model_name": "gpt-5.5", "litellm_params": dict(litellm_params), "model_info": {"id": model_id}} + + +def _deployment_of(ctx: RoutingContext, model_id: str) -> dict[str, object]: + deployment: Final = ctx.registry().deployment(model_id) + assert deployment is not None + return deployment.model_dump(exclude_none=True) + + +SUMMARY_ONLY_REASONING: Final = {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} +UNMARKED_REASONING: Final = {**SUMMARY_ONLY_REASONING, "id": "rs_test", "encrypted_content": "gAAAAA-blob"} +OPENAI_SHARED_KEY: Final = {"model": "openai/gpt-5.5", "api_key": "shared-key"} + + +@pytest.fixture +def no_process_wide_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + + +@pytest.mark.usefixtures("no_process_wide_api_base") +@pytest.mark.parametrize("marked", [True, False], ids=["marked origin", "unmarked origin"]) +@pytest.mark.parametrize( + "origin_params, target_params, target_can_decrypt", + [ + pytest.param(OPENAI_SHARED_KEY, OPENAI_SHARED_KEY, True, id="same key on the provider default endpoint"), + pytest.param( + {"model": "gpt-5.5", "custom_llm_provider": "openai", "api_key": "shared-key"}, + OPENAI_SHARED_KEY, + True, + id="provider named apart from the model", + ), + pytest.param(OPENAI_SHARED_KEY, {**OPENAI_SHARED_KEY, "api_key": "other-key"}, False, id="different key"), + pytest.param( + OPENAI_SHARED_KEY, + {**OPENAI_SHARED_KEY, "model": "openai/gpt-5.5-mini"}, + False, + id="same key on the default endpoint for another model", + ), + pytest.param( + {"model": "bedrock_mantle/us-west-2/openai.gpt-6-astra", "api_key": "shared-key"}, + {"model": "bedrock_mantle/us-west-2/openai.gpt-6-astra", "api_key": "shared-key"}, + False, + id="same key on Bedrock Mantle, which serves several vendors", + ), + pytest.param( + {"model": "openrouter/openai/gpt-6.1-sol", "api_key": "shared-key"}, + {"model": "openrouter/anthropic/claude-haiku-5.5", "api_key": "shared-key"}, + False, + id="same key on an aggregator that serves several vendors", + ), + pytest.param( + {"model": "anthropic/claude-opus-5-5", "api_key": "shared-key"}, + {"model": "anthropic/claude-opus-5-5", "api_key": "shared-key"}, + False, + id="provider without a Responses URL builder", + ), + pytest.param( + {"model": "azure/gpt-5.5", "api_key": "shared-key"}, + {"model": "azure/gpt-5.5", "api_key": "shared-key"}, + False, + id="provider whose builder needs an api_base", + ), + pytest.param( + {"model": "github_copilot/gpt-5.5", "api_key": "shared-key"}, + {"model": "github_copilot/gpt-5.5", "api_key": "shared-key"}, + False, + id="same key on a provider that authenticates with its own token", + ), + pytest.param( + {"model": "groq/llama-4-scout"}, + {"model": "groq/llama-4-scout"}, + False, + id="provider with no default key", + ), + pytest.param( + {"model": "vertex_ai/gemini-3-pro"}, + {"model": "vertex_ai/gemini-3-pro"}, + False, + id="provider whose default key is a credentials file", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "api_base": "https://openai-proxy.example.com/v1"}, + OPENAI_SHARED_KEY, + False, + id="explicit api_base beside the default endpoint", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Organization": "org-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"openai-organization": "org-b"}}, + False, + id="same key sent to different OpenAI organizations", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Project": "proj-a"}}, + OPENAI_SHARED_KEY, + False, + id="same key with an OpenAI project beside one without", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"Cookie": "session=account-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"cookie": "session=account-b"}}, + False, + id="same key signed in to different accounts by cookie", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "headers": {"api-key": "account-a"}}, + {**OPENAI_SHARED_KEY, "headers": {"api-key": "account-b"}}, + False, + id="explicit key beside different api-key headers", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "organization": "org-a"}, + {**OPENAI_SHARED_KEY, "organization": "org-b"}, + True, + id="organization param the Responses request never sends", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"x-team": "a", "OpenAI-Organization": "org-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"x-team": "b", "OpenAI-Organization": "org-a"}}, + True, + id="same OpenAI organization beside unrelated headers", + ), + pytest.param( + { + **OPENAI_SHARED_KEY, + "headers": {"OpenAI-Organization": "org-a"}, + "extra_headers": {"openai-organization": "org-b"}, + }, + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Organization": "org-b"}}, + True, + id="extra headers override the organization the headers set", + ), + ], +) +def test_hop_strip_keeps_reasoning_between_same_key_deployments_on_the_default_endpoint( + origin_params: dict[str, object], target_params: dict[str, object], target_can_decrypt: bool, marked: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", origin_params), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] if marked else dict(UNMARKED_REASONING) + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, + request_input, + None, + (_deployment_of(ctx, "target"),), + unmarked_origin=None if marked else "origin", + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +@pytest.mark.usefixtures("no_process_wide_api_base") +@pytest.mark.parametrize( + "target_api_key, target_can_decrypt", + [ + pytest.param(None, True, id="both on the environment key"), + pytest.param("env-key", True, id="explicit key equal to the environment key"), + pytest.param("other-key", False, id="explicit key other than the environment key"), + ], +) +def test_hop_strip_resolves_a_missing_key_from_the_environment( + monkeypatch: pytest.MonkeyPatch, target_api_key: str | None, target_can_decrypt: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + target_params: Final = {"model": "openai/gpt-5.5", **({"api_key": target_api_key} if target_api_key else {})} + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", {"model": "openai/gpt-5.5"}), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +def test_hop_strip_leaves_input_items_that_are_not_objects_as_sent(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + request_input: Final = ["and sunsets?", *_reasoning_input_minted_by("origin")] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + None, request_input, None, (), unmarked_origin=None + ) + + assert request_input == ["and sunsets?", SUMMARY_ONLY_REASONING, {"role": "user", "content": "and sunsets?"}] + + +GATEWAY_BASE: Final = {"model": "openai/gpt-5.5", "api_base": "https://gateway.example.com/v1"} + + +@pytest.mark.parametrize( + "origin_params, target_params, target_can_decrypt", + [ + pytest.param( + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-a"}}, + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-b"}}, + False, + id="environment key replaced by different Authorization headers", + ), + pytest.param( + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {"x-idempotency-key": "a"}}, + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {"x-idempotency-key": "b"}}, + True, + id="explicit pair beside different tag headers", + ), + pytest.param( + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-a"}}, + {**GATEWAY_BASE, "extra_headers": {"authorization": "Bearer sk-account-a"}}, + True, + id="same Authorization header on both", + ), + ], +) +def test_hop_strip_keeps_accounts_that_sign_in_through_headers_apart( + monkeypatch: pytest.MonkeyPatch, + origin_params: dict[str, object], + target_params: dict[str, object], + target_can_decrypt: bool, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", origin_params), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +@pytest.mark.parametrize( + "environment_key, target_model, target_can_decrypt", + [ + pytest.param("shared-key", GATEWAY_BASE["model"], True, id="environment key equal to the explicit key"), + pytest.param("other-key", GATEWAY_BASE["model"], False, id="environment key other than the explicit key"), + pytest.param("shared-key", "openai/gpt-5.5-mini", False, id="equal keys serving another model"), + ], +) +def test_hop_strip_matches_an_explicit_key_and_the_environment_key_on_one_api_base( + monkeypatch: pytest.MonkeyPatch, environment_key: str, target_model: str, target_can_decrypt: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", environment_key) + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", {**GATEWAY_BASE, "api_key": "shared-key"}), + _default_endpoint_deployment("target", {**GATEWAY_BASE, "model": target_model}), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +class _SecretReadCounter(CustomSecretManager): + def __init__(self) -> None: + super().__init__(secret_manager_name="read-counter") + self.reads: int = 0 + + async def async_read_secret( + self, secret_name: str, optional_params: Mapping[str, object] | None = None, timeout: object = None + ) -> str | None: + self.reads += 1 + return None + + def sync_read_secret( + self, secret_name: str, optional_params: Mapping[str, object] | None = None, timeout: object = None + ) -> str | None: + self.reads += 1 + return None + + +@pytest.mark.parametrize( + "deployment_params, reasoning_minted_by", + [ + pytest.param(OPENAI_SHARED_KEY, None, id="no reasoning in the input"), + pytest.param(OPENAI_SHARED_KEY, "target", id="reasoning the hop's target minted"), + pytest.param( + {"model": "groq/llama-4-scout"}, "origin", id="reasoning from a provider with no default boundary" + ), + pytest.param({**GATEWAY_BASE, "api_key": "shared-key"}, "origin", id="reasoning between explicit pairs"), + ], +) +def test_hop_strip_reads_no_secrets_when_no_default_boundary_needs_resolving( + monkeypatch: pytest.MonkeyPatch, deployment_params: dict[str, object], reasoning_minted_by: str | None +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", deployment_params), + _default_endpoint_deployment("target", deployment_params), + ] + ) + ctx: Final = RoutingContext(router) + secret_reads: Final = _SecretReadCounter() + monkeypatch.setattr(litellm, "secret_manager_client", secret_reads) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + request_input: Final = ( + _reasoning_input_minted_by(reasoning_minted_by) + if reasoning_minted_by + else [{"role": "user", "content": "and sunsets?"}] + ) + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin="origin" + ) + + assert secret_reads.reads == 0 + + +@pytest.mark.parametrize( + "unreadable_params", + [ + pytest.param( + {"model": "openai/gpt-5.5", "litellm_credential_name": "mistyped"}, + id="named credential with a value of the wrong type", + ), + pytest.param( + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {1: "one"}}, + id="extra header whose name is not a string", + ), + ], +) +def test_hop_strip_keeps_reasoning_a_target_can_decrypt_beside_a_target_with_unreadable_params( + monkeypatch: pytest.MonkeyPatch, unreadable_params: dict[str, object] +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="mistyped", + credential_values={"api_key": "shared-key", "api_version": datetime.date(2024, 10, 21)}, + credential_info={}, + ) + ], + ) + shared_pair: Final = {**GATEWAY_BASE, "api_key": "shared-key"} + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", shared_pair), + _default_endpoint_deployment("unreadable", unreadable_params), + _default_endpoint_deployment("target", shared_pair), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, + request_input, + None, + (_deployment_of(ctx, "unreadable"), _deployment_of(ctx, "target")), + unmarked_origin=None, + ) + + assert request_input[0] == original_reasoning + + +def test_boundary_key_puts_pydantic_and_plain_params_on_one_default_endpoint_boundary(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + from litellm.types.router import LiteLLM_Params + + pydantic_key: Final = EncryptedContentAffinityCheck._encryption_boundary_key(LiteLLM_Params(**OPENAI_SHARED_KEY)) + plain_key: Final = EncryptedContentAffinityCheck._encryption_boundary_key(dict(OPENAI_SHARED_KEY)) + + assert pydantic_key is not None + assert plain_key is not None + assert pydantic_key.is_shared_with(plain_key) + assert plain_key.is_shared_with(pydantic_key) + + +def test_boundary_key_has_no_default_endpoint_for_a_model_no_provider_claims(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + unmapped: Final = {"model": "unmapped-model-name", "api_key": "shared-key"} + + assert EncryptedContentAffinityCheck._encryption_boundary_key(unmapped) is None + + +def test_boundary_key_has_no_default_key_for_a_custom_provider(monkeypatch: pytest.MonkeyPatch): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "qa-custom"]) + monkeypatch.setattr(litellm, "_custom_providers", ["qa-custom"]) + + assert EncryptedContentAffinityCheck._encryption_boundary_key({"model": "qa-custom/some-model"}) is None + + +@pytest.mark.asyncio +async def test_affinity_pins_to_the_same_key_peer_on_the_default_endpoint_when_the_origin_cannot_serve(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", OPENAI_SHARED_KEY), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + reasoning_input: Final = _reasoning_input_minted_by("origin") + original_reasoning: Final = dict(reasoning_input[0]) + + result: Final[list[dict[str, object]]] = await EncryptedContentAffinityCheck(ctx=ctx).async_filter_deployments( + model="gpt-5.5", + healthy_deployments=[peer, _deployment_of(ctx, "other-org")], + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input}, + ) + + assert result == [peer] + assert reasoning_input[0] == original_reasoning + + +class _HandAdvancedClock: + def __init__(self) -> None: + self.now: float = 0.0 + + def __call__(self) -> float: + return self.now + + +def _count_secret_reads(monkeypatch: pytest.MonkeyPatch) -> _SecretReadCounter: + secret_reads: Final = _SecretReadCounter() + monkeypatch.setattr(litellm, "secret_manager_client", secret_reads) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + return secret_reads + + +async def _pin_a_turn_minted_by_the_origin( + affinity: object, healthy_deployments: list[dict[str, object]] +) -> tuple[list[dict[str, object]], Mapping[str, object]]: + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + assert isinstance(affinity, EncryptedContentAffinityCheck) + reasoning_input: Final = _reasoning_input_minted_by("origin") + result: Final[list[dict[str, object]]] = await affinity.async_filter_deployments( + model="gpt-5.5", + healthy_deployments=list(healthy_deployments), + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input}, + ) + return result, reasoning_input[0] + + +@pytest.mark.asyncio +async def test_repeated_pins_to_a_default_endpoint_peer_ask_the_secret_manager_once_a_minute( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", OPENAI_SHARED_KEY), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + healthy_deployments: Final = [peer, _deployment_of(ctx, "other-org")] + secret_reads: Final = _count_secret_reads(monkeypatch) + EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window.cache_clear() + clock: Final = _HandAdvancedClock() + affinity: Final = EncryptedContentAffinityCheck(ctx=ctx, clock=clock) + + first_pin, first_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + reads_for_the_first_pin: Final = secret_reads.reads + clock.now = 59.0 + second_pin, second_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + reads_within_the_minute: Final = secret_reads.reads + clock.now = 60.0 + third_pin, third_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + + assert first_pin == second_pin == third_pin == [peer] + assert first_reasoning == second_reasoning == third_reasoning == _reasoning_input_minted_by("origin")[0] + assert reads_for_the_first_pin > 0 + assert reads_within_the_minute == reads_for_the_first_pin + assert secret_reads.reads > reads_within_the_minute + + +@pytest.mark.asyncio +async def test_a_changed_named_credential_moves_the_pin_within_the_secret_manager_refresh_minute( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + def origin_credential(api_key: str) -> list[CredentialItem]: + return [ + CredentialItem(credential_name="origin-account", credential_values={"api_key": api_key}, credential_info={}) + ] + + monkeypatch.setattr(litellm, "credential_list", origin_credential("shared-key")) + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment( + "origin", {"model": "openai/gpt-5.5", "litellm_credential_name": "origin-account"} + ), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + other_org: Final = _deployment_of(ctx, "other-org") + _count_secret_reads(monkeypatch) + EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window.cache_clear() + affinity: Final = EncryptedContentAffinityCheck(ctx=ctx, clock=_HandAdvancedClock()) + + pin_before_the_change, _ = await _pin_a_turn_minted_by_the_origin(affinity, [peer, other_org]) + monkeypatch.setattr(litellm, "credential_list", origin_credential("other-key")) + pin_after_the_change, reasoning_after_the_change = await _pin_a_turn_minted_by_the_origin( + affinity, [peer, other_org] + ) + + assert pin_before_the_change == [peer] + assert pin_after_the_change == [other_org] + assert reasoning_after_the_change == _reasoning_input_minted_by("origin")[0] diff --git a/tests/unit/router_utils/test_prompt_caching_cache.py b/tests/unit/router_utils/test_prompt_caching_cache.py index ce8f04b8c54..1500d7433d2 100644 --- a/tests/unit/router_utils/test_prompt_caching_cache.py +++ b/tests/unit/router_utils/test_prompt_caching_cache.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Mapping from typing import Final, cast import pytest @@ -377,6 +378,32 @@ def test_extract_cacheable_prefix_mixed_string_and_list_content(): assert isinstance(result[2]["content"], list) +def test_prompt_cache_breakpoint_marks_the_same_cacheable_prefix_as_cache_control(): + def _messages(marker: Mapping[str, Mapping[str, str]]) -> list[AllMessageValues]: + return cast( # cast-ok: the text part type does not declare the prompt_cache_breakpoint key + list[AllMessageValues], + [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Here is the full text of a complex legal agreement" * 400, **marker} + ], + }, + {"role": "user", "content": "what are the key terms?"}, + ], + ) + + marked_key: Final = PromptCachingCache.get_prompt_caching_cache_key( + _messages({"prompt_cache_breakpoint": {"mode": "explicit"}}), None + ) + + assert marked_key is not None + assert marked_key == PromptCachingCache.get_prompt_caching_cache_key( + _messages({"cache_control": {"type": "ephemeral"}}), None + ) + assert PromptCachingCache.get_prompt_caching_cache_key(_messages({}), None) is None + + @pytest.mark.asyncio async def test_same_cacheable_prefix_routes_to_same_deployment() -> None: from litellm import Router diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index abf0e055e73..9d8b09dd1c4 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -138,3 +138,110 @@ def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handl assert additional["llm_provider-request-id"] == "req_upstream_123" assert additional["x-ratelimit-remaining-requests"] == "41" assert additional["request-id"] == additional["llm_provider-request-id"] + + +_ADAPTIVE_OPUS: Final = "anthropic/claude-opus-4-7" + + +@pytest.mark.parametrize( + ("model", "expected"), + [("claude-sonnet-4-5", "anthropic"), ("gpt-4o", "openai"), ("not-a-known-model", None)], +) +def test_inferred_provider_is_the_one_the_python_handler_resolves(model: str, expected: str | None) -> None: + assert messages_host.inferred_provider(model, None) == expected + + +@pytest.mark.parametrize( + ("request_effort", "levels", "drop_params", "expected"), + [ + ("max", ["low", "high"], False, "high"), + ("high", ["low"], True, "low"), + ], +) +def test_body_params_round_effort_to_the_levels_the_deployment_declares( + request_effort: str, levels: list[str], drop_params: bool, expected: str +) -> None: + request: Final[dict[str, object]] = { + "model": _ADAPTIVE_OPUS, + "messages": [], + "max_tokens": 8, + "output_config": {"effort": request_effort}, + "drop_params": drop_params, + "model_info": {"id": "deployment-1", "reasoning_effort_levels": levels}, + } + + assert messages_host.body_params(request)["output_config"] == {"effort": expected} + + +def test_body_params_keep_the_requested_effort_without_deployment_levels() -> None: + request: Final[dict[str, object]] = { + "model": _ADAPTIVE_OPUS, + "messages": [], + "max_tokens": 8, + "output_config": {"effort": "max"}, + } + + assert messages_host.body_params(request)["output_config"] == {"effort": "max"} + + +def _marked(phase: str | None = None, *args: object) -> ValueError: + error: Final = ValueError(*args) + if phase is not None: + error.messages_timeout = phase # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + return error + + +def test_a_stalled_stream_body_reaches_the_caller_as_the_transport_timeout() -> None: + import httpx + + mapped: Final = messages_host.map_failure(_marked("Body"), {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert isinstance(mapped, httpx.ReadTimeout) + assert str(mapped) == "Timeout on reading data from socket" + + +def test_a_native_read_timeout_maps_to_the_public_timeout_with_python_text() -> None: + mapped: Final = messages_host.map_failure(_marked("Read"), {"model": _ADAPTIVE_OPUS, "timeout": 1.0}, "anthropic") + + assert isinstance(mapped, litellm.Timeout) + assert "Read timed out. Timeout passed=1.0" in str(mapped) + + +def test_an_upstream_status_maps_through_the_provider_error_class() -> None: + upstream: Final = ValueError(429, "slow down") + upstream.headers = [("retry-after", "3")] # pyright: ignore[reportAttributeAccessIssue] # set by the native host + + mapped: Final = messages_host.map_failure(upstream, {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert isinstance(mapped, litellm.RateLimitError) + assert "slow down" in str(mapped) + assert mapped.__context__ is upstream + + +def test_a_late_rejection_tells_the_logger_the_resolved_call_first() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + logger: Final = Logging( + model=_ADAPTIVE_OPUS, + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id="late", + function_id="late", + ) + rejected: Final = ValueError("budget_tokens must be a number") + rejected.messages_late_rejection = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + rejected.messages_request_error = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = messages_host.map_failure( + rejected, + {"model": _ADAPTIVE_OPUS, "messages": [], "max_tokens": 8, "litellm_logging_obj": logger}, + "anthropic", + ) + + assert isinstance(mapped, litellm.BadRequestError) + assert logger.model_call_details["custom_llm_provider"] == "anthropic" + assert logger.model_call_details["optional_params"] == {"max_tokens": 8} diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index a40f8a77105..1ba0188e76d 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -5,6 +5,7 @@ from pathlib import Path from types import MappingProxyType from typing import Final +import httpx import pytest from pydantic import TypeAdapter @@ -244,3 +245,56 @@ def test_the_rust_contract_matches_the_shim_signatures() -> None: contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} + + +def _messages_logger() -> Logging: + return Logging( + model="claude-sonnet-4-5", + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id="messages", + function_id="messages", + ) + + +def test_native_provider_response_is_recorded_for_logging_replay() -> None: + logger: Final = _messages_logger() + + legacy.record_response(logger, 200, (("request-id", "req_native"),), '{"id": "msg_native"}') + + recorded: Final[object] = logger.model_call_details["httpx_response"] + assert isinstance(recorded, httpx.Response) + assert recorded.status_code == 200 + assert recorded.content == b'{"id": "msg_native"}' + assert logger.model_call_details["response_headers"]["request-id"] == "req_native" + + +def test_native_provider_exchange_window_is_recorded_like_the_timed_post() -> None: + logger: Final = _messages_logger() + logger.model_call_details["litellm_params"] = {"litellm_metadata": {"user_api_key": "hashed"}} + + legacy.record_api_window(logger, 1_700_000_000.0, 1_700_000_000.25) + + assert logger.model_call_details["llm_api_duration_ms"] == 250.0 + metadata: Final = logger.model_call_details["litellm_params"]["litellm_metadata"] + assert metadata["llm_api_timing_windows"] == ((1_700_000_000.0, 1_700_000_000.25),) + + +def test_stream_head_is_stamped_in_place_with_the_call_hidden_params() -> None: + logger: Final = _messages_logger() + head: Final[dict[str, object]] = {"additional_headers": {"request-id": "req_native"}} + started: Final = datetime.datetime.now() + + legacy.stamp_stream_head( + head, + logger, + {"model": "anthropic/claude-sonnet-4-5", "litellm_call_id": "messages"}, + started, + started + datetime.timedelta(milliseconds=5), + ) + + assert head["additional_headers"] == {"request-id": "req_native", "llm_provider-request-id": "req_native"} + assert head["litellm_call_id"] == "messages" + assert "_response_ms" in head diff --git a/tests/unit/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py index 64559d63193..5ee89fa7794 100644 --- a/tests/unit/rust_bridge/messages/test_secrets.py +++ b/tests/unit/rust_bridge/messages/test_secrets.py @@ -11,7 +11,7 @@ import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.anthropic.messages.native.handler import anthropic_messages from litellm.rust_bridge.host import settings -from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES +from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES from litellm.rust_bridge.public_call import NativeCall from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE @@ -72,21 +72,13 @@ async def _python_messages() -> object: return await anthropic_messages(**_public_kwargs()) -async def _rust_messages() -> object: - route: Final = NATIVE_MESSAGES.load() - assert route is not None - return route(_native_request()) - - async def _rust_amessages() -> object: route: Final = NATIVE_AMESSAGES.load() assert route is not None return await route(_native_request()) -@pytest.fixture( - params=(_python_messages, _rust_messages, _rust_amessages), ids=("python-async", "rust-sync", "rust-async") -) +@pytest.fixture(params=(_python_messages, _rust_amessages), ids=("python-async", "rust-async")) def messages(request: pytest.FixtureRequest) -> Messages: return cast(Messages, request.param) diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 634e69910b9..e4c39bb40f9 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -334,6 +334,17 @@ def test_stream_cost_injection_needs_a_stream(monkeypatch: pytest.MonkeyPatch) - assert not messages_detectors.stream_cost_injection(call()) +def test_foundry_beta_row_reads_the_api_base_the_anthropic_handler_reads(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + assert messages_detectors.foundry_beta_row(call(api_base="https://team.services.ai.azure.com/anthropic")) + assert messages_detectors.foundry_beta_row(call(base_url="https://team.cognitiveservices.azure.com/anthropic")) + assert not messages_detectors.foundry_beta_row(call(api_base="https://api.anthropic.com")) + assert not messages_detectors.foundry_beta_row(call()) + monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://team.services.ai.azure.com/anthropic") + assert messages_detectors.foundry_beta_row(call()) + + def test_model_alias_map_applies_only_to_the_requested_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_alias_map", {"fast": "openai/gpt-5-mini"}) assert chat_detectors.model_alias_map(call(model="fast")) diff --git a/tests/unit/secret_managers/test_aws_secret_manager.py b/tests/unit/secret_managers/test_aws_secret_manager.py index d9922b43ec8..e16da7a0844 100644 --- a/tests/unit/secret_managers/test_aws_secret_manager.py +++ b/tests/unit/secret_managers/test_aws_secret_manager.py @@ -1,4 +1,6 @@ import base64 +import re +from typing import Final import pytest from botocore.stub import Stubber @@ -34,3 +36,15 @@ def test_decrypt_value_turns_only_boolean_literals_into_bools(monkeypatch, plain assert decrypted == expected assert type(decrypted) is type(expected) + + +def test_kms_v2_without_a_license_tells_the_user_to_set_one(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + monkeypatch.delenv("LITELLM_LICENSE", raising=False) + monkeypatch.delenv("LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE", raising=False) + + expected_message: Final = ( + "AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment." + ) + with pytest.raises(ValueError, match=f"^{re.escape(expected_message)}$"): + AWSKeyManagementService_V2() diff --git a/tests/unit/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py index c8fef95cf98..a9bc741f800 100644 --- a/tests/unit/test_anthropic_beta_headers_filtering.py +++ b/tests/unit/test_anthropic_beta_headers_filtering.py @@ -10,21 +10,32 @@ This test validates: import json import os -from typing import Dict, Final, List +from typing import Dict, Final, List, get_args from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx +from anthropic.types.anthropic_beta_param import AnthropicBetaParam import litellm +from litellm import anthropic_beta_headers_manager from litellm.anthropic_beta_headers_manager import ( GetAnthropicBetaHeadersConfig, filter_and_transform_beta_headers, get_beta_headers_config, + resolve_beta_provider, update_headers_with_filtered_beta, update_request_with_filtered_beta, ) +def _literal_strings(annotation: object) -> frozenset[str]: + args: Final = get_args(annotation) + if not args: + return frozenset({annotation}) if isinstance(annotation, str) else frozenset() + return frozenset[str]().union(*(_literal_strings(arg) for arg in args)) + + class TestAnthropicBetaHeadersFiltering: """Test beta header filtering and mapping for all providers.""" @@ -50,21 +61,24 @@ class TestAnthropicBetaHeadersFiltering: """Get all beta headers from the anthropic provider config.""" return list(self.config.get("anthropic", {}).keys()) + def get_header_column(self, provider: str) -> Dict[str, str | None]: + return {**self.config.get(provider, {}), **self.config.get(f"{provider}_header_only", {})} + def get_supported_headers(self, provider: str) -> List[str]: """Get headers with non-null values for a provider.""" - provider_config = self.config.get(provider, {}) + provider_config = self.get_header_column(provider) return [ header for header, value in provider_config.items() if value is not None ] def get_unsupported_headers(self, provider: str) -> List[str]: """Get headers with null values for a provider.""" - provider_config = self.config.get(provider, {}) + provider_config = self.get_header_column(provider) return [header for header, value in provider_config.items() if value is None] def get_mapped_headers(self, provider: str) -> Dict[str, str]: """Get mapping of input headers to provider-specific headers.""" - provider_config = self.config.get(provider, {}) + provider_config = self.get_header_column(provider) return { header: value for header, value in provider_config.items() @@ -407,7 +421,15 @@ class TestAnthropicBetaHeadersFiltering: test_case["expected"] in filtered ), f"Header '{test_case['input']}' should be mapped to '{test_case['expected']}' for {test_case['provider']}, but got: {filtered}" - @pytest.mark.parametrize("beta", ["compact-2026-09-04", "inline-tools-2026-09-15"]) + @pytest.mark.parametrize( + "beta", + [ + "compact-2026-09-04", + "inline-tools-2026-09-15", + "prompt-caching-scope-2026-01-05", + "redact-thinking-2026-02-12", + ], + ) def test_vertex_ai_header_only_beta_is_forwarded_as_the_header_and_dropped_from_the_body( self, beta: str ) -> None: @@ -509,6 +531,44 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["inline-tools-2026-09-15"] + @pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) + def test_task_budgets_and_computer_use_betas_reach_every_platform(self, provider: str) -> None: + """Live probes on 2026-10-28 found all three platforms accept task-budgets-2026-03-13. On 2026-10-31 Vertex + ran the computer tool only with the beta each model takes (computer-use-2025-01-24 on Claude Sonnet 4.5 and + Haiku 4.5, computer-use-2025-11-24 on Claude Opus 4.5 and the 4.6 models) and answered 400 to the other, so + both reach Vertex and the model decides, the same as calling Vertex directly.""" + sent: Final = frozenset(("computer-use-2025-01-24", "computer-use-2025-11-24", "task-budgets-2026-03-13")) + + headers: Final = update_headers_with_filtered_beta({"anthropic-beta": ",".join(sorted(sent))}, provider) + + assert frozenset(headers["anthropic-beta"].split(",")) == sent + + @pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) + def test_claude_code_betas_reach_the_probed_platforms(self, provider: str) -> None: + """Claude Code 2.1.299 sends these betas outside the anthropic SDK's list, and live probes on + 2026-10-28 found Anthropic, Vertex, and Azure AI Foundry accept each of them as the anthropic-beta header.""" + claude_code_betas: Final = frozenset( + ("afk-mode-2026-01-31", "redact-thinking-2026-02-12", "structured-outputs-2025-12-15") + ) + + headers: Final = update_headers_with_filtered_beta({"anthropic-beta": ",".join(claude_code_betas)}, provider) + + assert frozenset(headers["anthropic-beta"].split(",")) == claude_code_betas + + @pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) + def test_every_anthropic_sdk_beta_has_a_decision_for_the_probed_platforms(self, provider: str) -> None: + sdk_betas: Final = _literal_strings(AnthropicBetaParam) + provider_column: Final = GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()[provider] + + assert sdk_betas + assert sdk_betas - provider_column.keys() == frozenset() + + @pytest.mark.parametrize("provider", ["azure_ai", "vertex_ai"]) + def test_every_anthropic_beta_has_a_decision_for_vertex_ai_and_azure_ai(self, provider: str) -> None: + config: Final = GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() + + assert config["anthropic"].keys() - config[provider].keys() == set() + def test_null_value_headers_filtered(self): """Test that headers with null values are always filtered out.""" for provider in [ @@ -605,3 +665,113 @@ def test_local_flags_load_bundled_beta_headers_without_fetching(monkeypatch, fla config = get_beta_headers_config(url="https://example.test/anthropic_beta_headers_config.json") remote_get.assert_not_called() assert config == GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() + + +FOUNDRY_API_BASE: Final = "https://my-resource.services.ai.azure.com/anthropic" +DIRECT_API_BASE: Final = "https://api.anthropic.com" +DIRECT_ONLY_BETA: Final = "advisor-tool-2026-03-01" +PRESERVED_THINKING_BETAS: Final = ( + "thinking-binding-controls-2026-08-01", + "mid-conversation-system-clear-at-2026-08-21", + "mid-conversation-tool-changes-2026-07-01", + "mid-conversation-output-config-2026-07-01", + "thinking-display-updates-2026-08-18", +) + + +@pytest.fixture +def bundled_config(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") + anthropic_beta_headers_manager._BETA_HEADERS_CONFIG = None + + +def test_requests_without_betas_never_fetch_the_beta_table(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + monkeypatch.delenv("LITELLM_OFFLINE_MODE", raising=False) + anthropic_beta_headers_manager._BETA_HEADERS_CONFIG = None + with respx.mock(assert_all_called=False) as router: + remote_table: Final = router.get(litellm.anthropic_beta_headers_url).respond( + json=GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() + ) + for api_base in (None, DIRECT_API_BASE, FOUNDRY_API_BASE): + update_request_with_filtered_beta( + headers={"anthropic-version": "2023-06-01"}, + request_data={"max_tokens": 8}, + provider="anthropic", + api_base=api_base, + ) + for provider, deployment_api_base in ( + ("anthropic", DIRECT_API_BASE), + ("anthropic", FOUNDRY_API_BASE), + ("openai", FOUNDRY_API_BASE), + ): + resolve_beta_provider(provider, deployment_api_base) + assert remote_table.call_count == 0 + + +def test_foundry_deployment_headers_follow_the_azure_ai_row(bundled_config: None) -> None: + headers: Final = {"anthropic-beta": f"{DIRECT_ONLY_BETA},thinking-binding-controls-2026-08-01"} + result: Final[dict[str, str]] = update_headers_with_filtered_beta( + headers=headers, provider="anthropic", api_base=FOUNDRY_API_BASE + ) + assert result == {"anthropic-beta": "thinking-binding-controls-2026-08-01"} + + +def test_direct_deployment_headers_keep_every_anthropic_row_beta(bundled_config: None) -> None: + headers: Final = {"anthropic-beta": f"{DIRECT_ONLY_BETA},thinking-binding-controls-2026-08-01"} + result: Final[dict[str, str]] = update_headers_with_filtered_beta( + headers=headers, provider="anthropic", api_base=DIRECT_API_BASE + ) + assert result == {"anthropic-beta": f"{DIRECT_ONLY_BETA},thinking-binding-controls-2026-08-01"} + + +# Foundry's /anthropic/v1/messages answered 200 to each of these on claude-sonnet-5-5, 2026-10-31 +FOUNDRY_ACCEPTED_ANTHROPIC_ROW_BETAS: Final = ( + "compact-2026-01-12", + "fast-mode-2026-02-01", + "fine-grained-tool-streaming-2025-05-14", + "inline-tools-2026-09-15", + "oauth-2025-04-20", + "output-128k-2025-02-19", + "token-efficient-tools-2025-02-19", +) + + +def test_foundry_deployment_keeps_the_anthropic_row_betas_foundry_accepts(bundled_config: None) -> None: + result: Final[dict[str, str]] = update_headers_with_filtered_beta( + headers={"anthropic-beta": ",".join((DIRECT_ONLY_BETA, *FOUNDRY_ACCEPTED_ANTHROPIC_ROW_BETAS))}, + provider="anthropic", + api_base=FOUNDRY_API_BASE, + ) + assert sorted(result["anthropic-beta"].split(",")) == sorted(FOUNDRY_ACCEPTED_ANTHROPIC_ROW_BETAS) + + +def test_foundry_deployment_request_filters_header_and_body_betas_together(bundled_config: None) -> None: + headers: dict[str, str] + data: dict[str, object] + headers, data = update_request_with_filtered_beta( + headers={"anthropic-beta": DIRECT_ONLY_BETA}, + request_data={"anthropic_beta": [DIRECT_ONLY_BETA, "thinking-binding-controls-2026-08-01"], "max_tokens": 8}, + provider="anthropic", + api_base=FOUNDRY_API_BASE, + ) + assert headers == {} + assert data == {"anthropic_beta": ["thinking-binding-controls-2026-08-01"], "max_tokens": 8} + + +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) +def test_live_verified_platforms_forward_every_preserved_thinking_beta(bundled_config: None, provider: str) -> None: + result: Final[dict[str, str]] = update_headers_with_filtered_beta( + headers={"anthropic-beta": ",".join(PRESERVED_THINKING_BETAS)}, provider=provider + ) + assert sorted(result["anthropic-beta"].split(",")) == sorted(PRESERVED_THINKING_BETAS), result + + +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai", "bedrock", "bedrock_converse"]) +def test_each_row_forwards_exactly_its_verified_preserved_thinking_cells(bundled_config: None, provider: str) -> None: + row: Final[dict[str, str | None]] = GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()[provider] + forwarded: Final = tuple(beta for beta in PRESERVED_THINKING_BETAS if row.get(beta) is not None) + result: Final[dict[str, str]] = update_headers_with_filtered_beta( + headers={"anthropic-beta": ",".join(PRESERVED_THINKING_BETAS)}, provider=provider + ) + assert sorted(result["anthropic-beta"].split(",")) == sorted(forwarded), (provider, result) diff --git a/tests/unit/test_check_openapi_docs.py b/tests/unit/test_check_openapi_docs.py index d7122541416..5d8803c4118 100644 --- a/tests/unit/test_check_openapi_docs.py +++ b/tests/unit/test_check_openapi_docs.py @@ -1,7 +1,10 @@ """Tests for scripts/check_openapi_docs.py, the scan behind scripts/openapi_docs_gate.py.""" +from collections.abc import Mapping from typing import Final +import pytest + import check_openapi_docs as checker _BODY: Final = {"content": {"application/json": {"schema": {"type": "object"}}}} @@ -9,8 +12,23 @@ _RESPONSES: Final = {"200": {"content": {"application/json": {"schema": {"$ref": _DOCUMENTED: Final = {"description": "Does a thing", "requestBody": _BODY, "responses": _RESPONSES} -def test_a_fully_documented_post_has_no_findings() -> None: - assert checker.violations({"paths": {"/v1/thing": {"post": _DOCUMENTED}}}) == () +def _post_with_media_types(request_media_type: str, response_media_type: str) -> Mapping[str, object]: + return { + **_DOCUMENTED, + "requestBody": {"content": {request_media_type: {"schema": {"type": "string"}}}}, + "responses": {"200": {"content": {response_media_type: {"schema": {"type": "string"}}}}}, + } + + +@pytest.mark.parametrize("media_type", ("application/json", "application/sdp")) +def test_a_fully_documented_post_has_no_findings(media_type: str) -> None: + operation: Final = _post_with_media_types(media_type, media_type) + assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == () + + +def test_a_schema_under_another_media_type_still_counts_as_undocumented() -> None: + operation: Final = _post_with_media_types("multipart/form-data", "audio/wav") + assert {v.code for v in checker.violations({"paths": {"/v1/thing": {"post": operation}}})} == {"OA001", "OA003"} def test_each_missing_piece_is_its_own_rule() -> None: diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index cca138e0bb5..e0ac07b3870 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -1451,6 +1451,47 @@ def test_transcription_usage_cost_returns_zero_for_unknown_type(): assert _transcription_usage_cost({}, {}) == 0.0 +@pytest.mark.parametrize("malformed_seconds", [-1.0, float("nan"), float("inf"), 10**1000, True]) +def test_transcription_ignores_invalid_provider_duration(malformed_seconds: float | int | bool) -> None: + from litellm.cost_calculator import _get_transcription_usage_duration + + response: Final = SimpleNamespace(usage={"type": "duration", "seconds": malformed_seconds}) + + assert _get_transcription_usage_duration(response) is None + + +def test_realtime_transcription_prices_each_clip_at_the_transcription_model_set_when_it_ran(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + from litellm.cost_calculator import handle_realtime_transcription_cost_calculation + + realtime_model = "gpt-realtime-2.1-mini" + first_asr, second_asr = "gpt-4o-transcribe", "gpt-transcribe" + first_rate = litellm.model_cost[first_asr]["input_cost_per_second"] + second_rate = litellm.model_cost[second_asr]["input_cost_per_second"] + assert first_rate != second_rate and first_rate > 0 and second_rate > 0 + + def session_frame(frame_type: str, asr_model: str | None) -> dict[str, object]: + transcription = None if asr_model is None else {"model": asr_model} + return {"type": frame_type, "session": {"model": realtime_model, "audio": {"input": {"transcription": transcription}}}} + + clip = {"type": "conversation.item.input_audio_transcription.completed", "usage": {"type": "duration", "seconds": 5}} + results: OpenAIRealtimeStreamList = [ + session_frame("session.created", None), + session_frame("session.updated", first_asr), + clip, + session_frame("session.updated", second_asr), + clip, + ] + + cost = handle_realtime_transcription_cost_calculation( + results=results, custom_llm_provider="openai", litellm_model_name=realtime_model + ) + + assert cost == pytest.approx(5 * first_rate + 5 * second_rate, rel=1e-9) + + def test_get_transcription_model_falls_back_to_session_model(monkeypatch): """session.model is used when transcription-specific model fields are absent.""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") @@ -7442,6 +7483,39 @@ def test_transcription_cost_falls_back_to_provider_billed_seconds( assert cost == pytest.approx(0.01 * expected_seconds) +@pytest.mark.parametrize( + ("provider", "measured_seconds", "expected_seconds"), + [("openai", 1.0, 1.0), ("azure", 1.0, 1.0), ("azure", None, 7.0)], + ids=["openai_measured_upload_wins", "azure_measured_upload_wins", "unmeasured_upload_bills_provider_seconds"], +) +def test_transcription_priced_by_provider_duration_still_bills_the_measured_upload( + monkeypatch: pytest.MonkeyPatch, provider: str, measured_seconds: float | None, expected_seconds: float +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "gpt-fake-transcribe", + { + "input_cost_per_second": 0.01, + "litellm_provider": provider, + "mode": "audio_transcription", + "provider_specific_entry": {"transcription_usage_duration": 1}, + }, + ) + measured: Final = {} if measured_seconds is None else {"audio_transcription_duration": measured_seconds} + transcription: Final = convert_to_model_response_object( + response_object={"text": "What is the weather in Paris?", "usage": {"type": "duration", "seconds": 7}}, + model_response_object=TranscriptionResponse(), + hidden_params={"model": "gpt-fake-transcribe", "custom_llm_provider": provider, **measured}, + response_type="audio_transcription", + ) + + cost: Final = completion_cost( + completion_response=transcription, model=f"{provider}/gpt-fake-transcribe", call_type="atranscription" + ) + + assert cost == pytest.approx(0.01 * expected_seconds) + + def test_custom_priced_transcription_deployment_bills_provider_seconds(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setitem( litellm.model_cost, @@ -7466,6 +7540,77 @@ def test_custom_priced_transcription_deployment_bills_provider_seconds(monkeypat assert cost == pytest.approx(0.000075 * 2) +@pytest.mark.parametrize( + "deployment_id,deployment_rates,expected_cost", + [ + ("per-second-only", {"input_cost_per_second": 0.0001}, 30 * 0.0001), + ("zero-token-rate", {"input_cost_per_second": 0.0001, "input_cost_per_token": 0.0}, 0.0), + ("token-rate", {"input_cost_per_second": 0.0001, "input_cost_per_token": 1e-6}, 500 * 1e-6), + ], +) +def test_transcription_deployment_bills_seconds_only_without_a_declared_token_rate( + monkeypatch: pytest.MonkeyPatch, + deployment_id: str, + deployment_rates: Mapping[str, float], + expected_cost: float, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + deployment_id, + {**deployment_rates, "litellm_provider": "openai", "mode": "audio_transcription"}, + ) + transcription: Final = convert_to_model_response_object( + response_object={ + "text": "What is the weather in Paris?", + "usage": { + "type": "tokens", + "input_tokens": 500, + "output_tokens": 100, + "total_tokens": 600, + "input_token_details": {"audio_tokens": 500, "text_tokens": 0}, + }, + }, + model_response_object=TranscriptionResponse(), + hidden_params={ + "model": "gpt-4o-transcribe", + "custom_llm_provider": "openai", + "audio_transcription_duration": 30.0, + }, + response_type="audio_transcription", + ) + + cost: Final = completion_cost( + completion_response=transcription, + model="openai/gpt-4o-transcribe", + call_type="atranscription", + custom_pricing=True, + router_model_id=deployment_id, + ) + + assert cost == pytest.approx(expected_cost) + + +@pytest.mark.parametrize("model", ["gpt-4o-transcribe", "gpt-4o-mini-transcribe"]) +def test_azure_transcription_without_usage_bills_only_tokens_while_openai_bills_the_measured_audio( + _local_model_cost_map: None, model: str +) -> None: + def cost_of_ten_seconds_without_usage(provider: str) -> float: + transcript: Final = TranscriptionResponse(text="Good") + transcript.set_audio_transcription_duration(10.0) + return completion_cost( + completion_response=transcript, + model=f"{provider}/{model}", + custom_llm_provider=provider, + call_type="atranscription", + ) + + assert cost_of_ten_seconds_without_usage("openai") > 0 + assert cost_of_ten_seconds_without_usage("azure") == 0.0, ( + "Azure lists only token prices for this model, https://azure.microsoft.com/en-us/pricing/details/azure-openai/ " + "(per the 2026-10-28 review)" + ) + + def test_gpt_image_2_azure_cost_tracking(): azure_image_generation_response: Final = { "created": 1758585600, @@ -10409,6 +10554,301 @@ def test_responses_ws_logging_object_carries_the_completed_responses_id_tier_and assert ResponsesWebSocketTokenUsageProcessor.create_logging_object([{"type": "error", "error": {}}]).id is None +@pytest.mark.parametrize("provider", ("openai", "azure")) +@pytest.mark.parametrize("model", ("gpt-realtime-2.1", "gpt-realtime-2.1-mini")) +def test_realtime_cached_multimodal_token_cost(_local_model_cost_map, provider: str, model: str): + model_name: Final = f"azure/{model}" if provider == "azure" else model + rates: Final = litellm.model_cost[model_name] + events: Final[OpenAIRealtimeStreamList] = [ + {"type": "session.created", "session": {"model": model}}, + { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 1000, + "output_tokens": 300, + "total_tokens": 1300, + "input_token_details": { + "text_tokens": 400, + "audio_tokens": 400, + "image_tokens": 200, + "cached_tokens": 300, + "cached_tokens_details": {"text_tokens": 100, "audio_tokens": 150, "image_tokens": 50}, + }, + "output_token_details": {"text_tokens": 100, "audio_tokens": 100, "reasoning_tokens": 100}, + } + }, + }, + ] + combined: Final = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(events) + actual: Final = handle_realtime_stream_cost_calculation( + results=events, + combined_usage_object=combined, + custom_llm_provider=provider, + litellm_model_name=model_name, + ) + expected: Final = ( + 300 * rates["input_cost_per_token"] + + 250 * rates["input_cost_per_audio_token"] + + 150 * rates["input_cost_per_image_token"] + + 100 * rates["cache_read_input_token_cost"] + + 150 * rates["cache_read_input_audio_token_cost"] + + 50 * rates["cache_read_input_image_token_cost"] + + 200 * rates["output_cost_per_token"] + + 100 * rates["output_cost_per_audio_token"] + ) + + assert actual == pytest.approx(expected) + + +def test_realtime_translation_duration_cost(_local_model_cost_map): + from litellm.cost_calculator import handle_realtime_translation_cost_calculation + + model: Final = "gpt-realtime-translate" + events: Final[OpenAIRealtimeStreamList] = [ + {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 2.0, "output_seconds": 1.5}} + ] + actual: Final = handle_realtime_translation_cost_calculation( + results=events, + custom_llm_provider="openai", + litellm_model_name=model, + ) + + assert actual == pytest.approx(2 * litellm.model_cost[model]["input_cost_per_second"]) + + +def test_realtime_translation_bills_session_audio_once_on_openai_and_azure(_local_model_cost_map) -> None: + from litellm.cost_calculator import handle_realtime_translation_cost_calculation + + one_minute_session: Final[OpenAIRealtimeStreamList] = [ + {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 60.0, "output_seconds": 60.0}} + ] + openai_cost: Final = handle_realtime_translation_cost_calculation( + results=one_minute_session, custom_llm_provider="openai", litellm_model_name="gpt-realtime-translate" + ) + azure_cost: Final = handle_realtime_translation_cost_calculation( + results=one_minute_session, custom_llm_provider="azure", litellm_model_name="gpt-realtime-translate" + ) + + assert openai_cost == pytest.approx(60 * litellm.model_cost["gpt-realtime-translate"]["input_cost_per_second"]) + assert openai_cost == pytest.approx(azure_cost) + + +@pytest.mark.parametrize("malformed_duration", [-1.0, float("nan"), float("inf"), 10**1000, True]) +def test_realtime_translation_duration_cost_ignores_malformed_provider_usage( + _local_model_cost_map, malformed_duration: float | int | bool +) -> None: + from litellm.cost_calculator import handle_realtime_translation_cost_calculation + + model: Final = "translation-deployment-priced-both-directions" + litellm.register_model( + model_cost={ + model: { + "litellm_provider": "openai", + "mode": "realtime", + "input_cost_per_second": 0.001, + "output_cost_per_second": 0.002, + } + } + ) + input_invalid_events: Final[OpenAIRealtimeStreamList] = [ + { + "type": "session.closed", + "usage": {"type": "duration", "input_seconds": malformed_duration, "output_seconds": 1.5}, + } + ] + output_invalid_events: Final[OpenAIRealtimeStreamList] = [ + { + "type": "session.closed", + "usage": {"type": "duration", "input_seconds": 2.0, "output_seconds": malformed_duration}, + } + ] + + input_invalid_cost: Final = handle_realtime_translation_cost_calculation( + results=input_invalid_events, custom_llm_provider="openai", litellm_model_name=model + ) + output_invalid_cost: Final = handle_realtime_translation_cost_calculation( + results=output_invalid_events, custom_llm_provider="openai", litellm_model_name=model + ) + + assert input_invalid_cost == pytest.approx(1.5 * litellm.model_cost[model]["output_cost_per_second"]) + assert output_invalid_cost == pytest.approx(2.0 * litellm.model_cost[model]["input_cost_per_second"]) + + +def _translation_session_cost( + deployment_id: str, + session_model: str | None = "gpt-realtime-translate", + model: str = "gpt-realtime-translate", + base_model: str | None = None, +) -> float: + session_created: Final[OpenAIRealtimeStreamList] = ( + [{"type": "session.created", "session": {"model": session_model}}] if session_model is not None else [] + ) + one_minute_session: Final = LiteLLMRealtimeStreamLoggingObject( + usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), + results=[ + *session_created, + {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 60.0, "output_seconds": 60.0}}, + ], + ) + return completion_cost( + completion_response=one_minute_session, + model=model, + call_type=CallTypes.arealtime.value, + custom_llm_provider="openai", + custom_pricing=True, + router_model_id=deployment_id, + base_model=base_model, + ) + + +@pytest.mark.parametrize("session_model", ["gpt-realtime-translate", None]) +@pytest.mark.parametrize( + "deployment_model,base_model", + [("gpt-realtime-translate", None), ("my-translator", "gpt-realtime-translate")], +) +@pytest.mark.parametrize( + "deployment_id,deployment_rates,bills_public_input_rate", + [ + ("translate-output-priced", {"output_cost_per_second": 0.001}, True), + ("translate-free-input", {"input_cost_per_second": 0.0, "output_cost_per_second": 0.001}, False), + ], +) +def test_realtime_translation_deployment_per_second_rates_fall_back_to_public_rates_per_key( + _local_model_cost_map: None, + deployment_id: str, + deployment_rates: Mapping[str, float], + bills_public_input_rate: bool, + deployment_model: str, + base_model: str | None, + session_model: str | None, +) -> None: + litellm.Router( + model_list=[ + { + "model_name": "translate", + "litellm_params": { + "model": f"openai/{deployment_model}", + "api_key": "sk-test", + **deployment_rates, + }, + "model_info": {"id": deployment_id, **({"base_model": base_model} if base_model else {})}, + } + ] + ) + public_input_rate: Final = litellm.model_cost["gpt-realtime-translate"]["input_cost_per_second"] + + assert public_input_rate > 0 + assert _translation_session_cost( + deployment_id, session_model=session_model, model=deployment_model, base_model=base_model + ) == pytest.approx(60 * 0.001 + (60 * public_input_rate if bills_public_input_rate else 0.0)) + + +def test_realtime_translation_custom_deployment_with_one_rate_bills_the_session_model_for_the_other( + _local_model_cost_map: None, +) -> None: + litellm.Router( + model_list=[ + { + "model_name": "translate", + "litellm_params": { + "model": "openai/my-translator", + "api_key": "sk-test", + "output_cost_per_second": 0.001, + }, + "model_info": {"id": "custom-translate-output-priced"}, + } + ] + ) + public_input_rate: Final = litellm.model_cost["gpt-realtime-translate"]["input_cost_per_second"] + + assert public_input_rate > 0 + assert _translation_session_cost( + "custom-translate-output-priced", model="my-translator" + ) == pytest.approx(60 * 0.001 + 60 * public_input_rate) + + +def test_realtime_translation_rate_ignores_a_session_model_the_deployment_does_not_run( + _local_model_cost_map: None, +) -> None: + litellm.register_model( + model_cost={ + "translation-session-echo": { + "litellm_provider": "openai", + "mode": "realtime", + "input_cost_per_second": 0.5, + } + } + ) + litellm.Router( + model_list=[ + { + "model_name": "translate", + "litellm_params": { + "model": "openai/gpt-realtime-translate", + "api_key": "sk-test", + "output_cost_per_second": 0.001, + }, + "model_info": {"id": "translate-output-priced"}, + } + ] + ) + + assert _translation_session_cost( + "translate-output-priced", session_model="translation-session-echo" + ) == pytest.approx(_translation_session_cost("translate-output-priced", session_model=None)) + + +def test_realtime_translation_deployment_without_per_second_rates_bills_public_rates(_local_model_cost_map) -> None: + litellm.Router( + model_list=[ + { + "model_name": "translate", + "litellm_params": { + "model": "openai/gpt-realtime-translate", + "api_key": "sk-test", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"id": "translate-token-priced"}, + } + ] + ) + public_entry: Final = litellm.model_cost["gpt-realtime-translate"] + public_cost: Final = 60 * (public_entry.get("input_cost_per_second") or 0.0) + 60 * ( + public_entry.get("output_cost_per_second") or 0.0 + ) + + assert public_cost > 0 + assert _translation_session_cost("translate-token-priced") == pytest.approx(public_cost) + + +def test_realtime_translation_spend_usage_carries_the_billed_input_seconds(_local_model_cost_map) -> None: + logging_obj: Final = Logging( + model="gpt-realtime-translate", + messages=[], + stream=False, + call_type=CallTypes.arealtime.value, + start_time=datetime.datetime.now(), + litellm_call_id="translation-spend-usage", + function_id="translation-spend-usage", + ) + logging_obj.update_environment_variables( + model="gpt-realtime-translate", user=None, optional_params={}, litellm_params={}, custom_llm_provider="openai" + ) + events: Final[OpenAIRealtimeStreamList] = [ + {"type": "session.closed", "usage": {"type": "duration", "input_seconds": 30.0, "output_seconds": 42.5}} + ] + + logging_obj.success_handler(result=events, start_time=datetime.datetime.now(), end_time=datetime.datetime.now()) + + payload: Final = logging_obj.model_call_details["standard_logging_object"] + assert payload["metadata"]["usage_object"]["seconds"] == pytest.approx(30.0) + assert payload["response_cost"] == pytest.approx( + 30.0 * litellm.model_cost["gpt-realtime-translate"]["input_cost_per_second"] + ) + + @pytest.mark.parametrize("profile_prefix", ("us", "global")) @pytest.mark.parametrize( "model_name", diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 85864b2024e..47629aab326 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -553,6 +553,63 @@ def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx assert sent_tool["function"]["name"] == "write_file" +_ALLOWED_TOOLS_TOOL_CHOICE: Final = { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, +} +_TWO_TOOLS: Final = [ + {"type": "function", "function": {"name": "get_time", "parameters": {"type": "object", "properties": {}}}}, + {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}}, +] + + +@pytest.mark.parametrize("model", ["openai/gpt-5.5", "azure/my-gpt-deployment"]) +def test_completion_forwards_allowed_tools_tool_choice_to_openai_compatible_chat( + respx_mock: respx.MockRouter, openai_api_response: dict[str, object], model: str +): + mock_route: Final = respx_mock.post(url__regex=r"http://localhost:12346/.*chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) + ) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "What time is it?"}], + tools=_TWO_TOOLS, + tool_choice=_ALLOWED_TOOLS_TOOL_CHOICE, + api_base="http://localhost:12346/v1", + api_key="fake_api_key", + api_version="2025-04-01-preview", + ) + + assert mock_route.called + sent: Final = json.loads(respx_mock.calls[0].request.content) + assert sent["tool_choice"] == _ALLOWED_TOOLS_TOOL_CHOICE + assert [tool["function"]["name"] for tool in sent["tools"]] == ["get_time", "get_weather"] + + +@pytest.mark.parametrize("drop_params", [False, True]) +@pytest.mark.parametrize("model", ["anthropic/claude-sonnet-4-5", "gemini/gemini-2.5-flash", "bedrock/converse/x"]) +def test_completion_allowed_tools_tool_choice_is_a_400_where_the_provider_cannot_enforce_it( + respx_mock: respx.MockRouter, model: str, drop_params: bool +): + with pytest.raises( + litellm.BadRequestError, match=r"allowed_tools.*only supported for azure, litellm_proxy, openai" + ) as exc: + litellm.completion( + model=model, + messages=[{"role": "user", "content": "What time is it?"}], + tools=_TWO_TOOLS, + tool_choice=_ALLOWED_TOOLS_TOOL_CHOICE, + drop_params=drop_params, + api_key="fake_api_key", + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + assert exc.value.status_code == 400 + assert not respx_mock.calls + + _CLAUDE_TOOL_SEARCH_CALL: Final = { "id": "srvtoolu_01search", "type": "function", @@ -6127,6 +6184,79 @@ async def test_groq_atranscription_honors_base_url_alias(respx_mock: respx.MockR assert response.text == "hello" +def test_transcription_keeps_the_prompt_passed_by_position(respx_mock: respx.MockRouter): + route: Final = respx_mock.post(f"{GROQ_INTERNAL_BASE}/audio/transcriptions").mock( + return_value=httpx.Response(200, json={"text": "hello"}) + ) + + litellm.transcription( + "groq/whisper-large-v3", GROQ_WAV_FILE, "en", "Four score", base_url=GROQ_INTERNAL_BASE, api_key="fake-key" + ) + + sent: Final = route.calls.last.request.content + assert b'name="language"\r\n\r\nen\r\n' in sent + assert b'name="prompt"\r\n\r\nFour score\r\n' in sent + + +@pytest.mark.parametrize( + "provider_body,expected_upload_measurements", + [({"text": "hello", "duration": 2.5}, 0), ({"text": "hello"}, 1)], + ids=["provider_reports_duration", "provider_omits_duration"], +) +def test_transcription_measures_the_upload_only_when_the_provider_omits_the_duration( + respx_mock: respx.MockRouter, provider_body: Mapping[str, object], expected_upload_measurements: int +): + respx_mock.post(f"{GROQ_INTERNAL_BASE}/audio/transcriptions").mock( + return_value=httpx.Response(200, json=provider_body) + ) + + with patch.object( + litellm_main, "calculate_request_duration", wraps=litellm_main.calculate_request_duration + ) as measure_upload: + response: Final = litellm.transcription( + model="groq/whisper-large-v3", file=GROQ_WAV_FILE, base_url=GROQ_INTERNAL_BASE, api_key="fake-key" + ) + + assert response.text == "hello" + assert measure_upload.call_count == expected_upload_measurements + + +@pytest.mark.parametrize("model_name", ("whisper-1", "gpt-4o-transcribe", "gpt-transcribe")) +def test_openai_compatible_transcription_with_an_openai_model_name_reaches_the_server( + respx_mock: respx.MockRouter, model_name: str +): + route: Final = respx_mock.post("https://vllm.internal/v1/audio/transcriptions").mock( + return_value=httpx.Response(200, json={"text": "served by vllm"}) + ) + + response: Final = litellm.transcription( + model=f"hosted_vllm/{model_name}", + file=GROQ_WAV_FILE, + api_base="https://vllm.internal/v1", + api_key="fake-key", + response_format="verbose_json", + ) + + assert route.called + assert response.text == "served by vllm" + + +def test_openai_gpt_transcribe_still_rejects_non_json_response_formats(respx_mock: respx.MockRouter): + route: Final = respx_mock.post("https://api.openai.com/v1/audio/transcriptions").mock( + return_value=httpx.Response(200, json={"text": "should not be sent"}) + ) + + with pytest.raises(litellm.UnsupportedParamsError, match="only supports response_format 'json' or 'text'"): + litellm.transcription( + model="openai/gpt-transcribe", + file=GROQ_WAV_FILE, + api_key="fake-key", + response_format="srt", + ) + + assert not route.called + + def test_groq_speech_honors_base_url_alias(respx_mock: respx.MockRouter): route: Final = respx_mock.post(f"{GROQ_INTERNAL_BASE}/audio/speech").mock( return_value=httpx.Response(200, content=b"mp3-bytes") diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index e4a2beb3a86..e9e7f17f80d 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -388,6 +388,34 @@ def test_vertex_gemini_flash_rows_carry_the_documented_retirement_date(prices: d assert prices[f"{prefix}{model}"]["deprecation_date"] == VERTEX_GEMINI_FLASH_RETIREMENT_DATES[model] +# https://developers.openai.com/api/docs/deprecations, the 2026-10-14 GPT-4 series notice +OPENAI_GPT_4_SERIES_SNAPSHOTS: Final = MappingProxyType( + { + "gpt-4o": ("gpt-4o-2024-08-06", "gpt-4o-2024-11-20"), + "gpt-4o-mini": ("gpt-4o-mini-2024-07-18",), + "gpt-4.1": ("gpt-4.1-2025-04-14",), + "gpt-4.1-mini": ("gpt-4.1-mini-2025-04-14",), + } +) + + +class _ShutdownRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + deprecation_date: str | None = None + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +@pytest.mark.parametrize("alias", sorted(OPENAI_GPT_4_SERIES_SNAPSHOTS)) +def test_openai_gpt_4_series_aliases_retire_with_their_snapshots(path: Path, alias: str) -> None: + prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + snapshot_dates: Final = frozenset( + _ShutdownRow.model_validate(prices[snapshot]).deprecation_date + for snapshot in OPENAI_GPT_4_SERIES_SNAPSHOTS[alias] + ) + assert snapshot_dates == {_ShutdownRow.model_validate(prices[alias]).deprecation_date} + + def _sampling_flags_across_provider_twins(prices: Mapping[str, Mapping[str, object]], key: str) -> frozenset[object]: bare: Final = key.split("/")[-1] twins: Final = tuple(twin for twin in (bare, f"vertex_ai/{bare}", f"gemini/{bare}") if twin in prices) @@ -838,6 +866,67 @@ def test_unregistered_provider_guard_flags_only_labels_nobody_registered(): ] +@pytest.mark.parametrize("model", ("gpt-realtime-2.1", "gpt-realtime-2.1-mini")) +def test_realtime_family_cache_image_rate_tracks_azure(prices: dict, model: str): + openai: Final = prices[model] + azure: Final = prices[f"azure/{model}"] + + assert openai["cache_read_input_image_token_cost"] > 0 + assert azure["cache_read_input_image_token_cost"] == openai["cache_read_input_image_token_cost"] + assert azure["input_cost_per_image_token"] >= azure["cache_read_input_image_token_cost"] + + +@pytest.mark.parametrize( + "model,mode", + ( + ("gpt-realtime-translate", "realtime"), + ("gpt-live-transcribe", "audio_transcription"), + ("gpt-transcribe", "audio_transcription"), + ), +) +def test_azure_realtime_specialized_models_follow_openai_modes(prices: dict, model: str, mode: str): + openai: Final = prices[model] + azure: Final = prices[f"azure/{model}"] + + assert openai["mode"] == azure["mode"] == mode + assert azure["supports_audio_input"] is True + assert set(openai.get("supported_endpoints", ())) <= set(azure.get("supported_endpoints", ())) + + +_DATED_SNAPSHOT: Final = re.compile(r"(?P.+)-\d{4}-\d{2}-\d{2}") + + +def _azure_realtime_snapshots_with_an_alias(prices: Mapping[str, dict]) -> tuple[tuple[str, str], ...]: + dated: Final = ( + (name, _DATED_SNAPSHOT.fullmatch(name)) + for name, entry in prices.items() + if name.startswith("azure/") and "/v1/realtime" in entry.get("supported_endpoints", ()) + ) + return tuple((name, match["alias"]) for name, match in dated if match is not None and match["alias"] in prices) + + +def _realtime_ga_only(entry: dict) -> object: + return (entry.get("provider_specific_entry") or {}).get("realtime_ga_only") + + +def test_dated_azure_realtime_models_select_the_same_protocol_as_their_alias(prices: dict): + snapshots: Final = _azure_realtime_snapshots_with_an_alias(prices) + disagreeing: Final = tuple( + (snapshot, alias) + for snapshot, alias in snapshots + if _realtime_ga_only(prices[snapshot]) != _realtime_ga_only(prices[alias]) + ) + + assert snapshots + assert disagreeing == () + + +def test_model_prices_backup_is_synchronized(prices: dict): + backup: Final = json.loads(BACKUP_PRICES_PATH.read_text()) + + assert backup == prices + + def _search_free_requests_rows(path: Path) -> Mapping[str, dict[str, JsonValue]]: prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) return MappingProxyType( @@ -845,6 +934,26 @@ def _search_free_requests_rows(path: Path) -> Mapping[str, dict[str, JsonValue]] ) +def _supports_continuation_rows(path: Path) -> frozenset[str]: + prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + return frozenset( + key for key, row in prices.items() if isinstance(row, dict) and row.get("supports_continuation") is True + ) + + +def test_gemini_4_argon_google_rows_carry_supports_continuation_in_both_cost_map_copies() -> None: + main_rows: Final = _supports_continuation_rows(PRICES_PATH) + # long decode models: https://ai.google.dev/gemini-api/docs/generate-content/long-decode and + # https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/guides/gemini-4-argon (2026-11-02) + long_decode_rows: Final = { + "gemini/gemini-4-argon", + "vertex_ai/gemini-4-argon", + "vertex_ai/gemini-4-argon-cyber-permissive", + } + assert long_decode_rows <= main_rows, sorted(main_rows) + assert _supports_continuation_rows(BACKUP_PRICES_PATH) == main_rows + + @pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) def test_search_free_requests_rows_are_grounding_rows_with_one_allowance_per_group(path: Path) -> None: rows = _search_free_requests_rows(path) diff --git a/tests/unit/test_responses_id_security.py b/tests/unit/test_responses_id_security.py index 7e74bfd0f17..789fdbd453e 100644 --- a/tests/unit/test_responses_id_security.py +++ b/tests/unit/test_responses_id_security.py @@ -4,23 +4,43 @@ Tests for ResponsesIDSecurity hook. Tests the security hook that prevents user B from seeing response from user A. """ +import datetime +from collections.abc import Callable, Mapping +from typing import Final from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException +from pydantic import TypeAdapter +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.hooks.responses_id_security import ( ResponsesIDSecurity, _streams_response_ids, ) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( GenericEvent, + ResponseAPIUsage, ResponseCompletedEvent, ResponseCreatedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, ) -from litellm.types.utils import SpecialEnums +from litellm.types.responses.main import DeleteResponseResult +from litellm.types.utils import ( + ADDRESSED_COMPARISON_RESPONSE_ID_FIELD, + CallTypesLiteral, + Choices, + Message, + ModelResponse, + SpecialEnums, + Usage, +) @pytest.fixture @@ -55,12 +75,12 @@ class TestIsEncryptedResponseId: import litellm.proxy.hooks.responses_id_security as responses_module with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: - mock_decrypt.return_value = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_123;user_id:user-456" - - result = responses_id_security._is_encrypted_response_id( - "resp_encrypted_value" + mock_decrypt.return_value = ( + f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_123;user_id:user-456" ) + result = responses_id_security._is_encrypted_response_id("resp_encrypted_value") + assert result is True mock_decrypt.assert_called_once() @@ -88,9 +108,7 @@ class TestDecryptResponseId: with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: mock_decrypt.return_value = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_original_123;user_id:user-456;team_id:team-789" - original_id, user_id, team_id = responses_id_security._decrypt_response_id( - "resp_encrypted_value" - ) + original_id, user_id, team_id = responses_id_security._decrypt_response_id("resp_encrypted_value") assert original_id == "resp_original_123" assert user_id == "user-456" @@ -104,9 +122,7 @@ class TestDecryptResponseId: with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: mock_decrypt.return_value = None - original_id, user_id, team_id = responses_id_security._decrypt_response_id( - "resp_plain_value" - ) + original_id, user_id, team_id = responses_id_security._decrypt_response_id("resp_plain_value") assert original_id == "resp_plain_value" assert user_id is None @@ -116,9 +132,7 @@ class TestDecryptResponseId: class TestCheckUserAccessToResponseId: """Test check_user_access_to_response_id function""" - def test_check_user_access_same_user( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_same_user(self, responses_id_security, mock_user_api_key_dict): """Test that same user has access to their response ID""" result = responses_id_security.check_user_access_to_response_id( response_id_user_id="test-user-123", @@ -128,9 +142,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_different_user_raises_exception( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_different_user_raises_exception(self, responses_id_security, mock_user_api_key_dict): """Test that different user is denied access to response ID""" with patch("litellm.proxy.proxy_server.general_settings", {}): with pytest.raises(HTTPException) as exc_info: @@ -143,9 +155,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "Forbidden" in exc_info.value.detail - def test_check_user_access_different_team_raises_exception( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_different_team_raises_exception(self, responses_id_security, mock_user_api_key_dict): """Test that different team is denied access to response ID""" with patch("litellm.proxy.proxy_server.general_settings", {}): with pytest.raises(HTTPException) as exc_info: @@ -158,9 +168,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "Forbidden" in exc_info.value.detail - def test_check_user_access_team_a_to_team_b_without_user_id( - self, responses_id_security - ): + def test_check_user_access_team_a_to_team_b_without_user_id(self, responses_id_security): """Test that key from team A (without user_id) cannot access response from team B (without user_id)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -179,9 +187,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "team" in exc_info.value.detail.lower() - def test_check_user_access_team_a_to_team_b_with_user_id( - self, responses_id_security - ): + def test_check_user_access_team_a_to_team_b_with_user_id(self, responses_id_security): """Test that key from team A (without user_id) cannot access response from team B (with user_id)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -217,9 +223,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_admin_can_access_any_response( - self, responses_id_security - ): + def test_check_user_access_admin_can_access_any_response(self, responses_id_security): """Test that proxy admin can access any response ID""" from litellm.proxy._types import LitellmUserRoles @@ -238,9 +242,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_security_disabled( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_security_disabled(self, responses_id_security, mock_user_api_key_dict): """Test that when security is disabled, any user can access any response""" with patch( "litellm.proxy.proxy_server.general_settings", @@ -266,9 +268,7 @@ class TestAsyncPreCallHook: """Test pre-call hook decrypts previous_response_id for aresponses call""" data = {"previous_response_id": "resp_encrypted_value"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -284,15 +284,11 @@ class TestAsyncPreCallHook: assert result["previous_response_id"] == "resp_original_123" @pytest.mark.asyncio - async def test_async_pre_call_hook_aget_responses( - self, responses_id_security, mock_user_api_key_dict, mock_cache - ): + async def test_async_pre_call_hook_aget_responses(self, responses_id_security, mock_user_api_key_dict, mock_cache): """Test pre-call hook decrypts response_id for aget_responses call""" data = {"response_id": "resp_encrypted_456"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -308,9 +304,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_456" @pytest.mark.asyncio - async def test_async_pre_call_hook_team_a_accessing_team_b_response( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_team_a_accessing_team_b_response(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from accessing team B response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -320,9 +314,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -341,9 +333,7 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_team_a_accessing_team_b_with_user( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_team_a_accessing_team_b_with_user(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A (no user) from accessing team B response (with user)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -353,9 +343,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b_with_user"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -375,9 +363,7 @@ class TestAsyncPreCallHook: assert "forbidden" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_same_team_access( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_same_team_access(self, responses_id_security, mock_cache): """Test pre-call hook allows team A to access their own team's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -387,9 +373,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_a"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -405,9 +389,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_team_a" @pytest.mark.asyncio - async def test_async_pre_call_hook_adelete_responses_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_adelete_responses_team_security(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from deleting team B's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -417,9 +399,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -438,9 +418,7 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_acancel_responses_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_acancel_responses_team_security(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from canceling team B's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -450,9 +428,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -470,16 +446,13 @@ class TestAsyncPreCallHook: assert exc_info.value.status_code == 403 assert "team" in exc_info.value.detail.lower() - @pytest.mark.asyncio async def test_async_pre_call_hook_alist_input_items_decrypts_response_id( self, responses_id_security, mock_user_api_key_dict, mock_cache ): data = {"response_id": "resp_encrypted_789"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -496,9 +469,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_789" @pytest.mark.asyncio - async def test_async_pre_call_hook_alist_input_items_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_alist_input_items_team_security(self, responses_id_security, mock_cache): mock_auth_team_a = MagicMock() mock_auth_team_a.user_id = None mock_auth_team_a.team_id = "team-a" @@ -506,9 +477,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -620,9 +589,7 @@ class TestAsyncPostCallStreamingIteratorHook: "/openai/v1/responses/resp_rawprovider123", ], ) - async def test_streamed_id_encrypted_on_all_responses_routes( - self, responses_id_security, route, monkeypatch - ): + async def test_streamed_id_encrypted_on_all_responses_routes(self, responses_id_security, route, monkeypatch): streamed_id = await self._drain_streamed_id(responses_id_security, route, monkeypatch) assert streamed_id != "resp_rawprovider123" @@ -634,12 +601,8 @@ class TestAsyncPostCallStreamingIteratorHook: ) @pytest.mark.asyncio - async def test_streamed_id_untouched_on_non_responses_route( - self, responses_id_security, monkeypatch - ): - streamed_id = await self._drain_streamed_id( - responses_id_security, "/chat/completions", monkeypatch - ) + async def test_streamed_id_untouched_on_non_responses_route(self, responses_id_security, monkeypatch): + streamed_id = await self._drain_streamed_id(responses_id_security, "/chat/completions", monkeypatch) assert streamed_id == "resp_rawprovider123" assert not responses_id_security._is_encrypted_response_id(streamed_id) @@ -698,12 +661,8 @@ class TestStreamedGenericEventIdEncryption: def _advertised_ids(events): nested = (getattr(event, "response", None) for event in events) return [ - payload["id"] if isinstance(payload, dict) else payload.id - for payload in nested - if payload is not None - ] + [ - event.id for event in events if isinstance(getattr(event, "id", None), str) - ] + payload["id"] if isinstance(payload, dict) else payload.id for payload in nested if payload is not None + ] + [event.id for event in events if isinstance(getattr(event, "id", None), str)] async def _drain(self, responses_id_security, monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-abcdefghij") @@ -723,9 +682,7 @@ class TestStreamedGenericEventIdEncryption: ] @pytest.mark.asyncio - async def test_every_event_advertises_the_same_encrypted_id( - self, responses_id_security, monkeypatch - ): + async def test_every_event_advertises_the_same_encrypted_id(self, responses_id_security, monkeypatch): events = await self._drain(responses_id_security, monkeypatch) advertised = self._advertised_ids(events) @@ -742,22 +699,14 @@ class TestStreamedGenericEventIdEncryption: ) @pytest.mark.asyncio - async def test_raw_provider_id_never_reaches_the_client( - self, responses_id_security, monkeypatch - ): + async def test_raw_provider_id_never_reaches_the_client(self, responses_id_security, monkeypatch): events = await self._drain(responses_id_security, monkeypatch) - assert [self.RAW_ID in event.model_dump_json() for event in events] == [ - False - ] * len(events) + assert [self.RAW_ID in event.model_dump_json() for event in events] == [False] * len(events) @pytest.mark.asyncio - async def test_sibling_fields_survive_the_rewrite( - self, responses_id_security, monkeypatch - ): - _, queued, keepalive, later, _ = await self._drain( - responses_id_security, monkeypatch - ) + async def test_sibling_fields_survive_the_rewrite(self, responses_id_security, monkeypatch): + _, queued, keepalive, later, _ = await self._drain(responses_id_security, monkeypatch) assert queued.response["status"] == "queued" assert later.response["status"] == "in_progress" @@ -769,27 +718,19 @@ class TestAsyncPostCallSuccessHook: """Test async_post_call_success_hook function""" @pytest.mark.asyncio - async def test_async_post_call_success_hook_encrypts_response( - self, responses_id_security, mock_user_api_key_dict - ): + async def test_async_post_call_success_hook_encrypts_response(self, responses_id_security, mock_user_api_key_dict): """Test post-call hook encrypts ResponsesAPIResponse""" - mock_response = ResponsesAPIResponse( - id="resp_789", created_at=1234567890, output=[], status="completed" - ) + mock_response = ResponsesAPIResponse(id="resp_789", created_at=1234567890, output=[], status="completed") data = {} - with patch.object( - responses_id_security, "_encrypt_response_id", return_value=mock_response - ) as mock_encrypt: + with patch.object(responses_id_security, "_encrypt_response_id", return_value=mock_response) as mock_encrypt: result = await responses_id_security.async_post_call_success_hook( data=data, user_api_key_dict=mock_user_api_key_dict, response=mock_response, ) - mock_encrypt.assert_called_once_with( - mock_response, mock_user_api_key_dict, request_cache=None - ) + mock_encrypt.assert_called_once_with(mock_response, mock_user_api_key_dict, request_cache=None) assert result == mock_response @pytest.mark.asyncio @@ -829,7 +770,9 @@ def salt_key_env(monkeypatch): return _UNIT_TEST_SALT_KEY -def _hook(general_settings=None, signing_key=_UNIT_TEST_SALT_KEY): +def _hook( + general_settings: Mapping[str, object] | None = None, signing_key: str | None = _UNIT_TEST_SALT_KEY +) -> ResponsesIDSecurity: settings = general_settings if general_settings is not None else {} return ResponsesIDSecurity( general_settings_reader=lambda: settings, @@ -837,21 +780,254 @@ def _hook(general_settings=None, signing_key=_UNIT_TEST_SALT_KEY): ) -def _auth(user_id="owner-user", team_id="owner-team", user_role=None): - from litellm.proxy._types import UserAPIKeyAuth - +def _auth( + user_id: str = "owner-user", team_id: str = "owner-team", user_role: LitellmUserRoles | None = None +) -> UserAPIKeyAuth: return UserAPIKeyAuth(user_id=user_id, team_id=team_id, user_role=user_role) -def _issue_managed_id(hook, owner, provider_response_id=_FABRICATED_PROVIDER_RESPONSE_ID): +def _issue_managed_id( + hook: ResponsesIDSecurity, owner: UserAPIKeyAuth, provider_response_id: str = _FABRICATED_PROVIDER_RESPONSE_ID +) -> str: """Mint an id exactly the way the proxy hands one to a client on create.""" - issued = hook._encrypt_response_id( - ResponsesAPIResponse( - id=provider_response_id, created_at=1234567890, output=[], status="completed" - ), + issued = hook._encrypt_response_id( # pyright: ignore[reportPrivateUsage] # mints ids the way create does + ResponsesAPIResponse(id=provider_response_id, created_at=1234567890, output=[], status="completed"), owner, ) - return issued.id + return ResponsesAPIResponse.model_validate(issued).id + + +def _comparison_options(data: object, nested: bool) -> Mapping[str, str]: + request: Final = TypeAdapter(dict[str, object]).validate_python(data) + container: Final = TypeAdapter(dict[str, object]).validate_python(request["extra_body"]) if nested else request + return TypeAdapter(dict[str, str]).validate_python(container["prompt_cache_options"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", (False, True)) +@pytest.mark.parametrize("stranger", ("user", "team")) +@pytest.mark.parametrize("allow_unmanaged", (False, True)) +async def test_regression_comparison_id_reentry_keeps_ownership_after_owner_retry( + salt_key_env: str, nested: bool, stranger: str, allow_unmanaged: bool +) -> None: + hook: Final = _hook({"allow_unmanaged_response_ids": allow_unmanaged}, signing_key=salt_key_env) + owner: Final = UserAPIKeyAuth(user_id="owner-user", team_id="owner-team") + caller: Final = UserAPIKeyAuth( + user_id="stranger-user" if stranger == "user" else owner.user_id, + team_id="stranger-team" if stranger == "team" else owner.team_id, + ) + cache: Final = DualCache() + issued: Final = ResponsesAPIResponse.model_validate( + await hook.async_post_call_success_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + data={}, + user_api_key_dict=owner, + response=ResponsesAPIResponse( + id=_FABRICATED_PROVIDER_RESPONSE_ID, created_at=1234567890, output=[], status="completed" + ), + ) + ) + assert issued.id != _FABRICATED_PROVIDER_RESPONSE_ID + options: Final = {"comparison_response_id": issued.id} + request: Final = {"extra_body": {"prompt_cache_options": options}} if nested else {"prompt_cache_options": options} + first: Final = TypeAdapter(dict[str, object]).validate_python( + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=owner, cache=cache, data=request, call_type="aresponses" + ) + ) + retried: Final = TypeAdapter(dict[str, object]).validate_python( + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=owner, cache=cache, data=request, call_type="aresponses" + ) + ) + with pytest.raises(HTTPException) as refusal: + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=caller, cache=cache, data=request, call_type="aresponses" + ) + assert refusal.value.status_code == 403 + assert f"not associated with the {stranger}" in refusal.value.detail + assert _comparison_options(first, nested) == {"comparison_response_id": _FABRICATED_PROVIDER_RESPONSE_ID} + assert _comparison_options(retried, nested) == {"comparison_response_id": _FABRICATED_PROVIDER_RESPONSE_ID} + assert _comparison_options(request, nested) == {"comparison_response_id": issued.id} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", (False, True)) +@pytest.mark.parametrize("replacement", ("resp_vendor_raw", "resp_!!!not-base64", "")) +async def test_guard_comparison_id_replaced_after_owner_hook_reaches_provider_unchanged( + salt_key_env: str, nested: bool, replacement: str +) -> None: + hook: Final = ResponsesIDSecurity( + general_settings_reader=lambda: {"allow_unmanaged_response_ids": True}, + signing_key_reader=lambda: salt_key_env, + ) + owner: Final = UserAPIKeyAuth(user_id="owner-user", team_id="owner-team") + cache: Final = DualCache() + issued: Final = ResponsesAPIResponse.model_validate( + await hook.async_post_call_success_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + data={}, + user_api_key_dict=owner, + response=ResponsesAPIResponse( + id=_FABRICATED_PROVIDER_RESPONSE_ID, created_at=1234567890, output=[], status="completed" + ), + ) + ) + first_options: Final = {"comparison_response_id": issued.id} + first_request: Final = ( + {"extra_body": {"prompt_cache_options": first_options}} if nested else {"prompt_cache_options": first_options} + ) + first: Final = TypeAdapter(dict[str, object]).validate_python( + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=owner, cache=cache, data=first_request, call_type="aresponses" + ) + ) + replaced_options: Final = {"comparison_response_id": replacement, "mode": "implicit"} + replaced: Final = ( + {**first, "extra_body": {"prompt_cache_options": replaced_options}} + if nested + else {**first, "prompt_cache_options": replaced_options} + ) + rerun: Final = TypeAdapter(dict[str, object]).validate_python( + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=owner, cache=cache, data=replaced, call_type="aresponses" + ) + ) + assert _comparison_options(rerun, nested) == replaced_options, rerun + + +def _comparison_request(comparison_response_id: str, nested: bool) -> dict[str, object]: + options: Final = {"comparison_response_id": comparison_response_id, "mode": "implicit"} + return {"extra_body": {"prompt_cache_options": options}} if nested else {"prompt_cache_options": options} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", (False, True)) +@pytest.mark.parametrize("replacement", ("resp_vendor_raw", "resp_!!!not-base64", "")) +async def test_comparison_id_replaced_after_owner_hook_is_rejected_by_default( + salt_key_env: str, nested: bool, replacement: str +) -> None: + hook: Final = _hook() + owner: Final = _auth() + first: Final = TypeAdapter(dict[str, object]).validate_python( + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=owner, + cache=DualCache(), + data=_comparison_request(_issue_managed_id(hook, owner), nested), + call_type="aresponses", + ) + ) + replaced: Final = {**first, **_comparison_request(replacement, nested)} + with pytest.raises(HTTPException) as refusal: + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=owner, cache=DualCache(), data=replaced, call_type="aresponses" + ) + assert refusal.value.status_code == 403 + assert "allow_unmanaged_response_ids" in refusal.value.detail + assert _comparison_options(replaced, nested)["comparison_response_id"] == replacement + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", (False, True)) +async def test_raw_comparison_id_is_rejected_by_default(salt_key_env: str, nested: bool) -> None: + request: Final = _comparison_request(_FABRICATED_UNMANAGED_ID, nested) + with pytest.raises(HTTPException) as refusal: + await _hook().async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=_auth(), cache=DualCache(), data=request, call_type="aresponses" + ) + assert refusal.value.status_code == 403 + assert "allow_unmanaged_response_ids" in refusal.value.detail + assert _comparison_options(request, nested)["comparison_response_id"] == _FABRICATED_UNMANAGED_ID + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", (False, True)) +@pytest.mark.parametrize( + ("general_settings", "signing_key", "user_role"), + [ + ({"allow_unmanaged_response_ids": True}, _UNIT_TEST_SALT_KEY, None), + ({"disable_responses_id_security": True}, _UNIT_TEST_SALT_KEY, None), + ({}, None, None), + ({}, _UNIT_TEST_SALT_KEY, LitellmUserRoles.PROXY_ADMIN), + ], +) +async def test_raw_comparison_id_follows_previous_response_id_escape_hatches( + salt_key_env: str, + nested: bool, + general_settings: dict[str, bool], + signing_key: str | None, + user_role: LitellmUserRoles | None, +) -> None: + hook: Final = _hook(general_settings=general_settings, signing_key=signing_key) + result: Final = TypeAdapter(dict[str, object]).validate_python( + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=_auth(user_role=user_role), + cache=DualCache(), + data=_comparison_request(_FABRICATED_UNMANAGED_ID, nested), + call_type="aresponses", + ) + ) + assert _comparison_options(result, nested) == { + "comparison_response_id": _FABRICATED_UNMANAGED_ID, + "mode": "implicit", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", (False, True)) +async def test_retained_comparison_id_of_another_owner_cannot_authorize_its_raw_id( + salt_key_env: str, nested: bool +) -> None: + hook: Final = _hook() + victim_id: Final = _issue_managed_id(hook, _auth()) + request: Final = { + **_comparison_request(_FABRICATED_PROVIDER_RESPONSE_ID, nested), + ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: victim_id, + } + with pytest.raises(HTTPException) as refusal: + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=_auth(user_id="stranger-user", team_id="stranger-team"), + cache=DualCache(), + data=request, + call_type="aresponses", + ) + assert refusal.value.status_code == 403 + assert "not associated with the user" in refusal.value.detail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", (False, True)) +async def test_retained_own_comparison_id_cannot_carry_a_different_raw_id(salt_key_env: str, nested: bool) -> None: + hook: Final = _hook() + stranger: Final = _auth(user_id="stranger-user", team_id="stranger-team") + stranger_id: Final = _issue_managed_id(hook, stranger, provider_response_id="resp_strangerownprovideridcccccccc") + request: Final = { + **_comparison_request("resp_victimprovideriddddddddddddddddddddd", nested), + ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: stranger_id, + } + with pytest.raises(HTTPException) as refusal: + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=stranger, cache=DualCache(), data=request, call_type="aresponses" + ) + assert refusal.value.status_code == 403 + assert "allow_unmanaged_response_ids" in refusal.value.detail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", (False, True)) +async def test_newly_supplied_comparison_id_wins_over_the_retained_one(salt_key_env: str, nested: bool) -> None: + hook: Final = _hook() + owner: Final = _auth() + stale_id: Final = _issue_managed_id(hook, owner, provider_response_id="resp_staleownerprovideridbbbbbbbbbb") + fresh_id: Final = _issue_managed_id(hook, owner) + result: Final = TypeAdapter(dict[str, object]).validate_python( + await hook.async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # Legacy dict boundary + user_api_key_dict=owner, + cache=DualCache(), + data={**_comparison_request(fresh_id, nested), ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: stale_id}, + call_type="aresponses", + ) + ) + assert _comparison_options(result, nested)["comparison_response_id"] == _FABRICATED_PROVIDER_RESPONSE_ID + assert result[ADDRESSED_COMPARISON_RESPONSE_ID_FIELD] == fresh_id class TestUnrecognizedResponseIdIsRejected: @@ -1080,9 +1256,7 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization: assert exc_info.value.status_code == 403 @pytest.mark.asyncio - async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response( - self, mock_cache, salt_key_env - ): + async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response(self, mock_cache, salt_key_env): hook = _hook() stranger = _auth(user_id="stranger-user", team_id="stranger-team") stranger_id = _issue_managed_id(hook, stranger, provider_response_id="resp_strangerownprovideridcccccccc") @@ -1098,3 +1272,435 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization: assert result["response_id"] == "resp_strangerownprovideridcccccccc" assert result["response_id"] != victim_provider_id + + +def _admin(): + from litellm.proxy._types import LitellmUserRoles + + return _auth(user_id="admin-user", team_id="admin-team", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _provider_answer(response_id: str) -> ResponsesAPIResponse: + return ResponsesAPIResponse(id=response_id, created_at=1234567890, output=[], status="completed") + + +class TestReadsAnswerWithTheIdTheClientHolds: + """A follow-up looks its history up by the id the client holds, so a read must not hand out a new one.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) + @pytest.mark.parametrize( + "issuer, reader", + [ + (_auth(), _auth()), + (_auth(team_id=None), _auth(team_id="another-team")), + (_auth(user_id=None), _auth(user_id=None)), + (_auth(user_id=None), _auth(user_id="teammate-user")), + (_auth(), _admin()), + (_auth(user_id=None), _auth(user_id=None, team_id="admin-team", user_role=_admin().user_role)), + ], + ids=[ + "same-key", + "same-user-new-team", + "same-team-no-user", + "teammate-reading-a-team-key-id", + "admin-user", + "admin-without-user", + ], + ) + async def test_a_caller_with_access_gets_back_the_exact_id_they_sent( + self, mock_cache, salt_key_env, call_type, issuer, reader + ): + hook = _hook() + issued_id = _issue_managed_id(hook, issuer) + data = await hook.async_pre_call_hook( + user_api_key_dict=reader, cache=mock_cache, data={"response_id": issued_id}, call_type=call_type + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=reader, response=_provider_answer(data["response_id"]) + ) + + assert result.id == issued_id + + @pytest.mark.asyncio + async def test_a_follow_up_gets_an_id_for_its_own_response_never_the_previous_one(self, mock_cache, salt_key_env): + hook = _hook() + owner = _auth() + previous_id = _issue_managed_id(hook, owner) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, + cache=mock_cache, + data={"model": "gpt-fake", "previous_response_id": previous_id}, + call_type="aresponses", + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=owner, response=_provider_answer("resp_followupprovideridzzzzzzzzzzzzzzz") + ) + + assert result.id != previous_id + assert hook._decrypt_response_id(result.id) == ( + "resp_followupprovideridzzzzzzzzzzzzzzz", + "owner-user", + "owner-team", + ) + + @pytest.mark.asyncio + async def test_an_admin_read_by_a_raw_provider_id_gets_a_managed_id_back(self, mock_cache, salt_key_env): + hook = _hook() + admin = _admin() + data = await hook.async_pre_call_hook( + user_api_key_dict=admin, + cache=mock_cache, + data={"response_id": _FABRICATED_UNMANAGED_ID}, + call_type="aget_responses", + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=admin, response=_provider_answer(_FABRICATED_UNMANAGED_ID) + ) + + assert hook._is_encrypted_response_id(result.id) + assert hook._decrypt_response_id(result.id) == (_FABRICATED_UNMANAGED_ID, "admin-user", "admin-team") + + +_LOGGED_MANAGED_ID = "resp_managedidlitellmbuiltbeforethehookran" + + +def _streamed_response(response_id): + return ResponsesAPIResponse( + id=response_id, + created_at=0, + model="gpt-5.1", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + usage=ResponseAPIUsage(input_tokens=3, output_tokens=2, total_tokens=5), + ) + + +def _upstream_frames(): + return ( + ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=_streamed_response(_LOGGED_MANAGED_ID) + ), + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_streamed_response(_LOGGED_MANAGED_ID) + ), + ) + + +def _native_logged_result(): + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_streamed_response(_LOGGED_MANAGED_ID) + ) + + +def _bridged_logged_result(): + return ModelResponse( + id="chatcmpl-bridgedchatcompletionid", + model="gpt-5.1", + choices=[Choices(message=Message(role="assistant", content="hi"), finish_reason="stop", index=0)], + usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +class _SpendRowRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.spend_row_ids = [] + self.standard_logging_ids = [] + self.callback_response_ids = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.spend_row_ids.append(get_logging_payload(kwargs, response_obj, start_time, end_time)["request_id"]) + self.callback_response_ids.append(response_obj.id) + if "standard_logging_object" in kwargs: + self.standard_logging_ids.append(kwargs["standard_logging_object"]["id"]) + + +async def _stream_logged_mid_stream( + hook: ResponsesIDSecurity, + route: str, + logged_result: ResponseCompletedEvent | ModelResponse, + caller: UserAPIKeyAuth | None = None, + request_data: Mapping[str, object] | None = None, +) -> tuple[str, _SpendRowRecorder]: + streaming_caller = caller if caller is not None else _auth() + recorder = _SpendRowRecorder() + logging_obj = LiteLLMLoggingObj( + model="gpt-5.1", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime.datetime.now(), + litellm_call_id="streamed-responses-call", + function_id="streamed-responses-fn", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.1", user="", optional_params={}, litellm_params={"api_base": ""} + ) + upstream = _upstream_frames() + + async def frames(): + for frame in upstream: + yield frame + + client_stream = hook.async_post_call_streaming_iterator_hook( + user_api_key_dict=streaming_caller.model_copy(update={"request_route": route}), + response=frames(), + request_data={**(request_data or {}), "litellm_logging_obj": logging_obj}, + ) + first_frame = await anext(client_stream) + now = datetime.datetime.now() + await logging_obj.async_success_handler(result=logged_result, start_time=now, end_time=now, cache_hit=False) + client_frames = [first_frame, *[frame async for frame in client_stream]] + return client_frames[-1].response.id, recorder + + +class TestStreamedSpendRowCarriesTheIdTheClientSaw: + @pytest.mark.asyncio + @pytest.mark.parametrize("logged_result", [_native_logged_result, _bridged_logged_result], ids=["native", "bridge"]) + @pytest.mark.parametrize("route", ["/v1/responses", "/responses", "/openai/v1/responses"]) + async def test_spend_row_and_callbacks_log_the_encrypted_id_of_the_terminal_frame( + self, salt_key_env, route, logged_result + ): + hook = _hook() + + client_id, recorder = await _stream_logged_mid_stream(hook, route, logged_result()) + + assert hook._decrypt_response_id(client_id) == (_LOGGED_MANAGED_ID, "owner-user", "owner-team") + assert recorder.spend_row_ids == [client_id] + assert recorder.standard_logging_ids == [client_id] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "hook, route", + [ + (_hook(), "/chat/completions"), + (_hook(general_settings={"disable_responses_id_security": True}), "/v1/responses"), + (_hook(signing_key=None), "/v1/responses"), + ], + ids=["not-a-create-route", "security-disabled", "no-signing-key"], + ) + async def test_stream_the_hook_leaves_alone_keeps_its_logged_id(self, salt_key_env, hook, route): + client_id, recorder = await _stream_logged_mid_stream(hook, route, _native_logged_result()) + + assert client_id == _LOGGED_MANAGED_ID + assert recorder.spend_row_ids == [_LOGGED_MANAGED_ID] + assert recorder.standard_logging_ids == [_LOGGED_MANAGED_ID] + + +class TestStreamedReadAnswersWithTheIdTheClientHolds: + @pytest.mark.asyncio + @pytest.mark.parametrize("route_prefix", ["/v1/responses", "/responses", "/openai/v1/responses"]) + async def test_the_owner_sees_their_id_while_callbacks_log_one_of_the_reads_own( + self, mock_cache, salt_key_env, route_prefix + ): + hook = _hook() + owner = _auth() + issued_id = _issue_managed_id(hook, owner, provider_response_id=_LOGGED_MANAGED_ID) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, recorder = await _stream_logged_mid_stream( + hook, f"{route_prefix}/{issued_id}", _native_logged_result(), caller=owner, request_data=data + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.standard_logging_ids + assert len(recorder.standard_logging_ids) == 1 + assert recorder.standard_logging_ids[0] != issued_id + assert hook._decrypt_response_id(recorder.standard_logging_ids[0]) == ( + _LOGGED_MANAGED_ID, + "owner-user", + "owner-team", + ) + + @pytest.mark.asyncio + async def test_an_admin_streamed_read_of_someone_elses_response_sees_the_id_sent_and_logs_one_of_its_own( + self, mock_cache, salt_key_env + ): + hook = _hook() + admin = _admin() + issued_id = _issue_managed_id(hook, _auth(), provider_response_id=_LOGGED_MANAGED_ID) + data = await hook.async_pre_call_hook( + user_api_key_dict=admin, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, recorder = await _stream_logged_mid_stream( + hook, f"/v1/responses/{issued_id}", _native_logged_result(), caller=admin, request_data=data + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.standard_logging_ids + assert len(recorder.standard_logging_ids) == 1 + assert hook._decrypt_response_id(recorder.standard_logging_ids[0]) == ( + _LOGGED_MANAGED_ID, + "admin-user", + "admin-team", + ) + + @pytest.mark.asyncio + async def test_a_frame_naming_the_held_response_under_another_wrapping_still_gets_the_id_sent( + self, mock_cache, salt_key_env + ): + hook = _hook() + owner = _auth() + issued_id = _issue_managed_id(hook, owner, provider_response_id=_routed_response_id(_LOGGED_MANAGED_ID)) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, _ = await _stream_logged_mid_stream( + hook, f"/v1/responses/{issued_id}", _native_logged_result(), caller=owner, request_data=data + ) + + assert client_id == issued_id + + +async def _read_logged_by_callbacks( + hook: ResponsesIDSecurity, + caller: UserAPIKeyAuth, + issued_id: str, + call_type: CallTypesLiteral, + cache: DualCache, + answer: Callable[[str], ResponsesAPIResponse | DeleteResponseResult] = _provider_answer, +) -> tuple[str, _SpendRowRecorder]: + recorder = _SpendRowRecorder() + now = datetime.datetime.now() + logging_obj = LiteLLMLoggingObj( + model="gpt-5.1", + messages=[], + stream=False, + call_type=call_type, + start_time=now, + litellm_call_id=f"{call_type}-call", + function_id=f"{call_type}-fn", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.1", user="", optional_params={}, litellm_params={"api_base": ""} + ) + data = await hook.async_pre_call_hook( + user_api_key_dict=caller, cache=cache, data={"response_id": issued_id}, call_type=call_type + ) + provider_answer = answer(data["response_id"]) + + client_response = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=caller, response=provider_answer + ) + await logging_obj.async_success_handler(result=provider_answer, start_time=now, end_time=now, cache_hit=False) + return client_response.id, recorder + + +class TestEveryReadLogsAnIdOfItsOwn: + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + async def test_callbacks_and_the_spend_row_log_a_new_id_per_read_never_the_one_the_client_holds( + self, mock_cache, salt_key_env, call_type, reader + ): + hook = _hook() + issued_id = _issue_managed_id(hook, _auth()) + + reads = [await _read_logged_by_callbacks(hook, reader, issued_id, call_type, mock_cache) for _ in range(2)] + + assert [client_id for client_id, _ in reads] == [issued_id, issued_id] + logged_ids = [id_ for _, recorder in reads for id_ in recorder.standard_logging_ids] + assert [recorder.spend_row_ids for _, recorder in reads] == [[logged_id] for logged_id in logged_ids] + assert len(set(logged_ids)) == 2 + assert issued_id not in logged_ids + assert {hook._decrypt_response_id(logged_id) for logged_id in logged_ids} == { + (_FABRICATED_PROVIDER_RESPONSE_ID, reader.user_id, reader.team_id) + } + + +_RAW_UPSTREAM_RESPONSE_ID = "resp_06de59f3upstreamiddddddddddddddddd" + + +def _routed_response_id(upstream_response_id: str) -> str: + return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + _provider_answer(upstream_response_id), "openai", {"model_info": {"id": "deployment-1"}} + ).id + + +def _deletion_of(response_id: str) -> DeleteResponseResult: + return DeleteResponseResult( + id=ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id), + object="response.deleted", + deleted=True, + ) + + +class TestDeleteAnswersWithTheIdTheClientHolds: + @pytest.mark.asyncio + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + @pytest.mark.parametrize( + "stored_id", + [_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID), _RAW_UPSTREAM_RESPONSE_ID], + ids=["routed-deployment", "plain-id"], + ) + async def test_a_caller_with_access_gets_back_the_exact_id_they_sent( + self, mock_cache, salt_key_env, reader, stored_id + ): + hook = _hook() + issued_id = _issue_managed_id(hook, _auth(), provider_response_id=stored_id) + data = await hook.async_pre_call_hook( + user_api_key_dict=reader, cache=mock_cache, data={"response_id": issued_id}, call_type="adelete_responses" + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=reader, response=_deletion_of(data["response_id"]) + ) + + assert (result.id, result.object, result.deleted) == (issued_id, "response.deleted", True) + + @pytest.mark.asyncio + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + async def test_callbacks_and_the_spend_row_log_an_id_of_its_own_never_the_one_the_client_holds( + self, mock_cache, salt_key_env, reader + ): + hook = _hook() + issued_id = _issue_managed_id( + hook, _auth(), provider_response_id=_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID) + ) + + client_id, recorder = await _read_logged_by_callbacks( + hook, reader, issued_id, "adelete_responses", mock_cache, answer=_deletion_of + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.callback_response_ids + assert len(recorder.callback_response_ids) == 1 + assert recorder.callback_response_ids[0] != issued_id + assert hook._decrypt_response_id(recorder.callback_response_ids[0]) == ( + _RAW_UPSTREAM_RESPONSE_ID, + reader.user_id, + reader.team_id, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "sent_id", + [_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID), _RAW_UPSTREAM_RESPONSE_ID], + ids=["routed-deployment", "plain-id"], + ) + async def test_an_admin_delete_by_a_raw_provider_id_gets_the_provider_id_back_and_logs_an_id_of_its_own( + self, mock_cache, salt_key_env, sent_id + ): + hook = _hook() + + client_id, recorder = await _read_logged_by_callbacks( + hook, _admin(), sent_id, "adelete_responses", mock_cache, answer=_deletion_of + ) + + assert client_id == _RAW_UPSTREAM_RESPONSE_ID + assert [hook._decrypt_response_id(logged_id) for logged_id in recorder.callback_response_ids] == [ + (_RAW_UPSTREAM_RESPONSE_ID, "admin-user", "admin-team") + ] diff --git a/tests/unit/test_router/core/test_fallbacks.py b/tests/unit/test_router/core/test_fallbacks.py index e4e5c3120c8..bc6c70543c8 100644 --- a/tests/unit/test_router/core/test_fallbacks.py +++ b/tests/unit/test_router/core/test_fallbacks.py @@ -1,6 +1,7 @@ """Cross-group fallbacks of a router core: where the chain comes from, how it is walked, and what is raised.""" from collections.abc import Callable, Mapping +from typing import Final import pytest @@ -355,21 +356,13 @@ async def test_an_error_specific_chain_wins_over_the_general_chain( @pytest.mark.asyncio -@pytest.mark.parametrize( - ("error", "field", "raised"), - [ - (context_window, "context_window_fallbacks", litellm.ContextWindowExceededError), - (content_policy, "content_policy_fallbacks", litellm.ContentPolicyViolationError), - ], - ids=["context-window", "content-policy"], -) -async def test_an_error_specific_chain_without_the_group_raises_without_trying_the_general_chain( - make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception], field: str, raised: type[Exception] +async def test_an_unmatched_content_policy_chain_does_not_try_the_general_chain( + make_core: MakeCore, host: ScriptedHost ) -> None: - core = build(make_core, fallbacks=[{"a": ["b"]}], **{field: [{"other": ["c"]}]}) - host.answer("openai/a", error()) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], content_policy_fallbacks=[{"other": ["c"]}]) + host.answer("openai/a", content_policy()) - with pytest.raises(raised): + with pytest.raises(litellm.ContentPolicyViolationError): await complete(core) assert host.deployments_called("acompletion") == ["a"] diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 9996ffcb098..4f35efdfd6e 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -6776,6 +6776,151 @@ async def test_aresponses_streaming_content_policy_error_event_routes_to_content assert fallback_call.await_args.kwargs["model"] == "gemini/gemini-2.5-flash" +@pytest.mark.asyncio +async def test_aresponses_stream_keeps_numbering_after_a_mid_stream_fallback(): + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=sequence_number, + ) + + from litellm.exceptions import MidStreamFallbackError + + router: Final = _InjectedFallbackRouter(_AsyncList([delta(0, "c"), delta(1, "d")])) + primary: Final = _make_responses_iterator( + chunks=[delta(0, "a"), delta(1, "b")], + error=MidStreamFallbackError( + message="primary died", + model="primary", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="ab", + ), + model="primary", + ) + + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, + ) + events: Final = [event async for event in wrapped] + + assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 1), ("c", 2), ("d", 3)] + + + +@pytest.mark.asyncio +async def test_aresponses_stream_fallback_keeps_higher_numbers_and_passes_unnumbered_events(): + from litellm.exceptions import MidStreamFallbackError + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + def delta(text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="msg_1", output_index=0, content_index=0, delta=text + ) + + def numbered(sequence_number: int, text: str) -> OutputTextDeltaEvent: + return delta(text).model_copy(update={"sequence_number": sequence_number}) + + router: Final = _InjectedFallbackRouter(_AsyncList([numbered(7, "c"), delta("d"), numbered(8, "e")])) + primary: Final = _make_responses_iterator( + chunks=[numbered(0, "a"), numbered(1, "b")], + error=MidStreamFallbackError( + message="primary died", + model="primary", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="ab", + ), + model="primary", + ) + + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, + ) + events: Final = [event async for event in wrapped] + + assert [(event.delta, event.get("sequence_number")) for event in events] == [ + ("a", 0), + ("b", 1), + ("c", 7), + ("d", None), + ("e", 8), + ] + + +@pytest.mark.asyncio +async def test_aresponses_stream_without_fallback_passes_upstream_numbers_unchanged(): + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=sequence_number, + ) + + router: Final = _InjectedFallbackRouter(_AsyncList([])) + primary: Final = _make_responses_iterator( + chunks=[delta(0, "a"), delta(1, "b"), delta(1, "c"), delta(0, "d")], model="primary" + ) + + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, + ) + events: Final = [event async for event in wrapped] + + assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 1), ("c", 1), ("d", 0)] + + +@pytest.mark.asyncio +async def test_aresponses_stream_fallback_continues_above_the_highest_number_the_primary_sent(): + from litellm.exceptions import MidStreamFallbackError + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=sequence_number, + ) + + router: Final = _InjectedFallbackRouter(_AsyncList([delta(0, "d")])) + primary: Final = _make_responses_iterator( + chunks=[delta(0, "a"), delta(3, "b"), delta(1, "c")], + error=MidStreamFallbackError( + message="primary died", + model="primary", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="abc", + ), + model="primary", + ) + + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, + ) + events: Final = [event async for event in wrapped] + + assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 3), ("c", 1), ("d", 4)] + + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_fallback_dispatch(): """The fallback dispatch matches on the trigger's own type, so the wrapper must hand it the diff --git a/tests/unit/test_router/test_router_fallbacks.py b/tests/unit/test_router/test_router_fallbacks.py index bccd6290f44..b25d5c35892 100644 --- a/tests/unit/test_router/test_router_fallbacks.py +++ b/tests/unit/test_router/test_router_fallbacks.py @@ -1207,3 +1207,156 @@ def test_context_stream_fallback_preserves_other_errors_for_existing_continuatio message="provider unavailable", model="primary", llm_provider="openai", generated_content="already visible" ) assert _context_window_stream_fallback_error(error, "primary", "openai", True) is error + + +def _context_fallback_precedence_router(context_key: str | None, ordered: bool) -> Router: + deployments: Final = ( + ("primary", "primary", 1), + ("primary" if ordered else "unused", "ordered", 2), + ("generic", "generic", 1), + ("context", "context", 1), + ) + return Router( + model_list=[ + { + "model_name": group, + "litellm_params": { + "model": f"openai/{model}", + "api_key": "test-key", + "api_base": "https://fallback-migration.local/v1", + "order": order, + }, + } + for group, model, order in deployments + ], + fallbacks=[{"primary": ["generic"]}], + context_window_fallbacks=None if context_key is None else [{context_key: ["context"]}], + num_retries=0, + disable_cooldowns=True, + ) + + +@pytest.mark.parametrize("sync_mode", (False, True)) +@pytest.mark.parametrize( + ("context_key", "ordered", "flat_fallbacks", "expected_model"), + ( + (None, True, False, "ordered"), + ("unrelated", True, False, "ordered"), + ("unrelated", False, False, "generic"), + ("primary", True, False, "context"), + ("*", True, False, "context"), + ("primary", True, True, "context"), + ("primary", False, True, "context"), + ), +) +@pytest.mark.asyncio +async def test_context_fallback_precedence_is_specific_to_the_requested_group( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + context_key: str | None, + ordered: bool, + flat_fallbacks: bool, + expected_model: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _context_fallback_precedence_router(context_key, ordered) + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=( + httpx.Response( + 400, + json={ + "error": { + "message": "This model's maximum context length is 10 tokens", + "type": "invalid_request_error", + "code": "context_length_exceeded", + } + }, + ), + _fallback_chat_response(expected_model, f"served by {expected_model}"), + ) + ) + + response: Final = ( + router.completion( + model="primary", + messages=[{"role": "user", "content": "long prompt"}], + fallbacks=["generic"] if flat_fallbacks else [{"primary": ["generic"]}], + ) + if sync_mode + else await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "long prompt"}], + fallbacks=["generic"] if flat_fallbacks else [{"primary": ["generic"]}], + ) + ) + + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", expected_model) + assert response.choices[0].message.content == f"served by {expected_model}" + + +@pytest.mark.parametrize("matching_context", (False, True)) +@pytest.mark.asyncio +async def test_context_fallback_exhaustion_preserves_the_selected_chain( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + matching_context: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _context_fallback_precedence_router("primary" if matching_context else "unrelated", True) + refusal: Final = httpx.Response( + 400, + json={"error": {"message": "This model's maximum context length is 10 tokens"}}, + ) + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=(refusal, refusal, _fallback_chat_response("generic", "generic answer")) + ) + + if matching_context: + with pytest.raises(litellm.ContextWindowExceededError): + await router.acompletion(model="primary", messages=[{"role": "user", "content": "long prompt"}]) + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "context") + return + + response: Final = await router.acompletion(model="primary", messages=[{"role": "user", "content": "long prompt"}]) + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "ordered", "generic") + assert response.choices[0].message.content == "generic answer" + + +@pytest.mark.parametrize("sync_mode", (False, True)) +@pytest.mark.parametrize("context_error", (False, True)) +@pytest.mark.parametrize("backup_status", (429, 503)) +@pytest.mark.asyncio +async def test_ordered_fallback_failures_preserve_the_original_context_error( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + context_error: bool, + backup_status: int, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _context_fallback_precedence_router("unrelated", True) + original_message: Final = ( + "This model's maximum context length is 10 tokens" if context_error else "primary unavailable" + ) + backup_error: Final = litellm.RateLimitError if backup_status == 429 else litellm.ServiceUnavailableError + expected_error: Final = litellm.ContextWindowExceededError if context_error else backup_error + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=( + httpx.Response(400 if context_error else 503, json={"error": {"message": original_message}}), + httpx.Response(backup_status, json={"error": {"message": "ordered deployment unavailable"}}), + httpx.Response(backup_status, json={"error": {"message": "generic deployment unavailable"}}), + ) + ) + + request: Final = ( + asyncio.to_thread(router.completion, model="primary", messages=[{"role": "user", "content": "long prompt"}]) + if sync_mode + else router.acompletion(model="primary", messages=[{"role": "user", "content": "long prompt"}]) + ) + with pytest.raises(expected_error) as raised: + await request + + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "ordered", "generic") + assert raised.value.status_code == (400 if context_error else backup_status) + assert (original_message if context_error else "generic deployment unavailable") in str(raised.value) diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 4964dea498f..5ad6eb819ef 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -3,8 +3,9 @@ import functools import gc import json import os +import uuid import weakref -from collections.abc import Awaitable, Callable, Iterator, Mapping +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping from datetime import datetime, timezone from types import MethodType from typing import Final @@ -13,6 +14,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from websockets.exceptions import ConnectionClosed +from websockets.frames import Close import litellm from litellm import Router @@ -21,8 +24,14 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.realtime_errors import close_after_upstream_handshake_refusal +from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.litellm_params import RoutingStrategyName from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( @@ -540,6 +549,242 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( assert await router.cache.async_get_cache(key=tpm_key) == 17 +class _DiscardingLoggingWorker: + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: + async_coroutine.close() + + +async def _client_never_sends() -> str: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + +def _realtime_client_ws() -> MagicMock: + client_ws: Final = MagicMock() + client_ws.headers = {} + client_ws.receive_text = AsyncMock(side_effect=_client_never_sends) + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + return client_ws + + +def _realtime_session_logging() -> LiteLLMLogging: + return LiteLLMLogging( + model="gpt-realtime-2.1", + messages=[], + stream=False, + call_type="_arealtime", + start_time=datetime.now(), + litellm_call_id=str(uuid.uuid4()), + function_id="realtime-session", + ) + + +async def _relay_upstream( + litellm_logging_obj: LiteLLMLogging, + *upstream_frames: bytes | Exception, + client_ws: MagicMock | None = None, + provider_config: BaseRealtimeConfig | None = None, + upstream_stays_open: bool = False, +) -> None: + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=_frames_then_silence(*upstream_frames) if upstream_stays_open else list(upstream_frames) + ) + backend_ws.send = AsyncMock() + relay: Final = RealTimeStreaming( + client_ws or _realtime_client_ws(), + backend_ws, + litellm_logging_obj, + provider_config=provider_config, + model="gpt-realtime-2.1", + logging_worker=_DiscardingLoggingWorker(), + ) + await asyncio.wait_for(relay.bidirectional_forward(), timeout=2) + + +_REALTIME_SESSION_CREATED: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + + +async def _clean_realtime_session(litellm_logging_obj: LiteLLMLogging) -> None: + await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(Close(1000, ""), None)) + + +async def _upstream_internal_error_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: + await _relay_upstream( + litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(Close(1011, "upstream failed"), None) + ) + + +async def _upstream_drop_without_close_frame_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: + await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(None, None)) + + +async def _upstream_close_before_any_frame(litellm_logging_obj: LiteLLMLogging) -> None: + await _relay_upstream(litellm_logging_obj, ConnectionClosed(Close(1008, "model not found"), None)) + + +def _frames_then_silence(*frames: bytes | str) -> Callable[[], Awaitable[bytes | str]]: + pending: Final = iter(frames) + + async def _next_frame() -> bytes | str: + frame: Final = next(pending, None) + if frame is None: + await asyncio.Event().wait() + raise AssertionError("unreachable") + return frame + + return _next_frame + + +def _client_ws_sending(*client_frames: str) -> MagicMock: + client_ws: Final = _realtime_client_ws() + client_ws.receive_text = AsyncMock(side_effect=_frames_then_silence(*client_frames)) + return client_ws + + +async def _vertex_access_token() -> str: + return "token" + + +async def _client_frame_the_provider_rejects(litellm_logging_obj: LiteLLMLogging) -> None: + deferred_setup_provider: Final = VertexChirpRealtimeConfig( + resolve_access_token=_vertex_access_token, project="proj-1", location="us" + ) + await _relay_upstream( + litellm_logging_obj, + client_ws=_client_ws_sending( + json.dumps( + { + "type": "session.update", + "session": {"type": "transcription", "audio": {"input": {"format": {"type": "audio/pcmu"}}}}, + } + ) + ), + provider_config=deferred_setup_provider, + upstream_stays_open=True, + ) + + +async def _client_hangs_up_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: + client_ws: Final = _realtime_client_ws() + client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(Close(1000, ""), None)) + await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, client_ws=client_ws, upstream_stays_open=True) + + +async def _refused_handshake(litellm_logging_obj: LiteLLMLogging) -> None: + await close_after_upstream_handshake_refusal(_realtime_client_ws(), 403) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("realtime_session", "counted_successes"), + [ + (_clean_realtime_session, 1), + (_upstream_close_before_any_frame, 0), + (_refused_handshake, 0), + (_upstream_internal_error_after_a_frame, 0), + (_upstream_drop_without_close_frame_after_a_frame, 0), + (_client_frame_the_provider_rejects, 0), + (_client_hangs_up_after_a_frame, 1), + ], + ids=[ + "clean-session", + "upstream-close-before-any-frame", + "refused-handshake", + "upstream-1011-after-a-frame", + "upstream-1006-after-a-frame", + "client-frame-the-provider-rejects", + "client-hangs-up-after-a-frame", + ], +) +async def test_realtime_session_counts_a_router_success_only_when_it_ended_without_a_failure( + realtime_session: Callable[[LiteLLMLogging], Awaitable[None]], counted_successes: int +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "realtime-voice", + "litellm_params": {"model": "openai/gpt-realtime-2.1"}, + "model_info": {"id": "realtime-a"}, + } + ], + num_retries=0, + ) + + async def _arealtime(litellm_logging_obj: LiteLLMLogging, **kwargs: object) -> None: + await realtime_session(litellm_logging_obj) + + await router._ageneric_api_call_with_fallbacks( + model="realtime-voice", original_function=_arealtime, litellm_logging_obj=_realtime_session_logging() + ) + + assert get_deployment_successes_for_current_minute(router, "realtime-a") == counted_successes + + +@pytest.mark.asyncio +@pytest.mark.parametrize("realtime_model", ["bedrock/amazon.nova-sonic-v1:0", "openai/gpt-realtime-2.1"]) +@pytest.mark.parametrize(("clean_sessions", "cooled_down"), [(10, False), (3, True)]) +async def test_realtime_cooldown_reads_the_failure_rate_of_the_sessions_a_deployment_served( + realtime_model: str, clean_sessions: int, cooled_down: bool +) -> None: + import time + + provider, _, provider_model = realtime_model.partition("/") + router: Final = Router( + model_list=[ + { + "model_name": "realtime-voice", + "litellm_params": {"model": realtime_model, "order": order}, + "model_info": {"id": deployment_id}, + } + for order, deployment_id in ((1, "realtime-a"), (2, "realtime-b")) + ], + num_retries=0, + ) + drop: Final = litellm.ServiceUnavailableError( + message="realtime provider stream ended before the session closed", + llm_provider=provider, + model=provider_model, + ) + + def realtime_session(failure: Exception | None): + async def _arealtime(litellm_logging_obj: LiteLLMLogging, **kwargs: object) -> None: + if failure is not None: + raise failure + await _clean_realtime_session(litellm_logging_obj) + + return _arealtime + + for _ in range(clean_sessions): + await router._ageneric_api_call_with_fallbacks( + model="realtime-voice", + original_function=realtime_session(None), + litellm_logging_obj=_realtime_session_logging(), + ) + + async def drop_session() -> bool: + with pytest.raises(litellm.ServiceUnavailableError): + await router._ageneric_api_call_with_fallbacks( + model="realtime-voice", + original_function=realtime_session(drop), + litellm_logging_obj=_realtime_session_logging(), + ) + return router.deployment_callback_on_failure( + kwargs={ + "exception": drop, + "litellm_params": {"metadata": {"model_group": "realtime-voice"}, "model_info": {"id": "realtime-a"}}, + }, + completion_response=None, + start_time=time.time(), + end_time=time.time(), + ) + + drop_cooldowns: Final = [await drop_session() for _ in range(5)] + + assert any(drop_cooldowns) is cooled_down + + async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]: routed: Final[dict[str, object]] = {} @@ -1497,6 +1742,51 @@ def test_set_model_group_info(model_list, user_facing_model_group_name): assert resp.model_group == user_facing_model_group_name +def test_set_model_group_info_reports_supports_continuation_when_any_deployment_has_it(): + router: Final = Router( + model_list=[ + { + "model_name": "long-decode", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, + }, + { + "model_name": "long-decode", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake"}, + "model_info": {"supports_continuation": True}, + }, + { + "model_name": "long-decode", + "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}, + }, + { + "model_name": "plain", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, + }, + ] + ) + long_decode: Final = router._set_model_group_info( + model_group="long-decode", user_facing_model_group_name="long-decode" + ) + plain: Final = router._set_model_group_info(model_group="plain", user_facing_model_group_name="plain") + assert long_decode is not None and long_decode.supports_continuation is True + assert plain is not None and plain.supports_continuation is False + + +@pytest.mark.parametrize("value", ["maybe", "", ["x"], "true", 1]) +def test_set_model_group_info_counts_only_a_boolean_true_supports_continuation(value: object) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "long-decode", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, + "model_info": {"supports_continuation": value}, + }, + ] + ) + group: Final = router._set_model_group_info(model_group="long-decode", user_facing_model_group_name="long-decode") + assert group is not None and group.supports_continuation is False + + @pytest.mark.asyncio async def test_set_response_headers(model_list): """Test if the 'set_response_headers' function is working correctly""" diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 91f0b2339c7..96519a0acbe 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -585,7 +585,10 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) -@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params"]) +@pytest.mark.parametrize( + "flag", + ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params", "supports_continuation"], +) def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: """A deployment's capability override must not land on the shared ``{provider}/{model}`` key, or every sibling deployment of that model would inherit it.""" @@ -672,7 +675,7 @@ def test_pdf_opt_out_on_one_deployment_does_not_block_its_sibling(respx_mock: re def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: - """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on + """A deployment's ``token_count_model`` and framing settings must not land on the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them.""" from litellm.proxy.common_utils.token_count_settings import cost_map_entry @@ -681,6 +684,9 @@ def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: configured: Final = { "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"}, "token_count_framing": 7, + "token_count_tool_framing": 3, + "token_count_responses_framing": 5, + "token_count_responses_tool_framing": 2, } builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured} assert all(builtin[field] != value for field, value in configured.items()), builtin diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5cf872f4ea7..c1c2593c4ef 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -108,6 +108,7 @@ from litellm.utils import ( client, get_non_default_completion_params, get_optional_params_image_gen, + get_optional_params_transcription, get_prompt_cache_min_tokens, get_supported_openai_params, get_token_count, @@ -476,6 +477,34 @@ def test_supports_video_input_reads_its_own_cost_map_flag( assert litellm.get_model_info(model)["supports_video_input"] is video_input +@pytest.mark.parametrize("continuation", [True, False]) +def test_supports_continuation_reads_its_own_cost_map_flag( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, continuation: bool +) -> None: + model: Final = "gemini/continuation-flag-model" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "litellm_provider": "gemini", + "mode": "chat", + "supports_continuation": continuation, + "supports_url_context": not continuation, + }, + ) + assert litellm.supports_continuation(model) is continuation + assert litellm.get_model_info(model).get("supports_continuation") is continuation + + +def test_supports_continuation_is_unset_when_the_cost_map_omits_it( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + model: Final = "gemini/no-continuation-flag-model" + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "gemini", "mode": "chat"}) + assert litellm.supports_continuation(model) is False + assert litellm.get_model_info(model).get("supports_continuation") is None + + def test_get_optional_params_image_gen(): from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig @@ -939,21 +968,11 @@ _COST_MAP_ROW_SCHEMA: Final = { "cache_creation_input_token_cost": {"type": "number"}, "cache_creation_input_token_cost_above_1hr": {"type": "number"}, "cache_creation_input_token_cost_above_12hr": {"type": "number"}, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": { - "type": "number" - }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": {"type": "number"}, "cache_creation_input_token_cost_above_12hr_flex": {"type": "number"}, "cache_creation_input_token_cost_above_12hr_priority": {"type": "number"}, "cache_creation_input_token_cost_above_12hr_batches": {"type": "number"}, @@ -1178,6 +1197,9 @@ _COST_MAP_ROW_SCHEMA: Final = { "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, "token_count_framing": {"type": "integer"}, + "token_count_tool_framing": {"type": "integer"}, + "token_count_responses_framing": {"type": "integer"}, + "token_count_responses_tool_framing": {"type": "integer"}, "token_count_model": { "type": "object", "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, @@ -1189,6 +1211,7 @@ _COST_MAP_ROW_SCHEMA: Final = { "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, "supports_bedrock_converse_tool_result_images": {"type": "boolean"}, "supports_url_context": {"type": "boolean"}, + "supports_continuation": {"type": "boolean"}, "supports_multimodal": {"type": "boolean"}, "uses_embed_content": {"type": "boolean"}, "supports_reasoning": {"type": "boolean"}, @@ -1247,6 +1270,9 @@ _COST_MAP_ROW_SCHEMA: Final = { "/v1/images/generations", "/v1/realtime", "/v1/realtime/transcription_sessions", + "/v1/realtime/translations", + "/v1/realtime/translations/client_secrets", + "/v1/realtime/translations/calls", "/v1/images/variations", "/v1/images/edits", "/v1/batch", @@ -1639,6 +1665,38 @@ def test_get_provider_rerank_config(): assert isinstance(config, HostedVLLMRerankConfig) +def test_get_provider_video_config_resolves_a_custom_provider_map_base_video_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A `BaseVideoConfig` registered in `custom_provider_map` answers for its provider name, a `CustomLLM` + registered there does not, and a built-in name still resolves to its built-in config even when a `CustomLLM` + shadows it in the map.""" + from litellm.llms.custom_llm import CustomLLM + from litellm.llms.openai.videos.transformation import OpenAIVideoConfig + + class QueueVideoConfig(OpenAIVideoConfig): + pass + + video_config: Final = QueueVideoConfig() + monkeypatch.setattr( + litellm, + "custom_provider_map", + [ + {"provider": "queue-video", "custom_handler": video_config}, + {"provider": "image-only", "custom_handler": CustomLLM()}, + {"provider": "openai", "custom_handler": CustomLLM()}, + ], + ) + monkeypatch.setattr(litellm, "_custom_providers", ["queue-video", "image-only", "openai"]) + + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="queue-video") is video_config + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="image-only") is None + assert isinstance( + ProviderConfigManager.get_provider_video_config(model="sora-2", provider="openai"), OpenAIVideoConfig + ) + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="not-registered") is None + + def test_get_provider_text_to_speech_config_vertex_gemini_skips_cloud_tts(): """Regression for LIT-6501: mapping vertex Gemini TTS params through Google Cloud TTS dropped response_format before the speech_to_completion bridge could honor it.""" @@ -7540,9 +7598,10 @@ async def test_oidc_google_api_key_fails_closed_when_the_metadata_server_fails_a side_effect=[metadata_failure(), httpx.Response(200, text=recovered_token)] ) chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) - with caplog.at_level(logging.WARNING, logger="LiteLLM"), pytest.raises( - (litellm.AuthenticationError, litellm.ServiceUnavailableError) - ) as excinfo: + with ( + caplog.at_level(logging.WARNING, logger="LiteLLM"), + pytest.raises((litellm.AuthenticationError, litellm.ServiceUnavailableError)) as excinfo, + ): await _ping_vllm_through(call_path, "vllm-sa") sent_while_failing: Final = chat.call_count await _ping_vllm_through(call_path, "vllm-sa") @@ -7791,9 +7850,7 @@ async def test_router_resolves_an_oidc_api_key_only_on_a_proxy_level_deployment( try: with respx.mock(assert_all_called=True) as upstream: upstream.post(_MISTRAL_CHAT_URL).respond(200, json=_MISTRAL_REPLY) - await router.acompletion( - model="devstral", messages=_PING, metadata={"user_api_key_team_id": "team-1"} - ) + await router.acompletion(model="devstral", messages=_PING, metadata={"user_api_key_team_id": "team-1"}) sent: Final = upstream.calls.last.request.headers["Authorization"] finally: router.discard() @@ -7861,7 +7918,9 @@ async def test_an_inline_oidc_api_key_reaches_the_endpoint_its_deployment_config with respx.mock(assert_all_called=False) as upstream: upstream.route().respond(200, json=_BEDROCK_REPLY) await router.acompletion( - model="claude", messages=_PING, aws_bedrock_runtime_endpoint="https://BEDROCK-runtime.us-east-1.amazonaws.com" + model="claude", + messages=_PING, + aws_bedrock_runtime_endpoint="https://BEDROCK-runtime.us-east-1.amazonaws.com", ) sent: Final = tuple( (call.request.url.host, call.request.headers["Authorization"]) for call in upstream.calls @@ -8070,7 +8129,9 @@ async def test_a_team_row_naming_an_oidc_credential_reaches_only_the_credentials ( lambda reference: LiteLLM_Params( **litellm.Router( - model_list=[{"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": reference}}] + model_list=[ + {"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": reference}} + ] ).get_model_list(model_name="openai/*")[0]["litellm_params"] ), lambda reference: "Bearer listing-token", @@ -8619,6 +8680,40 @@ def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): assert exc_info.value.model == MODEL +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +@pytest.mark.parametrize("mode", ["auto", "required"]) +def test_validate_tool_choice_allowed_tools_chat_layout_passes_through(mode: str): + tool_choice: Final = { + "type": "allowed_tools", + "allowed_tools": { + "mode": mode, + "tools": [{"type": "function", "function": {"name": "get_time"}}, {"type": "mcp", "server_label": "docs"}], + }, + } + assert validate_chat_completion_tool_choice(tool_choice, model=MODEL) is tool_choice + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "get_time"}]}, + {"type": "allowed_tools"}, + {"type": "allowed_tools", "allowed_tools": {"mode": "required"}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": []}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "none", "tools": [{"type": "function"}]}}, + {"type": "allowed_tools", "allowed_tools": {"tools": [{"type": "function", "function": {"name": "a"}}]}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": "get_time"}}, + ], +) +def test_validate_tool_choice_allowed_tools_malformed_is_a_400(tool_choice: dict[str, object]): + with pytest.raises(litellm.BadRequestError, match="allowed_tools") as exc_info: + validate_chat_completion_tool_choice(tool_choice, model=MODEL) + assert exc_info.value.status_code == 400 + assert exc_info.value.model == MODEL + assert '"mode": "auto" | "required"' in str(exc_info.value) + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize("tool_choice", [123, []]) def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): @@ -9042,7 +9137,9 @@ def test_basic_trimming_no_max_tokens_specified(): trimmed_messages = trim_messages(messages, model="gpt-4") print("trimmed messages for gpt-4") print(trimmed_messages) - assert (get_token_count(messages=trimmed_messages, model="gpt-4")) <= litellm.model_cost["gpt-4"]["max_input_tokens"] + assert (get_token_count(messages=trimmed_messages, model="gpt-4")) <= litellm.model_cost["gpt-4"][ + "max_input_tokens" + ] def test_multiple_messages_trimming(): @@ -9284,6 +9381,23 @@ def test_get_valid_models_does_not_ask_a_timed_out_provider_again_on_the_next_li assert (listings, catalog.call_count) == (([], []), 1) +@respx.mock +def test_get_valid_models_does_not_ask_a_timed_out_hosted_vllm_server_again_on_the_next_listing( + reset_mock_cache: None, +) -> None: + models_route: Final = respx.get("http://vllm-hung.test:8000/v1/models").mock( + side_effect=httpx.ReadTimeout("vLLM did not answer") + ) + deployment: Final = LiteLLM_Params(model="hosted_vllm/*", api_base="http://vllm-hung.test:8000/v1") + + listings: Final = tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="hosted_vllm", litellm_params=deployment) + for _ in range(2) + ) + + assert (listings, models_route.call_count) == (([], []), 1) + + @respx.mock def test_get_valid_models_lists_a_provider_again_right_after_a_fast_failure(reset_mock_cache: None) -> None: respx.get("https://openrouter.ai/api/v1/models/user").mock( @@ -10985,6 +11099,49 @@ def test_get_provider_audio_transcription_config_resolves_for_every_provider() - assert isinstance(configs[LlmProviders.OPENAI], litellm.OpenAIWhisperAudioTranscriptionConfig) + +@pytest.mark.parametrize( + ("custom_llm_provider", "model"), + [ + ("groq", "whisper-large-v3"), + ("deepgram", "nova-3"), + ("mistral", "voxtral-mini-latest"), + ("elevenlabs", "scribe_v1"), + ("fireworks_ai", "whisper-v3"), + ], +) +def test_openai_sdk_transcription_params_pass_through_to_other_providers(custom_llm_provider: str, model: str) -> None: + hints: Final = {"stream": True, "languages": ["en"], "keywords": ["LiteLLM"]} + + params: Final = get_optional_params_transcription(model=model, custom_llm_provider=custom_llm_provider, **hints) + + forwarded: Final = {**params.get("extra_body", {}), **params} + assert {key: forwarded.get(key) for key in hints} == hints + + +@respx.mock +def test_drop_params_keeps_openai_sdk_transcription_params_out_of_a_deepgram_request() -> None: + route: Final = respx.post(url__startswith="https://api.deepgram.com/v1/listen").mock( + return_value=httpx.Response( + 200, + json={"metadata": {"duration": 1.0}, "results": {"channels": [{"alternatives": [{"transcript": "hi"}]}]}}, + ) + ) + + litellm.transcription( + model="deepgram/nova-3", + file=("sample.wav", b"audio"), + api_key="dg-test", + punctuate=True, + stream=True, + languages=["en"], + keywords=["LiteLLM"], + drop_params=True, + ) + + assert dict(route.calls.last.request.url.params) == {"model": "nova-3", "punctuate": "true"} + + @pytest.mark.parametrize( ("model", "optional_params", "expected_response_format"), [ @@ -11276,6 +11433,26 @@ def test_vertex_llama3_model_info_resolves_without_the_meta_prefix(): assert prefixed_info["key"] == f"vertex_ai/{model}", model +def test_the_catalog_with_register_model_entries_lists_each_model_once_and_leaves_router_registrations_out() -> None: + from litellm.utils import catalog_with_runtime_registrations + + catalog: Final = catalog_with_runtime_registrations() + catalog_size: Final = len(catalog) + catalogued_model: Final = next(iter(catalog)) + catalogued_entry: Final = dict(catalog[catalogued_model]) + litellm.register_model({catalogued_model: {"supports_prompt_cache_breakpoint": True}}) + litellm.register_model({"qa-registered-only-model": {"litellm_provider": "openai", "mode": "chat"}}) + litellm.register_model({"qa-router-deployment-id": {"mode": "chat"}}, persist_across_reloads=False) + + merged: Final = catalog_with_runtime_registrations() + listed: Final = list(merged) + + assert len(listed) == len(merged) == catalog_size + 1 + assert "qa-registered-only-model" in listed + assert "qa-router-deployment-id" not in listed + assert dict(merged[catalogued_model]) == {**catalogued_entry, "supports_prompt_cache_breakpoint": True} + + @pytest.mark.parametrize( "key", ("openai.test-converse", "bedrock/openai.test-converse", "bedrock/converse/openai.test-converse") ) @@ -11344,7 +11521,9 @@ def test_declared_mid_conversation_system_support_prefers_the_deployment_over_th ) -def test_declared_mid_conversation_system_support_is_none_when_the_model_lookup_raises(monkeypatch: pytest.MonkeyPatch) -> None: +def test_declared_mid_conversation_system_support_is_none_when_the_model_lookup_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setitem(litellm.model_cost, "hosted_vllm/unreadable-probe", "not a model info entry") litellm.get_model_info.cache_clear() assert litellm.utils.declared_mid_conversation_system_support("unreadable-probe", "hosted_vllm", None) is None @@ -11365,7 +11544,6 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_ assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} - @pytest.mark.parametrize( ("provider", "expected_warnings"), (("streamlake", 1), ("veniceai", 1), ("my-own-llm", 0)), diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index e5459a5a0cd..571451164d9 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -1,11 +1,15 @@ -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Sequence +from itertools import chain from typing import Final import pytest +from openai.types.completion import Completion from litellm.exceptions import APIConnectionError +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper -from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, TextCompletionResponse, Usage def test_reasoning_content_preserved_in_text_completion_wrapper(): @@ -85,6 +89,146 @@ async def test_text_completion_stream_usage_is_opt_in( ] +def _completion_chunk(text: str, finish_reason: str | None, **choice_extras: object) -> Completion: + return Completion.construct( + id="cmpl-1", + created=1, + model="m", + object="text_completion", + choices=[{"index": 0, "text": text, "finish_reason": finish_reason, "logprobs": None, **choice_extras}], + ) + + +def _text_completion_stream(upstream: Sequence[Completion]) -> list[TextCompletionResponse]: + logging_obj: Final = Logging( + model="m", + messages=[{"role": "user", "content": "1, 2, 3, 4"}], + stream=True, + call_type="text_completion", + start_time=0.0, + litellm_call_id="call-1", + function_id="function-1", + ) + chat_stream: Final = CustomStreamWrapper( + completion_stream=iter(upstream), + model="m", + custom_llm_provider="text-completion-openai", + logging_obj=logging_obj, + ) + return list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m")) + + +@pytest.mark.parametrize("final_text", [", 5, 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text: str) -> None: + responses: Final = _text_completion_stream( + (_completion_chunk(", 5", None), _completion_chunk(final_text, "stop", stop_reason=", 7")) + ) + + assert responses[-1].choices[0].model_dump(exclude_none=True) == { + "finish_reason": "stop", + "index": 0, + "stop_reason": ", 7", + } + + +def test_text_completion_stream_puts_stop_reason_on_the_final_text_chunk() -> None: + responses: Final = _text_completion_stream( + (_completion_chunk(", 5", None), _completion_chunk(", 6", "stop", stop_reason=", 7")) + ) + + assert [(response.choices[0].get("text"), response.choices[0].get("stop_reason")) for response in responses] == [ + (", 5", None), + (", 6", ", 7"), + (None, ", 7"), + ] + + +@pytest.mark.parametrize("final_text", [", 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_held_back_token_ids_and_the_stop_reason(final_text: str) -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk("", None, token_ids=[11]), + _completion_chunk("", None, token_ids=[220]), + _completion_chunk(", 5", None, token_ids=[20]), + _completion_chunk(final_text, "stop", stop_reason=", 7", token_ids=[11, 220, 21]), + ) + ) + choices: Final = [response.choices[0] for response in responses] + + assert list(chain.from_iterable(choice.get("token_ids") or () for choice in choices)) == [11, 220, 20, 11, 220, 21] + assert [choice.get("stop_reason") for choice in choices if choice.finish_reason is not None] == [", 7"] + + +def test_text_completion_stream_sends_empty_text_on_token_id_only_chunks() -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk("", None, token_ids=[11]), + _completion_chunk("", None, token_ids=[220]), + _completion_chunk(", 5", None, token_ids=[20]), + ) + ) + + choices: Final = [response.choices[0] for response in responses] + + assert [(choice.text, choice.get("token_ids")) for choice in choices if choice.get("token_ids")] == [ + ("", [11]), + ("", [220]), + (", 5", [20]), + ] + + +def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own_index() -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk(" 5", None, index=0, token_ids=[20]), + _completion_chunk(" 7", None, index=1, token_ids=[22]), + _completion_chunk(", 6", None, index=0, token_ids=[11, 21]), + _completion_chunk("", None, index=1, token_ids=[11]), + _completion_chunk("", "stop", index=0, stop_reason=", 9", token_ids=[11, 220]), + _completion_chunk(", 8", "length", index=1, token_ids=[220, 23]), + ) + ) + choices: Final = tuple(chain.from_iterable(response.choices for response in responses)) + + assert { + index: ( + "".join(choice.text or "" for choice in choices if choice.index == index), + list(chain.from_iterable(choice.get("token_ids") or () for choice in choices if choice.index == index)), + ) + for index in (0, 1) + } == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])} + + +@pytest.mark.parametrize( + "upstream", + [ + ( + _completion_chunk(" a", None, index=0), + _completion_chunk(" b", None, index=1), + _completion_chunk("", "length", index=2), + _completion_chunk(" red", "stop", index=1, stop_reason=","), + _completion_chunk(" s", "stop", index=0, stop_reason=","), + ), + ( + _completion_chunk(" a", None, index=0), + _completion_chunk(" b", None, index=1), + _completion_chunk(" s", "stop", index=0, stop_reason=","), + _completion_chunk("", "length", index=2), + _completion_chunk(" red", "stop", index=1, stop_reason=","), + ), + ], + ids=["other-choices-finish-first", "choice-0-finishes-first"], +) +def test_text_completion_stream_finishes_each_choice_on_its_own_index(upstream: Sequence[Completion]) -> None: + finished: Final = [ + (choice.index, choice.finish_reason, choice.get("stop_reason")) + for choice in chain.from_iterable(response.choices for response in _text_completion_stream(upstream)) + if choice.finish_reason is not None + ] + + assert sorted(finished) == [(0, "stop", ","), (1, "stop", ","), (2, "length", None)] + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model") diff --git a/tests/unit/types/llms/test_types_llms_openai.py b/tests/unit/types/llms/test_types_llms_openai.py index 05801c9feec..8c37e86500b 100644 --- a/tests/unit/types/llms/test_types_llms_openai.py +++ b/tests/unit/types/llms/test_types_llms_openai.py @@ -16,8 +16,10 @@ from litellm.types.llms.openai import ( OpenAIChatCompletionLogprobsContent, OpenAIModerationResponse, OpenAIVideoObject, + OutputTextDeltaEvent, ResponseCompletedEvent, ResponsesAPIResponse, + ResponsesAPIStreamEvents, ) @@ -816,3 +818,32 @@ def test_rebuilding_a_dumped_responses_answer_keeps_the_tool_search_call_keys(ex rebuilt: Final = ResponsesAPIResponse(**answer.model_dump()) assert rebuilt.model_dump()["output"] == [search_call] + + +def _text_delta() -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="msg_1", output_index=0, content_index=0, delta="hi" + ) + + +def test_stream_event_dict_reads_see_a_field_set_after_construction(): + event: Final = _text_delta().model_copy(update={"sequence_number": 3}) + + assert event["sequence_number"] == 3 + assert event.get("sequence_number") == 3 + assert "sequence_number" in event + assert dict(event.items())["sequence_number"] == 3 + assert event["delta"] == "hi" + + +def test_stream_event_dict_reads_see_an_undeclared_field_named_keys(): + assert _text_delta().model_copy(update={"keys": "value"})["keys"] == "value" + + +def test_stream_event_dict_reads_still_miss_absent_keys(): + event: Final = _text_delta() + + assert event.get("sequence_number", "fallback") == "fallback" + assert "sequence_number" not in event + with pytest.raises(KeyError): + event["sequence_number"] diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 86a75ec65b9..d71791b6f11 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -270,6 +270,7 @@ INTERNAL_STATE_NAMES: Final = ( "github_copilot_user_session", "litellm_trusted_callback_vars", "_litellm_addressed_response_id", + "_litellm_addressed_comparison_response_id", "_litellm_strip_stream_usage", "_litellm_stream_usage_injected", "client_side_timeout", diff --git a/tests/unit/videos/test_main.py b/tests/unit/videos/test_main.py index 22e1e5c05eb..76e9b062476 100644 --- a/tests/unit/videos/test_main.py +++ b/tests/unit/videos/test_main.py @@ -32,7 +32,7 @@ reflect production. from contextlib import ExitStack from dataclasses import dataclass -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import MagicMock, patch import pytest @@ -205,9 +205,7 @@ def test_video_edit__dispatch_and_provider_from_id(seams): def test_video_extension__dispatch_and_provider_from_id(seams): - result = videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" - ) + result = videos_main.video_extension(video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5") seams.assert_only("video_extension_handler") assert result is seams.handler.video_extension_handler.return_value @@ -285,9 +283,7 @@ def test_generation__mock_response_short_circuits(seams): def test_list__mock_response_short_circuits(seams): - resp = videos_main.video_list( - mock_response=[{"id": "v1", "object": "video", "status": "completed"}] - ) + resp = videos_main.video_list(mock_response=[{"id": "v1", "object": "video", "status": "completed"}]) assert isinstance(resp, list) assert resp[0].id == "v1" @@ -318,7 +314,7 @@ def test_get_character__mock_response_short_circuits(seams): def test_unsupported_provider_raises_without_dispatch(seams): seams.get_config.return_value = None - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.BadRequestError, match="video status is not supported for azure"): videos_main.video_status(video_id=AZURE_VIDEO_ID) seams.handler.video_status_handler.assert_not_called() @@ -333,9 +329,7 @@ def test_unsupported_provider_raises_without_dispatch(seams): async def test_avideo_generation__delegates_with_async_flag(): sentinel = VideoObject(id="v-async", object="video", status="queued") with ( - patch.object( - videos_main, "video_generation", MagicMock(return_value=sentinel) - ) as sync, + patch.object(videos_main, "video_generation", MagicMock(return_value=sentinel)) as sync, patch.object( litellm, "get_llm_provider", @@ -352,9 +346,7 @@ async def test_avideo_generation__delegates_with_async_flag(): @pytest.mark.asyncio async def test_avideo_status__delegates_untouched(): sentinel = VideoObject(id="v-async", object="video", status="queued") - with patch.object( - videos_main, "video_status", MagicMock(return_value=sentinel) - ) as sync: + with patch.object(videos_main, "video_status", MagicMock(return_value=sentinel)) as sync: result = await videos_main.avideo_status(video_id="video_plain") assert result is sentinel @@ -367,9 +359,7 @@ async def test_avideo_content__pre_decodes_provider_before_delegating(): """avideo_content resolves the provider from the encoded id itself before handing off, so the sync worker receives the decoded provider, not None.""" sentinel = b"mp4-bytes" - with patch.object( - videos_main, "video_content", MagicMock(return_value=sentinel) - ) as sync: + with patch.object(videos_main, "video_content", MagicMock(return_value=sentinel)) as sync: result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) assert result is sentinel @@ -393,9 +383,7 @@ DB_YAML_CREDS = { CREDENTIAL_OPERATIONS = [ ( "video_generation_handler", - lambda: videos_main.video_generation( - prompt="p", model="sora-2", **DB_YAML_CREDS - ), + lambda: videos_main.video_generation(prompt="p", model="sora-2", **DB_YAML_CREDS), ), ( "video_status_handler", @@ -407,21 +395,15 @@ CREDENTIAL_OPERATIONS = [ ), ( "video_remix_handler", - lambda: videos_main.video_remix( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), + lambda: videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS), ), ( "video_edit_handler", - lambda: videos_main.video_edit( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), + lambda: videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS), ), ( "video_extension_handler", - lambda: videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS - ), + lambda: videos_main.video_extension(video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS), ), ( "video_list_handler", @@ -429,9 +411,7 @@ CREDENTIAL_OPERATIONS = [ ), ( "video_create_character_handler", - lambda: videos_main.video_create_character( - name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS - ), + lambda: videos_main.video_create_character(name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS), ), ( "video_get_character_handler", @@ -453,3 +433,33 @@ def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] + + +def test_video_generation__custom_provider_map_base_video_config_is_dispatched(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.openai.videos.transformation import OpenAIVideoConfig + + class QueueVideoConfig(OpenAIVideoConfig): + pass + + video_config: Final = QueueVideoConfig() + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "queue-video", "custom_handler": video_config}]) + monkeypatch.setattr(litellm, "_custom_providers", ["queue-video"]) + handler: Final = MagicMock(spec=BaseLLMHTTPHandler) + with patch.object(videos_main, "base_llm_http_handler", handler): + videos_main.video_generation(model="queue-video/seedance-1", prompt="a cat", api_base="http://backend.test") + kw: Final = handler.video_generation_handler.call_args.kwargs + assert kw["video_generation_provider_config"] is video_config + assert kw["custom_llm_provider"] == "queue-video" + assert kw["model"] == "seedance-1" + + +def test_video_generation__custom_provider_without_a_video_config_is_a_400(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.custom_llm import CustomLLM + + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "image-only", "custom_handler": CustomLLM()}]) + monkeypatch.setattr(litellm, "_custom_providers", ["image-only"]) + handler: Final = MagicMock(spec=BaseLLMHTTPHandler) + with patch.object(videos_main, "base_llm_http_handler", handler): + with pytest.raises(litellm.BadRequestError, match="video generation is not supported for image-only"): + videos_main.video_generation(model="image-only/seedance-1", prompt="a cat") + handler.video_generation_handler.assert_not_called() diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index d4ac60e501b..0cadc13a213 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -215,9 +215,6 @@ }, "max-lines": { "count": 1 - }, - "react/no-unescaped-entities": { - "count": 2 } }, "src/app/(dashboard)/guardrails/_components/content_filter/CompetitorIntentConfiguration.tsx": { @@ -247,6 +244,9 @@ } }, "src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx": { + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 6 }, @@ -281,6 +281,9 @@ "local/no-complex-jsx-arrow": { "count": 1 }, + "max-lines": { + "count": 1 + }, "max-params": { "count": 1 }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx index 11d75c0f750..f0529121b21 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx @@ -1,4 +1,5 @@ import { useAccessGroupDetails } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails"; +import type { TFunction } from "i18next"; import { ArrowLeftIcon, BotIcon, @@ -10,6 +11,7 @@ import { UsersIcon, } from "lucide-react"; import { useState } from "react"; +import { useTranslation } from "react-i18next"; import DefaultProxyAdminTag from "@/components/common_components/DefaultProxyAdminTag"; import { BadgeLink } from "@/components/shared/BadgeLink"; import CopyButton from "@/components/shared/CopyButton"; @@ -34,11 +36,11 @@ const MAX_PREVIEW = 5; const shortId = (id: string) => (id.length > 20 ? `${id.slice(0, 10)}...${id.slice(-6)}` : id); -const describeAllowedTools = (tools: readonly string[] | undefined): string => { +const describeAllowedTools = (tools: readonly string[] | undefined, t: TFunction): string => { if (tools === undefined) { - return "All tools"; + return t("accessGroups.detail.toolPermissions.allTools"); } - return tools.length === 0 ? "No tools" : tools.join(", "); + return tools.length === 0 ? t("accessGroups.detail.toolPermissions.noTools") : tools.join(", "); }; function MCPToolPermissionList({ @@ -48,16 +50,17 @@ function MCPToolPermissionList({ servers: readonly AccessGroupResource[]; toolPermissions: Readonly> | null; }) { + const { t } = useTranslation(); if (toolPermissions === null || servers.length === 0) { return null; } return ( -
+
{servers.map(({ id, name }) => { return (
{name ?? id} - {describeAllowedTools(toolPermissions[id])} + {describeAllowedTools(toolPermissions[id], t)}
); })} @@ -106,6 +109,7 @@ function ResourceBadge({ } export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailProps) { + const { t } = useTranslation(); const { data: accessGroup, isLoading } = useAccessGroupDetails(accessGroupId); const [isEditModalVisible, setIsEditModalVisible] = useState(false); const [showAllKeys, setShowAllKeys] = useState(false); @@ -124,10 +128,16 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr if (!accessGroup) { return (
- -

Access group not found

+

{t("accessGroups.detail.notFound")}

); } @@ -146,47 +156,50 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr
-

{accessGroup.access_group_name}

- ID: {accessGroup.access_group_id} - + + {t("accessGroups.detail.idLabel")} + {accessGroup.access_group_id} + +
- Group Details + {t("accessGroups.detail.groupDetails")}
-
Description
+
{t("accessGroups.detail.description")}
{accessGroup.description || "—"}
-
Created
+
{t("accessGroups.detail.created")}
{new Date(accessGroup.created_at).toLocaleString()} {accessGroup.created_by && ( <> - by + {t("accessGroups.detail.by")} )}
-
Last Updated
+
{t("accessGroups.detail.lastUpdated")}
{new Date(accessGroup.updated_at).toLocaleString()} {accessGroup.updated_by && ( <> - by + {t("accessGroups.detail.by")} )} @@ -200,13 +213,15 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr - Attached Keys + {t("accessGroups.detail.attachedKeys")} {keys.length} {keys.length > MAX_PREVIEW && ( )} @@ -219,7 +234,7 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr ))}
) : ( -

No keys attached

+

{t("accessGroups.detail.noKeys")}

)} @@ -228,13 +243,15 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr - Attached Teams + {t("accessGroups.detail.attachedTeams")} {teams.length} {teams.length > MAX_PREVIEW && ( )} @@ -247,7 +264,7 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr ))} ) : ( -

No teams attached

+

{t("accessGroups.detail.noTeams")}

)} @@ -259,37 +276,37 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr - Models + {t("accessGroups.detail.tabs.models")} {models.length} - MCP Servers + {t("accessGroups.detail.tabs.mcpServers")} {mcpServers.length} - Agents + {t("accessGroups.detail.tabs.agents")} {agents.length} - Skills + {t("accessGroups.detail.tabs.skills")} {skills.length} - + - + - + - + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx index 43eaf11b60d..35089234d06 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx @@ -2,6 +2,7 @@ import { BotIcon, InfoIcon, LayersIcon, ServerIcon, SparklesIcon } from "lucide-react"; import type { UseFormReturn } from "react-hook-form"; +import { useTranslation } from "react-i18next"; import { z } from "zod"; import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; @@ -16,9 +17,10 @@ import SkillSelector from "@/components/skills/SkillSelector"; import { Input } from "@/components/ui/input"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Textarea } from "@/components/ui/textarea"; +import i18n from "@/i18n/i18n"; export const accessGroupFormSchema = z.object({ - name: z.string().min(1, "Please enter the access group name"), + name: z.string().min(1, { error: () => i18n.t("accessGroups.form.name.required") }), description: z.string(), modelIds: z.array(z.string()), mcpServerIds: z.array(z.string()), @@ -80,6 +82,7 @@ export function AccessGroupBaseForm({ onTabChange, initialServerIds, }: AccessGroupBaseFormProps) { + const { t } = useTranslation(); const { accessToken } = useAuthorized(); const { data: agentsData } = useAgents(); const { data: mcpServersData } = useMCPServers(); @@ -98,56 +101,61 @@ export function AccessGroupBaseForm({ - General Info + {t("accessGroups.form.tabs.general")} - Models + {t("accessGroups.form.tabs.models")} - MCP Servers + {t("accessGroups.form.tabs.mcpServers")} - Agents + {t("accessGroups.form.tabs.agents")} - Skills + {t("accessGroups.form.tabs.skills")} - + {({ ref, ...field }) => ( - + )} - + {({ ref, ...field }) => ( -