From cf9b5e4fa75ad5e18d3a803ee0b0d0a792df3534 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 16 May 2026 18:31:43 -0700 Subject: [PATCH 01/11] [Infra] Bump versions (#28094) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * bump: version 0.1.40 → 0.1.41 * bump: version 1.85.0 → 1.86.0 * add uv lock --- enterprise/pyproject.toml | 4 ++-- pyproject.toml | 6 +++--- uv.lock | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 9698e7912d4..9f37b52d94c 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.40" +version = "0.1.41" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.40" +version = "0.1.41" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/pyproject.toml b/pyproject.toml index a397c10ce66..f63770105dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.85.0" +version = "1.86.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -57,7 +57,7 @@ proxy = [ "azure-storage-blob==12.28.0", "mcp==1.26.0", "litellm-proxy-extras==0.4.72", - "litellm-enterprise==0.1.40", + "litellm-enterprise==0.1.41", "RestrictedPython==8.1", "rich==13.9.4", "polars==1.38.1", @@ -251,7 +251,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.85.0" +version = "1.86.0" version_files = [ "pyproject.toml:^version", ] diff --git a/uv.lock b/uv.lock index dbf4d85f536..f3eaf6ca88c 100644 --- a/uv.lock +++ b/uv.lock @@ -3189,7 +3189,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.85.0" +version = "1.86.0" source = { editable = "." } dependencies = [ { name = "aiohttp" }, @@ -3534,7 +3534,7 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.40" +version = "0.1.41" source = { editable = "enterprise" } [[package]] From fe63650ebd0945933f080d1f79fd73db45b341dd Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 18 May 2026 08:56:25 -0700 Subject: [PATCH 02/11] fix(proxy): gate team allowed_passthrough_routes to proxy admins (#28097) * fix(proxy): gate team allowed_passthrough_routes to proxy admins allowed_passthrough_routes short-circuits the role-based route gate, so the keys endpoints already restrict it to proxy admins. The team writers (/team/new, /team/update) had no equivalent check, letting an org admin (a non-proxy-admin who clears the route gate and _verify_team_access) self-grant pass-through routes on their team. Lift the keys check into a shared helper and apply it to both team endpoints. Resolves LIT-3019 * docs(proxy): note view-only admins are intentionally excluded from passthrough gate Clarifies the proxy-admin guard per review feedback; no behavior change. Refs LIT-3019 --- .../management_endpoints/common_utils.py | 32 ++++++ .../key_management_endpoints.py | 31 +----- .../management_endpoints/team_endpoints.py | 9 ++ .../test_team_endpoints.py | 100 ++++++++++++++++++ 4 files changed, 142 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index a43d15a580f..dc27e87726a 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from fastapi import HTTPException, status +from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -53,6 +54,37 @@ def require_caller_user_id_for_non_admin( return user_api_key_dict.user_id +def _check_passthrough_routes_caller_permission( + data: BaseModel, + user_api_key_dict: UserAPIKeyAuth, + *, + entity: str = "key", +) -> None: + """ + Only proxy admins may set `allowed_passthrough_routes` (top-level or under + `metadata`) — it short-circuits the role-based route gate, so keys and teams + must be gated identically. + """ + # view-only admins excluded by design; blocked upstream from writes anyway + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + if getattr(data, "allowed_passthrough_routes", None): + raise HTTPException( + status_code=403, + detail={ + "error": f"Only proxy admins can set `allowed_passthrough_routes` on a {entity}." + }, + ) + metadata = getattr(data, "metadata", None) + if isinstance(metadata, dict) and metadata.get("allowed_passthrough_routes"): + raise HTTPException( + status_code=403, + detail={ + "error": f"Only proxy admins can set `metadata.allowed_passthrough_routes` on a {entity}." + }, + ) + + def _is_user_team_admin( user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable ) -> bool: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7ff706eb91e..2ab147043d9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -55,6 +55,7 @@ from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_k from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, _is_user_org_admin_for_team, _is_user_team_admin, _set_object_metadata_field, @@ -548,36 +549,6 @@ def _check_allowed_routes_caller_permission( ) -def _check_passthrough_routes_caller_permission( - data: BaseModel, - user_api_key_dict: UserAPIKeyAuth, -) -> None: - """ - Only proxy admins may set `allowed_passthrough_routes` on a key, either at - the top level of the request or nested under `metadata`. - - The route gate evaluates passthrough access ahead of the standard role - gate, so the field is restricted to admins to keep that ordering safe. - """ - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return - if getattr(data, "allowed_passthrough_routes", None): - raise HTTPException( - status_code=403, - detail={ - "error": "Only proxy admins can set `allowed_passthrough_routes` on a key." - }, - ) - metadata = getattr(data, "metadata", None) - if isinstance(metadata, dict) and metadata.get("allowed_passthrough_routes"): - raise HTTPException( - status_code=403, - detail={ - "error": "Only proxy admins can set `metadata.allowed_passthrough_routes` on a key." - }, - ) - - async def validate_team_id_used_in_service_account_request( team_id: Optional[str], prisma_client: Optional[PrismaClient], diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 35e3d196e9e..86c4d6dcd9a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -73,6 +73,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, _is_user_org_admin_for_team, _is_user_team_admin, _set_object_metadata_field, @@ -1049,6 +1050,10 @@ async def new_team( # noqa: PLR0915 Member(role="admin", user_id=user_api_key_dict.user_id) ) + _check_passthrough_routes_caller_permission( + data, user_api_key_dict, entity="team" + ) + ## ADD TO MODEL TABLE _model_id = None if data.model_aliases is not None and isinstance(data.model_aliases, dict): @@ -1646,6 +1651,10 @@ async def update_team( # noqa: PLR0915 user_api_key_dict=user_api_key_dict, ) + _check_passthrough_routes_caller_permission( + data, user_api_key_dict, entity="team" + ) + if data.soft_budget is not None: max_budget_to_check = ( data.max_budget diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 5c7bbc46c95..b450a262907 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -7943,3 +7943,103 @@ async def test_team_member_me_returns_404_for_unknown_team(mock_db_client): user_api_key_dict=caller_auth, ) assert exc_info.value.status_code == 404 + + +def _non_admin_auth(): + return UserAPIKeyAuth( + user_id="u-team-admin", user_role=LitellmUserRoles.INTERNAL_USER + ) + + +def test_check_passthrough_routes_caller_permission_team(): + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, + ) + + admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + non_admin = _non_admin_auth() + + _check_passthrough_routes_caller_permission( + NewTeamRequest(allowed_passthrough_routes=["/foo/*"]), admin, entity="team" + ) + + _check_passthrough_routes_caller_permission( + NewTeamRequest(), non_admin, entity="team" + ) + _check_passthrough_routes_caller_permission( + NewTeamRequest(allowed_passthrough_routes=[]), non_admin, entity="team" + ) + + with pytest.raises(HTTPException) as exc: + _check_passthrough_routes_caller_permission( + NewTeamRequest(allowed_passthrough_routes=["/admin/*"]), + non_admin, + entity="team", + ) + assert exc.value.status_code == 403 + assert "allowed_passthrough_routes" in str(exc.value.detail) + assert "team" in str(exc.value.detail) + + with pytest.raises(HTTPException) as exc: + _check_passthrough_routes_caller_permission( + NewTeamRequest(metadata={"allowed_passthrough_routes": ["/admin/*"]}), + non_admin, + entity="team", + ) + assert exc.value.status_code == 403 + assert "metadata.allowed_passthrough_routes" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_new_team_blocks_non_admin_passthrough_routes(mock_db_client): + """A non-proxy-admin cannot self-grant pass-through routes via /team/new.""" + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with patch( + "litellm.proxy.management_endpoints.team_endpoints._check_user_team_limits", + AsyncMock(return_value=None), + ): + with pytest.raises(ProxyException) as exc: + await new_team( + data=NewTeamRequest( + team_alias="t", allowed_passthrough_routes=["/admin/*"] + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=_non_admin_auth(), + ) + assert str(exc.value.code) == "403" + assert "allowed_passthrough_routes" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client): + """Even a team manager (non-proxy-admin) cannot set pass-through routes via + /team/update — the gate runs after _verify_team_access.""" + from fastapi import Request + + from litellm.proxy._types import ProxyException, UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import update_team + + existing = MagicMock() + existing.model_dump.return_value = {"team_id": "t1"} + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints._verify_team_access", + AsyncMock(return_value=None), + ): + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest( + team_id="t1", allowed_passthrough_routes=["/admin/*"] + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=_non_admin_auth(), + ) + assert str(exc.value.code) == "403" + assert "allowed_passthrough_routes" in str(exc.value.message) From bb448b0031c8b5fc20db52383291798fd5ea001d Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 18 May 2026 09:15:39 -0700 Subject: [PATCH 03/11] fix(tests): stabilize image-edit VCR cassettes to stop live gpt-image-1 spend (#28110) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(tests): stabilize image-edit VCR cassettes to stop live gpt-image-1 spend The image-edit cassettes for ``gpt-image-1`` were accumulating >50 episodes and being refused by the persister (``tests/_vcr_redis_persister.py``), so every CI run was hitting the real OpenAI endpoint. The async parametrize was the clearest tell: ``test_openai_image_edit_litellm_sdk[True]`` cached to 1 entry, but the ``[False]`` (async) sibling grew to 51 entries and never replayed. Two non-deterministic sources were fueling the growth, both fixed here. After this patch, the cassettes settle at one episode per unique call and replay for the 24-hour TTL like every other suite. 1. Pin httpx's multipart boundary at the source. The existing ``_normalize_multipart_boundary`` rewrites the boundary in the ``Content-Type`` header reliably, but on the async transport path the body is not always a contiguous ``bytes`` object when ``before_record_request`` runs, so the body-side replacement silently no-ops and the recorded cassette retains the random ``boundary=`` string. The next CI run gets a fresh random boundary, the ``safe_body`` matcher misses, and ``record_mode="new_episodes"`` appends another episode. Wrapping ``httpx._multipart.MultipartStream.__init__`` so it always uses ``vcr-static-boundary`` when no boundary is supplied eliminates the variance for both sync and async paths and leaves the normalizer in place as a backstop. Exposed as ``pin_httpx_multipart_boundary`` so other multipart-heavy suites (audio, ocr, batches) can adopt the same fixture later. 2. Pass raw ``bytes`` (not ``BytesIO`` streams) through the image-edit fixtures. A ``BytesIO`` whose file pointer is at EOF after the first multipart upload silently encodes an empty image on the next SDK / Router retry — yet another divergent body that VCR records as a new episode. ``bytes`` are immutable and position-less, so retries re-encode an identical payload every time. This is also a small production-correctness improvement: a customer passing ``BytesIO`` today would hit the same empty-body retry bug. The BytesIO-specific smoke test (``test_openai_image_edit_with_bytesio``) is preserved by giving ``get_test_images_as_bytesio`` its own factory instead of aliasing the bytes one. 3. Add ``scripts/flush_image_edit_vcr_cassettes.py`` — a one-shot Redis SCAN/DEL helper that clears the bloated pre-fix cassettes under ``litellm:vcr:cassette:tests/image_gen_tests/test_image_edits/*``. Without this, the next CI run still loads the existing 51-entry cassette, the new fixed-boundary body still doesn't match any of the stale entries, the persister still refuses to save, and the bleed continues. Run once with the production ``CASSETTE_REDIS_URL`` after merge (dry-run by default). * DIAGNOSTIC: log VCR body mismatches + per-episode body hashes Temporary observability boost so we can root-cause why ``test_image_edits.py`` async parametrizes still record fresh episodes on every CI run even though the multipart boundary is now pinned (sync parametrizes cache cleanly as VCR HIT). The matcher currently raises ``AssertionError("request bodies differ")`` with zero context, so we cannot tell whether the live body genuinely varies, the matcher is comparing a bytes object to a stream object, or the normalizer is silently skipping the body because it is not bytes/str. Three logs added; the first two are worth keeping permanently, the third is intended to be reverted after the diagnosis lands: 1. ``_safe_body_matcher`` now emits a structured stderr block on mismatch (type of each side, length, SHA-256, first divergent byte offset, ±100-byte window). Always-on -- mismatches are signal, not noise, and the existing per-test verdict already logs once per test. PERMANENT. 2. ``_normalize_multipart_boundary`` now logs to stderr when the body type is not bytes/bytearray/str -- the silent ``else: return`` branch was masking exactly the case we suspect is firing on async (httpx ``MultipartStream`` handed to vcrpy before the body is read). PERMANENT. 3. ``_RedisPersister.save_cassette`` now logs every episode's body SHA-256, length, and 120-byte preview at save time. This lets two consecutive CI runs be diffed: if the same test records a different hash run-to-run, the live body genuinely varies; if both runs record the same hash but the matcher still misses, the bug is in the matcher itself. TEMPORARY -- revert once the async variance is identified and fixed. Once a single ``image_gen_testing`` CI run produces these logs, revert this commit (or just the persister hash block) with a force push so the cassette save path is not noisy in steady-state. * DIAGNOSTIC: route VCR diagnostics through per-PID files (bypass xdist capture) Re-push of the diagnostic logging from the previous commit, this time wired so the output actually survives to the CI log. xdist captures stdout/stderr from every passing test in the worker process; the body-matcher and normalizer-skip diagnostics fire from inside vcrpy machinery during the test, so for any test that ultimately passes (which is all of them once the cassettes are recorded), the diagnostic lines are silently swallowed. Fix: write each diagnostic line to a per-PID file under ``test-results/vcr-diagnostics/.log`` instead of writing to stderr. The controller's ``pytest_terminal_summary`` aggregates those files and writes them through ``terminalreporter.write_line``, which is not subject to per-test capture. As a bonus, ``test-results/`` is already collected by the ``store_test_results`` step in CircleCI, so the raw per-worker logs survive as build artifacts even after the test session ends. Three call sites updated: 1. ``_emit_body_mismatch_diagnostic`` (matcher) -- writes the structured type/length/sha/window block via ``vcr_diag_write_line``. 2. ``_normalize_multipart_boundary`` -- logs the silent-skip path (body not bytes/bytearray/str) the same way. 3. ``_maybe_log_episode_body_hashes`` (persister) -- replaces the ``_log.warning`` calls (which the root-logger config also swallows in CI) with ``vcr_diag_write_line``. Image-gen conftest is the only suite wired to dump the aggregated log at session end. Other suites can opt in by adding ``emit_vcr_diagnostic_log(terminalreporter)`` to their own ``pytest_terminal_summary``. The diagnostic dir is cleared at the start of each session (controller-only) so a local rerun does not mix output from prior runs. Same revert plan as the previous diagnostic commit: keep the matcher + normalizer skip diagnostics permanently (they only fire on signal events), revert the persister body-hash dump once the async variance is identified. * fix(tests): coalesce iterable request bodies before matching/recording Root cause of the residual async image-edit cassette leak. The diagnostic run for ``ba3915d9`` printed: [vcr-safe-body-matcher] request body mismatch body[a]: type='list_iterator' length=unknown sha256=N/A body[b]: type='list_iterator' length=unknown sha256=N/A httpx's async transport hands vcrpy a ``request.body`` that is a ``list_iterator`` over multipart chunks rather than a contiguous ``bytes`` blob. Two consequences: 1. ``_safe_body_matcher`` compares the two iterator objects with ``==``, which is identity comparison for arbitrary iterators - semantically identical multipart bodies never compare equal, and ``record_mode="new_episodes"`` appends a new episode on every CI run until the cassette crosses ``MAX_EPISODES_PER_CASSETTE`` and the persister refuses to save (this is exactly what the OVERFLOW warning has been catching). 2. ``_normalize_multipart_boundary`` short-circuits its ``else: return`` branch because the body is neither bytes nor str, so any residual random boundary characters in the body bytes are never rewritten. Sync requests do not hit this code path: httpx's sync transport hands vcrpy a single ``bytes`` body, so ``==`` works and the boundary normalizer runs as intended. That is why ``test_openai_image_edit_litellm_sdk[True]`` records to ``entries=1`` and replays cleanly while ``[False]`` (async) kept growing by one episode per run. Fix: add ``_materialize_iterable_body`` which coalesces an iterable ``request.body`` into ``bytes`` in-place. Call it from two places: * The top of ``_before_record_request``, so the boundary normalizer and the cassette serializer both see bytes from then on. * The top of ``_safe_body_matcher``, as defense in depth in case a future vcrpy code path invokes the matcher without first going through ``_before_record_request``. The vcrpy ``Request`` is a wrapper used for matching and recording; the underlying httpx transport sends its own request body separately, so replacing the iterator on the vcrpy wrapper does not starve the live HTTP send. After this lands the async parametrizes should flip from ``[VCR MISS:RECORDED] entries=N+1`` to ``[VCR HIT] entries=N`` on the next CI run, matching the sync side and dropping the residual ~$3/day to $0. * fix(tests): handle bytes_iterator + never leave an exhausted body Follow-up to 8e08272b. The previous attempt at coalescing iterable request bodies bailed out (``return`` without writing ``request.body``) whenever it could not classify the chunk type. That was the wrong failure mode for one critical case: vcrpy sometimes presents the body as ``iter(some_bytes)``, whose Python type is ``bytes_iterator`` and which yields ``int`` byte values (0-255), not byte chunks. The old code saw an ``int`` chunk, hit the ``else: return`` branch, and left ``request.body`` pointing at the now-exhausted iterator. The post-fix diagnostic run made this loud: [vcr-safe-body-matcher] request body mismatch body[a]: type='bytes_iterator' length=unknown sha256=N/A body[b]: type='bytes_iterator' length=unknown sha256=N/A Every async image-edit test then ballooned from entries=2 to entries=10 in that single CI run -- the exhausted iterator meant the live multipart upload went out as an empty body, OpenAI returned 400, the SDK + flaky retries fired, each retry got a fresh iterator that my hook exhausted again, and ``new_episodes`` recorded each failed attempt as a new cassette episode. This patch: * Recognizes ``bytes_iterator`` (chunks are ``int``) and reconstructs the buffer via ``bytes(chunks)``. * Keeps the existing ``list_iterator``-over-bytes-chunks handling via ``b"".join(...)``. * **Always writes a bytes value back to ``request.body`` after consuming the iterator.** If the chunk shape is unrecognized, ``request.body`` is set to ``b""`` rather than left as an exhausted iterator. That is wrong in the sense of "we lost the body" but right in the sense of "the failure mode is now visible (live API call sends empty body and fails fast) instead of invisible (corrupt cassette grows silently)". Combined with the matcher diagnostic, any future regression in this code path will surface in the CI log immediately. Local verification covers ``bytes_iterator``, ``list_iterator`` over bytes chunks, generator over bytes chunks, empty iterator, already-bytes (idempotent), identical-content iterator equality in the matcher (now matches), and differing-content iterator inequality (still raises). * fix(tests): clear vcrpy's sticky _was_iter flag so materialized bodies stay bytes Actual root cause of the async image-edit cassette leak. The previous diagnostic run produced this dead giveaway: [vcr-episode-body-hash] ... episode[0]: body type='bytes_iterator' is not bytes/bytearray/str -- cannot hash [vcr-safe-body-matcher] request body mismatch body[a]: type='bytes_iterator' length=unknown sha256=N/A body[b]: type='bytes_iterator' length=unknown sha256=N/A Both sides of the matcher were ``bytes_iterator`` **after** the materializer had supposedly converted them to bytes. That made no sense until I read vcrpy's ``Request`` class. vcrpy's ``Request`` keeps two private flags that are set in ``__init__`` from the original body's type and **never cleared by the setter**: def __init__(self, method, uri, body, headers): self._was_file = hasattr(body, "read") self._was_iter = _is_nonsequence_iterator(body) ... @property def body(self): if self._was_file: return BytesIO(self._body) if self._was_iter: return iter(self._body) return self._body @body.setter def body(self, value): if isinstance(value, str): value = value.encode("utf-8") self._body = value # <-- does NOT touch _was_iter / _was_file So when httpx's async transport hands vcrpy an iterator body, ``_was_iter`` becomes ``True`` and stays there forever. Even after ``_materialize_iterable_body`` writes plain bytes via ``request.body = out``, the next read of ``.body`` re-wraps the stored bytes in ``iter()`` -- producing a fresh ``bytes_iterator`` that compares unequal to any other ``bytes_iterator`` via object identity. The matcher missed every time, the cassette grew by one episode per run, and the persister saw the same iterator type when trying to hash the body for the diagnostic log. Fix: after writing the materialized bytes, also force ``_was_iter`` and ``_was_file`` to ``False``. vcrpy exposes no public API for this, so we touch the private flags directly -- acknowledged as a pragmatic test-only hack with a clear unit boundary (the only call site is ``_materialize_iterable_body``). Local repro reproduces the exact production setup: ``Request('POST', url, iter(b'multipart-content'), {})`` on two sides, runs the matcher, asserts HIT. Verified the matcher hits on identical content and still raises on differing content. Should be the last fix needed. Existing cassettes that contain oddly-shaped bodies (lists of int chunks, etc. from the previous ``_was_iter=True`` save path) still match because the materializer canonicalises both sides to bytes before comparison -- no fourth re-flush required. * revert(tests): drop the temp per-episode body-hash diagnostic Removed now that 1c51ad13 has confirmed the root cause (vcrpy's sticky ``_was_iter`` flag making the body getter re-wrap stored bytes in ``iter()`` on every access). The hash dump did its job -- the post-1c51ad13 image_gen_testing run shows all five async image-edit tests as ``[VCR HIT]`` with stable entry counts and zero billing errors -- and is too noisy to keep on by default (over 100 lines per session at steady state). Kept permanently: * ``_safe_body_matcher`` mismatch diagnostic in ``_vcr_conftest_common.py``. Only fires on a body mismatch, which is signal worth surfacing whenever it happens. * ``_normalize_multipart_boundary`` "skipped" log line. Same rationale -- only fires when the body shape is something the normalizer cannot rewrite in place. * The ``test-results/vcr-diagnostics/.log`` per-PID file plumbing (``vcr_diag_write_line`` / ``emit_vcr_diagnostic_log``). Useful for any future diagnostic that needs to bypass xdist stdout/stderr capture; cheap to keep. * chore(tests): delete unused flush script + wire VCR diagnostic dump everywhere * Remove ``scripts/flush_image_edit_vcr_cassettes.py``. It was a one-shot helper for the initial cassette flush; the iterator and ``_was_iter`` fixes mean no future flush should be required, and the script was never run anywhere (the actual flushes happened inside the CI conftest via the temp hacks that have since been reverted). * The matcher mismatch + normalizer skip diagnostics already write per-PID files for every suite that imports the shared VCR plumbing, but ``emit_vcr_diagnostic_log`` -- the controller-side dump that surfaces those files into the CI log at session end -- was only wired into ``image_gen_tests``. Add the one-line call to the 12 sibling conftests that already use VCR so the diagnostics surface in any suite's terminal output if a body matcher ever misses. No new output in steady state -- the dump is a no-op when no diagnostics were recorded that session. * chore(tests): trim non-essential comments per project comment policy Strips docstrings, inline comments, and block comments that this PR introduced where the code itself was already self-evident. Keeps the few lines that document non-obvious behaviour (raw-bytes-not-BytesIO rationale on the image fixtures, the per-PID-files-bypass-xdist note on the diagnostic directory). Touches only comments this PR added -- no pre-existing comment is removed. Net: -161 lines of comment/docstring across 3 files, no code behaviour change. * chore(tests): forward **kwargs in pin_httpx_multipart_boundary wrapper Defensive against future httpx MultipartStream.__init__ adding new optional kwargs. Without the forward, the wrapper would silently drop them. No behaviour change today. * chore(tests): canonicalize VCR matchers and surface shouldn't-happen branches Bundles the "follow-up cleanup PR" into this one so it does not get lost. Four small changes: 1. Introduce ``_canonical_body(req) -> (bytes, pre_type)`` and route ``_safe_body_matcher`` through it. The matcher now operates on bytes by construction; the "compare two iterator objects via ``==`` and silently get object-identity semantics" failure mode (which cost us this entire PR to diagnose) is structurally impossible to reintroduce. ``pre_type`` is the body type *before* canonicalization, surfaced by the mismatch diagnostic so a future regression involving a new body shape is still visible. 2. Add a structured diagnostic to ``_key_fingerprint_matcher``. It was previously raising a bare ``AssertionError("API key fingerprints differ")`` with zero context -- exactly the anti-pattern the body matcher had before this PR. 3. Surface "shouldn't-happen" branches via ``vcr_diag_write_line``: * ``_strip_image_b64_payloads`` -- logs when ``response``, ``response['body']``, or ``response['body']['string']`` arrives in an unexpected shape (vcrpy contract violation). * ``_compute_key_fingerprint`` -- logs the ``"no-key"`` fallback with the request method/URL so a stripped-auth-header bug is visible instead of masked. * ``_canonical_body`` -- logs its own empty-bytes fallback when a body has a shape ``_materialize_iterable_body`` did not handle. 4. Re-introduce per-episode body-hash logging in ``_RedisPersister.save_cassette`` (was reverted in 927c5548 as "noisy"). Quantified cost: ~25 KB of CI log per session at peak, ~ms-scale CPU, zero output in steady state (no save = no log). Trade-off favours keeping it: lets two consecutive CI runs be diffed by body hash, which is how we will spot the next regression in the same class. All call sites still work: local repro confirms iter==iter HIT, iter!=iter raises, plain-bytes HIT, body-hash log emits via the same per-PID file plumbing as the matcher diagnostics. * chore(tests): symmetrize diag-log cleanup across every VCR-using conftest ``image_gen_tests/conftest.py`` was the only suite that cleared ``test-results/vcr-diagnostics/*.log`` at session start. The other 12 VCR-using conftests inherited any stale per-PID logs from a previous local run and would dump them in the terminal summary -- harmless in CI (fresh container) but confusing locally when running multiple suites in sequence. Extracts the cleanup into a ``reset_vcr_diag_dir`` helper in ``tests/_vcr_conftest_common.py`` and calls it from every VCR-using conftest's ``pytest_configure``. Same single source of truth, no inline duplication. * fix(tests): gate body materialization on __next__ and strip PR comments aiohttp/vcrpy stores the json kwarg as a dict; _materialize_iterable_body was iterating it via __iter__ and joining the keys, replacing the request body with concatenated key names ("textlanguageentities"). Gate on __next__ so containers (dict/list/tuple) are left alone — only single-use iterators like httpx's bytes_iterator / list_iterator are materialized. Log diagnostic line when chunk type is unrecognized. * fix(tests): JSON-encode dict bodies in canonical_body for stable matching aiohttp stubs store the json kwarg as a dict; the fallback that compared all dicts as b"" caused concurrent presidio analyze calls to be served the wrong cassette episode. JSON-encode with sort_keys for stable bytes. * fix(tests): guard emit_vcr_diagnostic_log against multi-conftest re-emission Co-authored-by: Yassin Kortam * fix(tests): globalize multipart-boundary pin + stabilize whisper fixtures Diagnostic shows audio_testing was silently re-recording 50+ live Whisper episodes per CI run (over MAX_EPISODES_PER_CASSETTE, so the persister refused to save). Two changes: * Move the session-autouse _pin_multipart_boundary fixture into the shared _vcr_conftest_common module so every VCR-using suite picks it up via a single import. image_gen had it inline; the other 12 suites silently lacked it. * Replace the module-level open("rb") audio file handles in test_whisper with cached bytes + a per-call (filename, bytes, mimetype) tuple, mirroring the image_edits raw-bytes pattern. Stops the file-pointer- at-EOF bug where the second test got an empty multipart body. * chore(tests): drop per-episode body-hash dump and redundant emit guard --------- Co-authored-by: shin-berri Co-authored-by: yuneng-jiang Co-authored-by: Cursor Agent Co-authored-by: Yassin Kortam --- tests/_vcr_conftest_common.py | 269 ++++++++++++++++++-- tests/audio_tests/conftest.py | 7 +- tests/audio_tests/test_whisper.py | 43 ++-- tests/guardrails_tests/conftest.py | 7 +- tests/image_gen_tests/conftest.py | 7 +- tests/image_gen_tests/test_image_edits.py | 37 +-- tests/litellm_utils_tests/conftest.py | 7 +- tests/llm_responses_api_testing/conftest.py | 7 +- tests/llm_translation/conftest.py | 7 +- tests/local_testing/conftest.py | 7 +- tests/logging_callback_tests/conftest.py | 7 +- tests/ocr_tests/conftest.py | 7 +- tests/pass_through_unit_tests/conftest.py | 7 +- tests/router_unit_tests/conftest.py | 7 +- tests/search_tests/conftest.py | 7 +- tests/unified_google_tests/conftest.py | 7 +- 16 files changed, 366 insertions(+), 74 deletions(-) diff --git a/tests/_vcr_conftest_common.py b/tests/_vcr_conftest_common.py index a179a21ba69..cb43f1abbdd 100644 --- a/tests/_vcr_conftest_common.py +++ b/tests/_vcr_conftest_common.py @@ -36,6 +36,75 @@ SAFE_BODY_MATCHER_NAME = "safe_body" KEY_FINGERPRINT_MATCHER_NAME = "key_fingerprint" KEY_FINGERPRINT_HEADER = "x-litellm-key-fp" +VCR_DIAG_DIR_ENV = "LITELLM_VCR_DIAG_DIR" +VCR_DIAG_DIR_DEFAULT = "test-results/vcr-diagnostics" + + +def _vcr_diag_dir() -> str: + return os.environ.get(VCR_DIAG_DIR_ENV) or VCR_DIAG_DIR_DEFAULT + + +def vcr_diag_write_line(msg: str) -> None: + try: + directory = _vcr_diag_dir() + os.makedirs(directory, exist_ok=True) + path = os.path.join(directory, f"{os.getpid()}.log") + with open(path, "a", encoding="utf-8") as fh: + fh.write(msg.rstrip("\n") + "\n") + except OSError: + pass + + +def reset_vcr_diag_dir() -> None: + if os.environ.get("PYTEST_XDIST_WORKER"): + return + directory = _vcr_diag_dir() + if not os.path.isdir(directory): + return + try: + names = os.listdir(directory) + except OSError: + return + for name in names: + if name.endswith(".log"): + try: + os.remove(os.path.join(directory, name)) + except OSError: + pass + + +def emit_vcr_diagnostic_log(terminalreporter) -> None: + directory = _vcr_diag_dir() + if not os.path.isdir(directory): + return + try: + files = sorted(f for f in os.listdir(directory) if f.endswith(".log")) + except OSError: + return + if not files: + return + terminalreporter.write_sep("=", "VCR DIAGNOSTIC LOG", bold=True) + terminalreporter.write_line( + f" source dir: {directory} (also archived as a CI artifact)" + ) + for name in files: + path = os.path.join(directory, name) + try: + with open(path, "r", encoding="utf-8") as fh: + content = fh.read() + except OSError as exc: + terminalreporter.write_line( + f" [failed to read {name}: {type(exc).__name__}: {exc}]" + ) + continue + if not content.strip(): + continue + terminalreporter.write_sep("-", name, bold=False) + for line in content.splitlines(): + terminalreporter.write_line(line) + terminalreporter.write_sep("=", bold=True) + + # Intentionally narrower than ``FILTERED_REQUEST_HEADERS``: AWS SigV4 headers # carry secrets but their values rotate on every call, so fingerprinting them # would defeat caching. @@ -91,6 +160,32 @@ VCR_IMAGE_B64_PLACEHOLDER = "dGVzdA==" VCR_FIXED_MULTIPART_BOUNDARY = "vcr-static-boundary" +def pin_httpx_multipart_boundary(monkeypatch) -> None: + try: + import httpx._multipart as _httpx_multipart + except ImportError: + return + + _original_init = _httpx_multipart.MultipartStream.__init__ + + def _init_with_fixed_boundary(self, data, files, boundary=None, **kwargs): + if boundary is None: + boundary = VCR_FIXED_MULTIPART_BOUNDARY.encode("ascii") + return _original_init(self, data=data, files=files, boundary=boundary, **kwargs) + + monkeypatch.setattr( + _httpx_multipart.MultipartStream, "__init__", _init_with_fixed_boundary + ) + + +@pytest.fixture(scope="session", autouse=True) +def _pin_multipart_boundary(): + monkeypatch = pytest.MonkeyPatch() + pin_httpx_multipart_boundary(monkeypatch) + yield + monkeypatch.undo() + + def _scrub_response(response): if not isinstance(response, dict): return response @@ -139,9 +234,17 @@ def _strip_image_b64_payloads(response): preserves all those checks while shrinking cassettes by ~99%. """ if not isinstance(response, dict): + vcr_diag_write_line( + f"[vcr-strip-b64] response is {type(response).__name__!r}, not " + "dict; skipping b64 scrub" + ) return response body = response.get("body") if not isinstance(body, dict): + vcr_diag_write_line( + f"[vcr-strip-b64] response['body'] is {type(body).__name__!r}, " + "not dict; skipping b64 scrub" + ) return response raw = body.get("string") if raw is None: @@ -151,12 +254,20 @@ def _strip_image_b64_payloads(response): try: text = bytes(raw).decode("utf-8") except UnicodeDecodeError: + vcr_diag_write_line( + "[vcr-strip-b64] response body bytes are not valid UTF-8; " + "skipping b64 scrub" + ) return response was_bytes = True elif isinstance(raw, str): text = raw was_bytes = False else: + vcr_diag_write_line( + f"[vcr-strip-b64] response['body']['string'] is " + f"{type(raw).__name__!r}, not bytes/str; skipping b64 scrub" + ) return response try: @@ -186,6 +297,35 @@ def _before_record_response(response): return filter_non_2xx_response(_scrub_response(_strip_image_b64_payloads(response))) +def _canonical_body(request) -> tuple[bytes, str]: + pre_type = type(getattr(request, "body", None)).__name__ + _materialize_iterable_body(request) + body = getattr(request, "body", None) + if body is None: + return b"", pre_type + if isinstance(body, bytes): + return body, pre_type + if isinstance(body, bytearray): + return bytes(body), pre_type + if isinstance(body, str): + return body.encode("utf-8"), pre_type + if isinstance(body, (dict, list)): + try: + return ( + json.dumps(body, sort_keys=True, separators=(",", ":")).encode("utf-8"), + pre_type, + ) + except (TypeError, ValueError): + pass + method = getattr(request, "method", "?") + uri = getattr(request, "uri", getattr(request, "url", "?")) + vcr_diag_write_line( + f"[vcr-canonical-body] FALLBACK: {method} {uri} body type " + f"{type(body).__name__!r} not coerced to bytes; comparing as b''" + ) + return b"", pre_type + + def _safe_body_matcher(r1, r2) -> None: """Compare request bodies as bytes; never invokes ``json.loads``. @@ -195,27 +335,47 @@ def _safe_body_matcher(r1, r2) -> None: This matcher is strictly more conservative — the only equivalence it gives up vs. the default is "JSON key order doesn't matter". """ - body1 = getattr(r1, "body", None) - body2 = getattr(r2, "body", None) + body1, pre1 = _canonical_body(r1) + body2, pre2 = _canonical_body(r2) if body1 == body2: return - - def _to_bytes(b): - if b is None: - return b"" - if isinstance(b, bytes): - return b - if isinstance(b, str): - return b.encode("utf-8") - return None - - n1 = _to_bytes(body1) - n2 = _to_bytes(body2) - if n1 is not None and n2 is not None and n1 == n2: - return + _emit_body_mismatch_diagnostic(r1, r2, body1, body2, pre1, pre2) raise AssertionError("request bodies differ") +def _emit_body_mismatch_diagnostic(r1, r2, body1, body2, pre1, pre2) -> None: + def _describe(label, asbytes, pre_type): + return ( + f" {label}: pre_canonical_type={pre_type!r} length={len(asbytes)} " + f"sha256={hashlib.sha256(asbytes).hexdigest()} " + f"preview={asbytes[:120]!r}" + ) + + method_a = getattr(r1, "method", "?") + method_b = getattr(r2, "method", "?") + url_a = getattr(r1, "uri", getattr(r1, "url", "?")) + url_b = getattr(r2, "uri", getattr(r2, "url", "?")) + lines = [ + "[vcr-safe-body-matcher] request body mismatch", + f" request[a]: {method_a} {url_a}", + f" request[b]: {method_b} {url_b}", + _describe("body[a]", body1, pre1), + _describe("body[b]", body2, pre2), + ] + if body1 != body2: + offset = next( + (i for i in range(min(len(body1), len(body2))) if body1[i] != body2[i]), + min(len(body1), len(body2)), + ) + start = max(0, offset - 100) + end_a = min(len(body1), offset + 100) + end_b = min(len(body2), offset + 100) + lines.append(f" first divergent byte offset: {offset}") + lines.append(f" window[a] @ {start}..{end_a}: {body1[start:end_a]!r}") + lines.append(f" window[b] @ {start}..{end_b}: {body2[start:end_b]!r}") + vcr_diag_write_line("\n".join(lines)) + + def _iter_header_values(headers, name: str): if headers is None: return @@ -271,6 +431,13 @@ def _compute_key_fingerprint(request) -> str: stable = _stable_key_value(header_name, text) parts.append(f"{header_name}={stable}") if not parts: + method = getattr(request, "method", "?") + uri = getattr(request, "uri", getattr(request, "url", "?")) + vcr_diag_write_line( + f"[vcr-key-fingerprint] no API key header found on {method} " + f"{uri}; falling back to 'no-key'. If this request should have " + "carried auth, something earlier in the pipeline stripped it." + ) return "no-key" digest = hashlib.sha256("\n".join(parts).encode("utf-8")).hexdigest() return digest[:16] @@ -360,6 +527,13 @@ def _normalize_multipart_boundary(request) -> None: elif isinstance(body, str): new_body = body.replace(current_boundary, VCR_FIXED_MULTIPART_BOUNDARY) else: + vcr_diag_write_line( + f"[vcr-multipart-normalize] body normalization SKIPPED: " + f"body type {type(body).__name__!r} is not bytes/bytearray/str. " + f"content-type={content_type_value!r}. " + f"Recorded body will retain the random boundary substring " + f"and the safe_body matcher will miss on the next run." + ) return try: @@ -389,6 +563,7 @@ def _before_record_request(request): headers = getattr(request, "headers", None) if headers is None: return request + _materialize_iterable_body(request) if not any(_iter_header_values(headers, KEY_FINGERPRINT_HEADER)): fingerprint = _compute_key_fingerprint(request) try: @@ -400,6 +575,56 @@ def _before_record_request(request): return request +def _materialize_iterable_body(request) -> None: + body = getattr(request, "body", None) + if body is None or isinstance(body, (bytes, bytearray, str)): + return + if not hasattr(body, "__next__"): + return + try: + chunks = list(body) + except TypeError: + return + + out = _coalesce_chunks_to_bytes(chunks) + if out is None: + method = getattr(request, "method", "?") + uri = getattr(request, "uri", getattr(request, "url", "?")) + first_type = type(chunks[0]).__name__ if chunks else "empty" + vcr_diag_write_line( + f"[vcr-materialize] FALLBACK: {method} {uri} chunk type " + f"{first_type!r} not coerced to bytes; storing b''" + ) + out = b"" + + try: + request.body = out + except (AttributeError, TypeError): + pass + + for attr in ("_was_iter", "_was_file"): + try: + setattr(request, attr, False) + except (AttributeError, TypeError): + pass + + +def _coalesce_chunks_to_bytes(chunks): + if not chunks: + return b"" + first = chunks[0] + try: + if isinstance(first, int): + return bytes(chunks) + if isinstance(first, (bytes, bytearray)): + return b"".join(c if isinstance(c, bytes) else bytes(c) for c in chunks) + if isinstance(first, str): + return "".join(chunks).encode("utf-8") + except (TypeError, ValueError): + return None + return None + + def _key_fingerprint_matcher(r1, r2) -> None: def _fp(req): for value in _iter_header_values( @@ -410,7 +635,17 @@ def _key_fingerprint_matcher(r1, r2) -> None: return value if isinstance(value, str) else str(value) return "no-key" - if _fp(r1) != _fp(r2): + fp1, fp2 = _fp(r1), _fp(r2) + if fp1 != fp2: + method_a = getattr(r1, "method", "?") + method_b = getattr(r2, "method", "?") + url_a = getattr(r1, "uri", getattr(r1, "url", "?")) + url_b = getattr(r2, "uri", getattr(r2, "url", "?")) + vcr_diag_write_line( + "[vcr-key-fingerprint-matcher] API key fingerprints differ\n" + f" request[a]: {method_a} {url_a} fingerprint={fp1!r}\n" + f" request[b]: {method_b} {url_b} fingerprint={fp2!r}" + ) raise AssertionError("API key fingerprints differ") diff --git a/tests/audio_tests/conftest.py b/tests/audio_tests/conftest.py index ff47853d494..c4ff576e5bd 100644 --- a/tests/audio_tests/conftest.py +++ b/tests/audio_tests/conftest.py @@ -5,14 +5,17 @@ import pytest sys.path.insert(0, os.path.abspath("../..")) -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -44,6 +47,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -57,3 +61,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/audio_tests/test_whisper.py b/tests/audio_tests/test_whisper.py index cdf079f8cb4..243d27614b1 100644 --- a/tests/audio_tests/test_whisper.py +++ b/tests/audio_tests/test_whisper.py @@ -23,12 +23,21 @@ pwd = os.path.dirname(os.path.realpath(__file__)) print(pwd) file_path = os.path.join(pwd, "gettysburg.wav") - -audio_file = open(file_path, "rb") - - file2_path = os.path.join(pwd, "eagle.wav") -audio_file2 = open(file2_path, "rb") + +with open(file_path, "rb") as _f: + _GETTYSBURG_BYTES = _f.read() +with open(file2_path, "rb") as _f: + _EAGLE_BYTES = _f.read() + + +def _audio_file(): + return ("gettysburg.wav", _GETTYSBURG_BYTES, "audio/wav") + + +def _audio_file2(): + return ("eagle.wav", _EAGLE_BYTES, "audio/wav") + load_dotenv() @@ -44,7 +53,7 @@ async def _run_transcription( ): transcript = await litellm.atranscription( model=model, - file=audio_file, + file=_audio_file(), api_key=api_key, api_base=api_base, response_format=response_format, @@ -101,7 +110,7 @@ async def test_transcription_caching(): response_1 = await litellm.atranscription( model="whisper-1", - file=audio_file, + file=_audio_file(), ) await asyncio.sleep(5) @@ -110,7 +119,7 @@ async def test_transcription_caching(): response_2 = await litellm.atranscription( model="whisper-1", - file=audio_file, + file=_audio_file(), ) print("response_1", response_1) @@ -122,7 +131,7 @@ async def test_transcription_caching(): response_3 = await litellm.atranscription( model="whisper-1", - file=audio_file2, + file=_audio_file2(), ) print("response_3", response_3) print("response3 hidden params", response_3._hidden_params) @@ -146,7 +155,7 @@ async def test_whisper_log_pre_call(): with patch.object(custom_logger, "log_pre_api_call") as mock_log_pre_call: await litellm.atranscription( model="whisper-1", - file=audio_file, + file=_audio_file(), ) mock_log_pre_call.assert_called_once() @@ -165,7 +174,7 @@ async def test_whisper_log_pre_call(): with patch.object(custom_logger, "log_pre_api_call") as mock_log_pre_call: await litellm.atranscription( model="whisper-1", - file=audio_file, + file=_audio_file(), ) mock_log_pre_call.assert_called_once() @@ -177,7 +186,7 @@ async def test_gpt_4o_transcribe(): from unittest.mock import patch, MagicMock await litellm.atranscription( - model="openai/gpt-4o-transcribe", file=audio_file, response_format="json" + model="openai/gpt-4o-transcribe", file=_audio_file(), response_format="json" ) @@ -187,7 +196,9 @@ async def test_gpt_4o_transcribe_model_mapping(): # Test GPT-4o mini transcribe response = await litellm.atranscription( - model="openai/gpt-4o-mini-transcribe", file=audio_file, response_format="json" + model="openai/gpt-4o-mini-transcribe", + file=_audio_file(), + response_format="json", ) # Check that the response contains the correct model in hidden params @@ -198,7 +209,7 @@ async def test_gpt_4o_transcribe_model_mapping(): # Test GPT-4o transcribe response2 = await litellm.atranscription( - model="openai/gpt-4o-transcribe", file=audio_file, response_format="json" + model="openai/gpt-4o-transcribe", file=_audio_file(), response_format="json" ) # Check that the response contains the correct model in hidden params @@ -209,7 +220,7 @@ async def test_gpt_4o_transcribe_model_mapping(): # Test traditional whisper-1 still works response3 = await litellm.atranscription( - model="openai/whisper-1", file=audio_file, response_format="json" + model="openai/whisper-1", file=_audio_file(), response_format="json" ) # Check that the response contains the correct model in hidden params @@ -262,7 +273,7 @@ async def test_azure_transcribe_model_mapping(): # Make the transcription call response = await litellm.atranscription( model="azure/whisper-1", - file=audio_file, + file=_audio_file(), response_format="json", api_key="test-api-key", api_base="https://my-endpoint-europe-berri-992.openai.azure.com/", diff --git a/tests/guardrails_tests/conftest.py b/tests/guardrails_tests/conftest.py index eb563699b2b..f2f65645c3d 100644 --- a/tests/guardrails_tests/conftest.py +++ b/tests/guardrails_tests/conftest.py @@ -16,14 +16,17 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -55,6 +58,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -160,3 +164,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/image_gen_tests/conftest.py b/tests/image_gen_tests/conftest.py index 93dec98e708..9f808c11161 100644 --- a/tests/image_gen_tests/conftest.py +++ b/tests/image_gen_tests/conftest.py @@ -9,14 +9,17 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm # noqa: E402,F401 -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -58,6 +61,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -71,3 +75,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 656b8a69117..ca8ec3bbe32 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/tests/image_gen_tests/test_image_edits.py @@ -103,12 +103,6 @@ class BaseLLMImageEditTest(ABC): pwd = os.path.dirname(os.path.realpath(__file__)) -# Image fixtures must be regenerated per access — module-level -# ``open(...)`` handles get consumed after a single multipart upload, leaving -# subsequent tests in the same process to send empty bodies. That non-determinism -# (a) blows the recorded cassette past ``MAX_EPISODES_PER_CASSETTE`` so the -# persister refuses to save (see ``tests/_vcr_redis_persister.py``), and -# (b) re-bills the live image edit endpoint on every CI run. def _read_image_bytes(filename: str) -> bytes: with open(os.path.join(pwd, filename), "rb") as f: return f.read() @@ -119,32 +113,20 @@ _LITELLM_SITE_BYTES = _read_image_bytes("litellm_site.png") def _make_test_images() -> list: - """Return a fresh pair of image streams seeded with the fixture bytes. + return [_ISHAAN_GITHUB_BYTES, _LITELLM_SITE_BYTES] - Use this everywhere you'd previously have used the module-level - ``TEST_IMAGES``. Each call returns brand new ``BytesIO`` objects whose - file pointers start at 0, so multipart uploads encode the full image - bytes on every test invocation. Parametrized and ``flaky``-retried - test methods call ``get_base_image_edit_call_args`` once per - invocation, so a fresh stream per call is sufficient — the factory - must not auto-rewind on EOF or the SDK's multipart writer will read - the same bytes forever (worker OOM). - """ + +def _make_single_test_image() -> bytes: + return _ISHAAN_GITHUB_BYTES + + +def get_test_images_as_bytesio(): return [ BytesIO(_ISHAAN_GITHUB_BYTES), BytesIO(_LITELLM_SITE_BYTES), ] -def _make_single_test_image() -> BytesIO: - return BytesIO(_ISHAAN_GITHUB_BYTES) - - -def get_test_images_as_bytesio(): - """Helper function to get test images as BytesIO objects""" - return _make_test_images() - - class TestOpenAIImageEditGPTImage1(BaseLLMImageEditTest): """ Concrete implementation of BaseLLMImageEditTest for OpenAI image edits. @@ -710,10 +692,9 @@ async def test_multiple_image_edit_with_different_formats(): try: prompt = "Create a cohesive artistic style across all images" - # Test with mixed BytesIO and file objects mixed_images = [ - _make_single_test_image(), # File object - get_test_images_as_bytesio()[1], # BytesIO object + _make_single_test_image(), + get_test_images_as_bytesio()[1], ] result = await aimage_edit( diff --git a/tests/litellm_utils_tests/conftest.py b/tests/litellm_utils_tests/conftest.py index 08745c99c07..418ee76a399 100644 --- a/tests/litellm_utils_tests/conftest.py +++ b/tests/litellm_utils_tests/conftest.py @@ -12,14 +12,17 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm # noqa: E402,F401 -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -86,6 +89,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -116,3 +120,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/llm_responses_api_testing/conftest.py b/tests/llm_responses_api_testing/conftest.py index 2a08db57149..1928b540dad 100644 --- a/tests/llm_responses_api_testing/conftest.py +++ b/tests/llm_responses_api_testing/conftest.py @@ -13,14 +13,17 @@ sys.path.insert( import litellm # noqa: E402 -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -52,6 +55,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -116,3 +120,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/llm_translation/conftest.py b/tests/llm_translation/conftest.py index 5fcd31aa32d..d346dae4308 100644 --- a/tests/llm_translation/conftest.py +++ b/tests/llm_translation/conftest.py @@ -18,14 +18,17 @@ sys.path.insert( import litellm # noqa: E402 -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -73,6 +76,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -82,6 +86,7 @@ def pytest_runtest_logreport(report): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) # --------------------------------------------------------------------------- diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index 0ff7dff668a..6a746041f15 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -22,14 +22,17 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -84,6 +87,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -93,6 +97,7 @@ def pytest_runtest_logreport(report): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) # --------------------------------------------------------------------------- diff --git a/tests/logging_callback_tests/conftest.py b/tests/logging_callback_tests/conftest.py index cdb9200bc83..6dde85f2ca7 100644 --- a/tests/logging_callback_tests/conftest.py +++ b/tests/logging_callback_tests/conftest.py @@ -19,14 +19,17 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -79,6 +82,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -229,3 +233,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/ocr_tests/conftest.py b/tests/ocr_tests/conftest.py index 66970b8579f..94790bd7aa3 100644 --- a/tests/ocr_tests/conftest.py +++ b/tests/ocr_tests/conftest.py @@ -12,14 +12,17 @@ import pytest sys.path.insert(0, os.path.abspath("../..")) -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -51,6 +54,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -64,3 +68,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/pass_through_unit_tests/conftest.py b/tests/pass_through_unit_tests/conftest.py index 42a95343eb7..390e14b7f11 100644 --- a/tests/pass_through_unit_tests/conftest.py +++ b/tests/pass_through_unit_tests/conftest.py @@ -5,14 +5,17 @@ import pytest sys.path.insert(0, os.path.abspath("../..")) -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -56,6 +59,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -71,3 +75,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/router_unit_tests/conftest.py b/tests/router_unit_tests/conftest.py index fe976515c92..6a8f3e589f4 100644 --- a/tests/router_unit_tests/conftest.py +++ b/tests/router_unit_tests/conftest.py @@ -12,14 +12,17 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm # noqa: E402,F401 -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -97,6 +100,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -123,3 +127,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/search_tests/conftest.py b/tests/search_tests/conftest.py index e06d3e95eee..78ba19a7724 100644 --- a/tests/search_tests/conftest.py +++ b/tests/search_tests/conftest.py @@ -13,14 +13,17 @@ import pytest sys.path.insert(0, os.path.abspath("../..")) -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -52,6 +55,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -65,3 +69,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/unified_google_tests/conftest.py b/tests/unified_google_tests/conftest.py index d28f89a77b0..5b4f57b8036 100644 --- a/tests/unified_google_tests/conftest.py +++ b/tests/unified_google_tests/conftest.py @@ -12,14 +12,17 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm # noqa: E402,F401 -from tests._vcr_conftest_common import ( # noqa: E402 +from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, + _pin_multipart_boundary, apply_vcr_auto_marker_to_items, emit_cassette_cache_session_banner, emit_vcr_classification_summary, + emit_vcr_diagnostic_log, install_live_call_probe, record_vcr_outcome, register_persister_if_enabled, + reset_vcr_diag_dir, vcr_config_dict, ) @@ -84,6 +87,7 @@ def _vcr_outcome_gate(request, vcr): def pytest_configure(config): _verbose_state.remember_pluginmanager(config) + reset_vcr_diag_dir() def pytest_runtest_logreport(report): @@ -110,3 +114,4 @@ def pytest_collection_modifyitems(config, items): def pytest_terminal_summary(terminalreporter, exitstatus, config): emit_cassette_cache_session_banner(terminalreporter) emit_vcr_classification_summary(terminalreporter) + emit_vcr_diagnostic_log(terminalreporter) From 8c6625216bbe7790f32775849d428206ff558e06 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Mon, 18 May 2026 12:16:20 -0700 Subject: [PATCH 04/11] fix(bedrock/cohere): send embedding_types as JSON array, not string (#28172) * fix(bedrock/cohere): wrap embedding_types as list in map_openai_params Bedrock Cohere expects embedding_types as a JSON array but encoding_format was passed through as a raw string, causing: Malformed input request: #/embedding_types: expected type: JSONArray, found: String * test(bedrock/cohere): assert embedding_types is sent as JSON array --------- Co-authored-by: Ishaan Jaffer --- .../bedrock/embed/cohere_transformation.py | 2 +- .../bedrock/embed/test_bedrock_embedding.py | 47 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/embed/cohere_transformation.py b/litellm/llms/bedrock/embed/cohere_transformation.py index d00cb74aae0..2c0dc834144 100644 --- a/litellm/llms/bedrock/embed/cohere_transformation.py +++ b/litellm/llms/bedrock/embed/cohere_transformation.py @@ -22,7 +22,7 @@ class BedrockCohereEmbeddingConfig: ) -> dict: for k, v in non_default_params.items(): if k == "encoding_format": - optional_params["embedding_types"] = v + optional_params["embedding_types"] = v if isinstance(v, list) else [v] elif k == "dimensions": optional_params["output_dimension"] = v return optional_params diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index c67a8712340..9955851132c 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -957,3 +957,50 @@ def test_titan_image_embedding_cost_uses_per_image_rate(): assert response.usage is not None assert response.usage.prompt_tokens_details is not None assert response.usage.prompt_tokens_details.image_count == 1 + + +@pytest.mark.parametrize( + "encoding_format,expected_embedding_types", + [ + ("float", ["float"]), + ("base64", ["base64"]), + (["float", "int8"], ["float", "int8"]), + ], +) +def test_bedrock_cohere_embedding_types_wrapped_as_list( + encoding_format, expected_embedding_types +): + """ + Bedrock Cohere expects `embedding_types` as a JSON array, not a raw string. + + Regression test for: Bedrock returns + Malformed input request: #/embedding_types: expected type: JSONArray, found: String + when `encoding_format` is passed as a string. + """ + litellm.set_verbose = True + client = HTTPHandler() + model = "bedrock/cohere.embed-multilingual-v3" + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(cohere_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model=model, + input=test_input, + encoding_format=encoding_format, + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key="test-bearer-token-12345", + ) + + assert isinstance(response, litellm.EmbeddingResponse) + + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + assert "embedding_types" in request_body + assert request_body["embedding_types"] == expected_embedding_types + assert isinstance(request_body["embedding_types"], list) From ce87c411bfb33a8b37acaa630a39e4e4c8685add Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 18 May 2026 15:41:51 -0700 Subject: [PATCH 05/11] fix(tests): migrate realtime + rerank tests off shut-down upstream models (#28191) * fix(tests): use gpt-realtime in realtime guardrails test OpenAI shut down gpt-4o-realtime-preview-2024-12-17 on 2026-05-07, so the live OpenAI realtime guardrails integration test now fails with model_not_found (session.created never arrives, _wait_for_event times out). Point OPENAI_REALTIME_URL at the current GA model, gpt-realtime. Scope limited to this test: the pricing-catalog JSON keeps the retired entries intentionally (historical cost calc + separate Azure timeline), and the Azure realtime cost-calc test is unaffected. * fix(tests): mock nvidia_nim rerank instead of hitting EOL'd endpoint NVIDIA reached end-of-life for the hosted nvidia/llama-3.2-nv-rerankqa-1b-v2 rerank API on 2026-05-18 with no published replacement, so the live BaseLLMRerankTest.test_basic_rerank for nvidia_nim now returns HTTP 410 ("Gone"). NVIDIA's hosted catalog rotates on a schedule, so swapping in another live model would only defer the failure. Override test_basic_rerank in TestNvidiaNim to mock the sync/async HTTP transport (same pattern as test_nvidia_nim_rerank_ranking_endpoint in this file) and inject a fake NVIDIA_NIM_API_KEY via monkeypatch. The request/response transformation and cost calculation stay covered offline. Scope limited to nvidia_nim; other BaseLLMRerankTest providers untouched. * fix(tests): migrate remaining realtime tests off shut-down gpt-4o-realtime-preview OpenAI's 2026-05-07 shutdown removed the entire gpt-4o-realtime-preview family, including the undated 'gpt-4o-realtime-preview' alias (not just the dated snapshot fixed earlier). Three live tests still connected with the dead alias and failed with messages_received=1 (an error event instead of session.created): - test_openai_realtime_simple.py: get_model() -> gpt-realtime (drives TestOpenAIRealtime.test_realtime_connection / test_realtime_with_query_params) - test_openai_realtime.py: test_openai_realtime_direct_call_no_intent and test_openai_realtime_direct_call_with_intent -> openai/gpt-realtime (the with_intent test shares the same dead alias even though it was not in the failing set this run) Mocked unit tests (test_realtime_query_params_construction, test_realtime_query_params_use_normalized_model_name) are left as-is: they never hit the network and assert string plumbing only. Also fixes test_text_message_blocked_by_guardrail_no_ai_response, which now connects (the earlier URL swap worked) but tripped a model-wording-brittle assertion. The guardrail flow asks the model to voice the block message verbatim; gpt-4o-realtime-preview complied (output contained 'blocked'), gpt-realtime refuses verbatim-repeat instructions ('I'm sorry, but I can't repeat that message.'). Since the original user message is blocked before it reaches OpenAI, the refusal is still a safe outcome. Assertion #3 now accepts both voicing and refusal, and adds a hard check that the blocked phrase never leaks into AI output. --- .../realtime/test_openai_realtime.py | 10 +++-- .../realtime/test_openai_realtime_simple.py | 5 ++- .../test_realtime_guardrails_openai.py | 44 ++++++++++++++----- tests/llm_translation/test_nvidia_nim.py | 41 +++++++++++++++++ 4 files changed, 86 insertions(+), 14 deletions(-) diff --git a/tests/llm_translation/realtime/test_openai_realtime.py b/tests/llm_translation/realtime/test_openai_realtime.py index c5f77de6beb..fc9f938b4cd 100644 --- a/tests/llm_translation/realtime/test_openai_realtime.py +++ b/tests/llm_translation/realtime/test_openai_realtime.py @@ -101,7 +101,9 @@ async def test_openai_realtime_direct_call_no_intent(): try: await litellm._arealtime( - model="openai/gpt-4o-realtime-preview", + # OpenAI shut down the gpt-4o-realtime-preview family (incl. the + # undated alias) on 2026-05-07; gpt-realtime is the GA successor. + model="openai/gpt-realtime", websocket=websocket_client, api_key=os.environ.get("OPENAI_API_KEY"), timeout=60, @@ -249,14 +251,16 @@ async def test_openai_realtime_direct_call_with_intent(): websocket_client = RealTimeWebSocketClient() caught_exception = None + # OpenAI shut down the gpt-4o-realtime-preview family (incl. the undated + # alias) on 2026-05-07; gpt-realtime is the GA successor. query_params: RealtimeQueryParams = { - "model": "openai/gpt-4o-realtime-preview", + "model": "openai/gpt-realtime", "intent": "chat", } try: await litellm._arealtime( - model="openai/gpt-4o-realtime-preview", + model="openai/gpt-realtime", websocket=websocket_client, api_key=os.environ.get("OPENAI_API_KEY"), query_params=query_params, diff --git a/tests/llm_translation/realtime/test_openai_realtime_simple.py b/tests/llm_translation/realtime/test_openai_realtime_simple.py index 5522d843e42..073c1ce11af 100644 --- a/tests/llm_translation/realtime/test_openai_realtime_simple.py +++ b/tests/llm_translation/realtime/test_openai_realtime_simple.py @@ -21,7 +21,10 @@ class TestOpenAIRealtime(BaseRealtimeTest): """ def get_model(self) -> str: - return "gpt-4o-realtime-preview" + # OpenAI shut down the entire gpt-4o-realtime-preview family + # (including the undated alias) on 2026-05-07. gpt-realtime is the + # current GA realtime model. + return "gpt-realtime" def get_api_key_env_var(self) -> str: return "OPENAI_API_KEY" diff --git a/tests/llm_translation/realtime/test_realtime_guardrails_openai.py b/tests/llm_translation/realtime/test_realtime_guardrails_openai.py index 170440f6b9f..ec9d73e2d60 100644 --- a/tests/llm_translation/realtime/test_realtime_guardrails_openai.py +++ b/tests/llm_translation/realtime/test_realtime_guardrails_openai.py @@ -26,9 +26,7 @@ from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.types.guardrails import GuardrailEventHooks OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") -OPENAI_REALTIME_URL = ( - "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-12-17" -) +OPENAI_REALTIME_URL = "wss://api.openai.com/v1/realtime?model=gpt-realtime" pytestmark = pytest.mark.skipif( not OPENAI_API_KEY, @@ -192,10 +190,35 @@ async def test_text_message_blocked_by_guardrail_no_ai_response(): len(transcript_deltas) >= 1 ), f"Expected guardrail message in transcript delta, got: {event_types}" - # 3. No *real* AI response should have been generated. - # The guardrail may produce its own response (e.g. "Content blocked: ...") - # via response.cancel + conversation.item.create + response.create. - # We allow the guardrail's own block message but NOT original AI content. + # 3. No *real* AI response to the blocked content should have been + # generated. The original user message is blocked BEFORE it is + # forwarded to OpenAI, so the only thing the model ever sees is the + # guardrail's "say exactly: " prompt + # (see realtime_streaming.py). Two safe outcomes are possible: + # - the model voices the block message verbatim (older realtime + # snapshots did this -> text contains "blocked"), or + # - the model declines to repeat it (gpt-realtime tends to refuse + # verbatim-repeat instructions, e.g. "I'm sorry, but I can't + # repeat that message."). + # Both mean the blocked prompt itself was never answered, so we + # accept either. The hard invariant is that the blocked phrase must + # never leak into AI output, and the model must not have produced a + # normal answer to the user (which would have neither a block nor a + # refusal marker). + safe_markers = ( + "block", + "guardrail", + "content filter", + "policy", + "can't repeat", + "cannot repeat", + "won't repeat", + "can't assist", + "can't help", + "unable to", + "i'm sorry", + "i am sorry", + ) done_events = [e for e in client_events if e.get("type") == "response.done"] for done in done_events: output = done.get("response", {}).get("output", []) @@ -205,11 +228,12 @@ async def test_text_message_blocked_by_guardrail_no_ai_response(): for c in item.get("content", []) ] real_ai_text = " ".join(ai_texts).strip() - # Allow guardrail-generated block messages (contain "Content blocked" or "blocked") if real_ai_text: assert ( - "blocked" in real_ai_text.lower() - or "guardrail" in real_ai_text.lower() + BLOCKED_PHRASE not in real_ai_text + ), f"Blocked phrase leaked into AI response: {real_ai_text!r}" + assert any( + marker in real_ai_text.lower() for marker in safe_markers ), f"AI responded with non-guardrail content even though message was blocked: {real_ai_text!r}" finally: diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 469516407c8..80e764147bb 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -262,3 +262,44 @@ class TestNvidiaNim(BaseLLMRerankTest): def get_expected_cost(self) -> float: """Nvidia NIM rerank models are free (cost = 0.0)""" return 0.0 + + @pytest.mark.asyncio() + @pytest.mark.parametrize("sync_mode", [True, False]) + async def test_basic_rerank(self, sync_mode, monkeypatch): + """ + Override the base live rerank test with a mocked HTTP layer. + + NVIDIA reached end-of-life for the hosted + nvidia/llama-3.2-nv-rerankqa-1b-v2 rerank API on 2026-05-18 and + published no replacement model, so a live call now returns HTTP 410 + ("Gone"). NVIDIA's hosted catalog rotates on a schedule, so pointing + at another live model would only defer the same failure. Mock the + transport instead (same pattern as + test_nvidia_nim_rerank_ranking_endpoint above) so the request/response + transformation and cost calculation stay covered offline. + """ + monkeypatch.setenv("NVIDIA_NIM_API_KEY", "fake-api-key") + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.text = "" + mock_response.json.return_value = { + "rankings": [ + {"index": 0, "logit": 0.95}, + {"index": 1, "logit": 0.75}, + ], + "usage": {"total_tokens": 7}, + } + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + return_value=mock_response, + ), + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=mock_response, + ), + ): + await super().test_basic_rerank(sync_mode=sync_mode) From 477b63c5ead39bb38b804552ccd0aa04b5fde118 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 19 May 2026 04:57:06 +0530 Subject: [PATCH 06/11] fix(caching): replay openai/responses bridge cache hits as chat streams (#28158) * fix(caching): replay openai/responses bridge cache hits as chat streams When chat completions route through openai/responses, cached ModelResponse payloads under aresponses keys were deserialized as ResponsesAPIResponse (500) or re-translated as responses events (empty streaming deltas). Deserialize chat-shaped cache entries as acompletion and bypass the responses stream iterator for cached CustomStreamWrapper replay. Co-authored-by: Cursor * fix(caching): map responses bridge call_type for sync vs async stream replay Co-authored-by: Yassin Kortam * fix: handle ModelResponse cache return in responses bridge and drop dead acompletion check Co-authored-by: Yassin Kortam * fix(caching): detect chat cache hits via object field before choices fallback Prefer chat.completion object type over the broad choices-key heuristic so Responses API cached payloads are not misclassified if their schema changes. Co-authored-by: Cursor * test(caching): cover responses bridge cache-hit paths in CI-tracked test suite The new bridge cache replay logic in caching_handler.py and the preformatted-stream guard in litellm_responses_transformation/handler.py were exercised only by tests under tests/local_testing/, which the responses-caching-types and misc shards do not run. Codecov flagged the patch as 29.72% covered. Add equivalent unit tests under tests/test_litellm/ so the responses, caching, types, and misc shards execute them and ship their coverage data to Codecov: - _is_chat_completion_cached_dict happy/sad paths - aresponses streaming bridge cache hit -> CustomStreamWrapper - responses non-streaming bridge cache hit -> ModelResponse - legacy ResponsesAPIResponse stream + non-stream replay - _is_preformatted_cached_chat_stream true/false - completion/acompletion early return on cached ModelResponse - completion/acompletion skip rewrap on preformatted cached stream * fix: add negative guard on object field in _is_chat_completion_cached_dict Co-authored-by: Yassin Kortam * fix(vcr): treat corrupt cassette payloads as cache miss * test: bump EOL'd NVIDIA rerank and OpenAI realtime models in CI The NVIDIA hosted rerank endpoint for nvidia/llama-3_2-nv-rerankqa-1b-v2 reached end-of-life on 2026-05-18 and now returns HTTP 410 Gone, breaking TestNvidiaNim::test_basic_rerank. Switch to nvidia/nv-rerankqa-mistral-4b-v3, which is still hosted on the NVIDIA API catalog and is already listed in model_prices_and_context_window.json. OpenAI also retired the gpt-4o-realtime-preview-2024-12-17 model used by test_realtime_guardrails_openai (now returns model_not_found). Switch the realtime test URL to the GA gpt-realtime alias. Unrelated to the responses-bridge cache fix in this PR, but committing here to unblock CI per maintainer guidance. Co-authored-by: Mateo Wang * test(realtime): switch retired gpt-4o-realtime-preview to gpt-realtime OpenAI removed gpt-4o-realtime-preview and all its date snapshots on 2026-05-18 (every variant now returns model_not_found), breaking the live-WebSocket OpenAI realtime tests in CI: - test_openai_realtime_direct_call_no_intent - test_openai_realtime_direct_call_with_intent - TestOpenAIRealtime.test_realtime_connection - TestOpenAIRealtime.test_realtime_with_query_params Point each of those to the current GA alias gpt-realtime (verified live). Pure unit/mock tests that just assert the string value (e.g. in test_realtime_query_params_construction and the test_realtime_query_params_use_normalized_model_name mock) are left alone since they do not depend on model availability. Also relax the AI-response assertion in test_text_message_blocked_by_guardrail_no_ai_response: gpt-realtime occasionally produces a polite refusal ("I'm sorry, but I can't say that") when the cancel arrives after the model has already started generating, which is the expected outcome (no real AI content) but does not contain the words 'blocked' or 'guardrail'. The primary guardrail behaviour (guardrail_violation error event + transcript_delta block message) is still asserted unchanged. Co-authored-by: Mateo Wang * test(nvidia_nim): mock rerank live API instead of hitting EOL'd endpoint NVIDIA reached end-of-life for the hosted nvidia/llama-3.2-nv-rerankqa-1b-v2 rerank API on 2026-05-18 (returns HTTP 410 Gone), and the proposed replacement nv-rerankqa-mistral-4b-v3 returns HTTP 404 for the CI account, breaking TestNvidiaNim::test_basic_rerank. Override test_basic_rerank to mock the HTTP transport (same pattern as test_nvidia_nim_rerank_ranking_endpoint above) so the request/response transformation and cost calculation stay covered without depending on NVIDIA's hosted catalog rotation. The model identifier reverts to the original llama-3.2-nv-rerankqa-1b-v2 since the request never leaves the test process. --------- Co-authored-by: Cursor Co-authored-by: Yassin Kortam Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Mateo Wang --- litellm/caching/caching_handler.py | 70 ++++-- .../handler.py | 25 +++ .../transformation.py | 8 + tests/_vcr_redis_persister.py | 17 +- tests/local_testing/test_caching_handler.py | 66 ++++++ .../caching/test_caching_handler.py | 204 ++++++++++++++++++ ...itellm_responses_transformation_handler.py | 150 +++++++++++++ ...responses_transformation_transformation.py | 27 +++ 8 files changed, 544 insertions(+), 23 deletions(-) create mode 100644 tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 3cf1d911d7f..3f4e54382c9 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -87,6 +87,16 @@ class CachingHandlerResponse(BaseModel): in_memory_cache_obj = InMemoryCache() +def _is_chat_completion_cached_dict(cached_result: dict) -> bool: + cached_id = cached_result.get("id") + if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"): + return True + obj = cached_result.get("object") + if isinstance(obj, str): + return obj.startswith("chat.completion") + return "choices" in cached_result + + def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool: """ When stream=True, do not run success callbacks at cache-hit time. @@ -861,27 +871,47 @@ class LLMCachingHandler: elif (call_type == "aresponses" or call_type == "responses") and isinstance( cached_result, dict ): - from litellm.responses.streaming_iterator import ( - CachedResponsesAPIStreamingIterator, - ) - - response_obj = ResponsesAPIResponse(**cached_result) - if ( - hasattr(response_obj, "_hidden_params") - and response_obj._hidden_params is not None - and isinstance(response_obj._hidden_params, dict) - ): - response_obj._hidden_params["cache_hit"] = True - - if kwargs.get("stream", False) is True: - cached_result = CachedResponsesAPIStreamingIterator( - response=response_obj, - logging_obj=logging_obj, - request_data=kwargs, - call_type=call_type, - ) + use_chat_completion_cache = _is_chat_completion_cached_dict(cached_result) + if use_chat_completion_cache: + if kwargs.get("stream", False) is True: + bridge_call_type = ( + CallTypes.acompletion.value + if call_type == "aresponses" + else CallTypes.completion.value + ) + cached_result = self._convert_cached_stream_response( + cached_result=cached_result, + call_type=bridge_call_type, + logging_obj=logging_obj, + model=model, + ) + else: + cached_result = convert_to_model_response_object( + response_object=cached_result, + model_response_object=ModelResponse(), + ) else: - cached_result = response_obj + from litellm.responses.streaming_iterator import ( + CachedResponsesAPIStreamingIterator, + ) + + response_obj = ResponsesAPIResponse(**cached_result) + if ( + hasattr(response_obj, "_hidden_params") + and response_obj._hidden_params is not None + and isinstance(response_obj._hidden_params, dict) + ): + response_obj._hidden_params["cache_hit"] = True + + if kwargs.get("stream", False) is True: + cached_result = CachedResponsesAPIStreamingIterator( + response=response_obj, + logging_obj=logging_obj, + request_data=kwargs, + call_type=call_type, + ) + else: + cached_result = response_obj if ( hasattr(cached_result, "_hidden_params") diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index ce398ee8288..2de7bda6467 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -37,6 +37,15 @@ class ResponsesToCompletionBridgeHandler: stream = litellm_params.get("stream", False) return bool(stream) + @staticmethod + def _is_preformatted_cached_chat_stream(result: Any) -> bool: + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + return ( + isinstance(result, CustomStreamWrapper) + and result.custom_llm_provider == "cached_response" + ) + @staticmethod def _coerce_response_object( response_obj: Any, @@ -177,6 +186,8 @@ class ResponsesToCompletionBridgeHandler: **request_data, ) + from litellm.types.utils import ModelResponse + stream = self._resolve_stream_flag(optional_params, litellm_params) if isinstance(result, ResponsesAPIResponse): return self.transformation_handler.transform_response( @@ -192,6 +203,8 @@ class ResponsesToCompletionBridgeHandler: api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), ) + elif isinstance(result, ModelResponse): + return result elif not stream: responses_api_response = self._collect_response_from_stream(result) return self.transformation_handler.transform_response( @@ -208,6 +221,10 @@ class ResponsesToCompletionBridgeHandler: json_mode=kwargs.get("json_mode"), ) else: + if self._is_preformatted_cached_chat_stream(result): + return self._apply_post_stream_processing( + result, model, custom_llm_provider + ) completion_stream = self.transformation_handler.get_model_response_iterator( streaming_response=result, # type: ignore sync_stream=True, @@ -256,6 +273,8 @@ class ResponsesToCompletionBridgeHandler: aresponses=True, ) + from litellm.types.utils import ModelResponse + stream = self._resolve_stream_flag(optional_params, litellm_params) if isinstance(result, ResponsesAPIResponse): return self.transformation_handler.transform_response( @@ -271,6 +290,8 @@ class ResponsesToCompletionBridgeHandler: api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), ) + elif isinstance(result, ModelResponse): + return result elif not stream: responses_api_response = await self._collect_response_from_stream_async( result @@ -289,6 +310,10 @@ class ResponsesToCompletionBridgeHandler: json_mode=kwargs.get("json_mode"), ) else: + if self._is_preformatted_cached_chat_stream(result): + return self._apply_post_stream_processing( + result, model, custom_llm_provider + ) completion_stream = self.transformation_handler.get_model_response_iterator( streaming_response=result, # type: ignore sync_stream=False, diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 32423f23314..e3cbf422e5d 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1141,6 +1141,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): event_type = parsed_chunk.get("type") if isinstance(event_type, ResponsesAPIStreamEvents): event_type = event_type.value + + if parsed_chunk.get("object") == "chat.completion.chunk" or ( + event_type is None + and isinstance(parsed_chunk.get("choices"), list) + and parsed_chunk.get("choices") + ): + return ModelResponseStream(**parsed_chunk) + verbose_logger.debug(f"Chat provider: Processing event type: {event_type}") if event_type == "response.created": diff --git a/tests/_vcr_redis_persister.py b/tests/_vcr_redis_persister.py index 7fdb7267a38..373cb66696a 100644 --- a/tests/_vcr_redis_persister.py +++ b/tests/_vcr_redis_persister.py @@ -159,9 +159,20 @@ def make_redis_persister( raise CassetteNotFoundError() from exc if data is None: raise CassetteNotFoundError() - if isinstance(data, bytes): - data = data.decode("utf-8") - return deserialize(data, serializer) + try: + if isinstance(data, bytes): + data = data.decode("utf-8") + return deserialize(data, serializer) + except Exception as exc: + _record_cache_failure("load", exc) + msg = ( + f"VCR redis load failed for {cassette_path}; cached " + f"payload is corrupt, treating as cache miss: " + f"{type(exc).__name__}: {exc}" + ) + _log.warning(msg) + warnings.warn(msg, VCRCassetteCacheWarning, stacklevel=2) + raise CassetteNotFoundError() from exc @staticmethod def save_cassette(cassette_path, cassette_dict, serializer): diff --git a/tests/local_testing/test_caching_handler.py b/tests/local_testing/test_caching_handler.py index 2b6712cbaa3..0f4539162a2 100644 --- a/tests/local_testing/test_caching_handler.py +++ b/tests/local_testing/test_caching_handler.py @@ -25,6 +25,7 @@ from unittest.mock import AsyncMock, patch, MagicMock from litellm.caching.caching_handler import ( LLMCachingHandler, CachingHandlerResponse, + _is_chat_completion_cached_dict, _should_defer_streaming_cache_hit_callbacks, ) from litellm.caching.caching import LiteLLMCacheType @@ -40,6 +41,7 @@ from litellm.types.utils import ( from litellm.types.llms.openai import ResponsesAPIResponse from datetime import timedelta, datetime from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm._logging import verbose_logger import logging @@ -1072,6 +1074,70 @@ def test_convert_cached_streaming_responses_result_to_iterator(): ) +def test_is_chat_completion_cached_dict(): + assert _is_chat_completion_cached_dict( + {"id": "chatcmpl-abc", "object": "chat.completion", "choices": []} + ) + assert _is_chat_completion_cached_dict( + {"id": "other", "object": "chat.completion.chunk", "choices": []} + ) + assert not _is_chat_completion_cached_dict( + {"id": "resp_abc", "object": "response", "output": []} + ) + + +def test_convert_cached_aresponses_bridge_chat_completion_stream(): + """ + openai/responses chat-completions bridge caches ModelResponse JSON on aresponses + cache keys; replay must not call ResponsesAPIResponse(**chatcmpl_dict). + """ + caching_handler = LLMCachingHandler( + original_function=aresponses, request_kwargs={}, start_time=datetime.now() + ) + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.aresponses.value, + model="gpt-5.4", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + cached_result = { + "id": "chatcmpl-bridge-cache-test", + "object": "chat.completion", + "created": int(time.time()), + "model": "gpt-5.4", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hi!"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 7, + "completion_tokens": 11, + "total_tokens": 18, + }, + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.aresponses.value, + kwargs={ + "model": "gpt-5.4", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + logging_obj=logging_obj, + model="gpt-5.4", + args=(), + ) + + assert isinstance(result, CustomStreamWrapper) + + def test_convert_cached_streaming_reasoning_result_to_iterator(): caching_handler = LLMCachingHandler( original_function=responses, request_kwargs={}, start_time=datetime.now() diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 742a4f410d4..3eb949d7f29 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -232,3 +232,207 @@ def test_combine_usage_handles_none_details(): combined = llm_caching_handler.combine_usage(usage_a, usage_c) assert combined.prompt_tokens_details is not None assert combined.prompt_tokens_details.image_count == 1 + + +def test_is_chat_completion_cached_dict(): + from litellm.caching.caching_handler import _is_chat_completion_cached_dict + + assert _is_chat_completion_cached_dict( + {"id": "chatcmpl-abc", "object": "chat.completion", "choices": []} + ) + assert _is_chat_completion_cached_dict( + {"id": "other", "object": "chat.completion.chunk", "choices": []} + ) + assert _is_chat_completion_cached_dict( + {"id": "no-object", "choices": [{"index": 0}]} + ) + assert not _is_chat_completion_cached_dict( + {"id": "resp_abc", "object": "response", "output": []} + ) + + +def _build_logging_obj(call_type: str, stream: bool): + import uuid as _uuid + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + return LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=call_type, + model="gpt-5.4", + messages=[], + function_id=str(_uuid.uuid4()), + stream=stream, + start_time=datetime.now(), + ) + + +def test_convert_cached_aresponses_bridge_chat_completion_stream(): + """openai/responses chat-completions bridge: streaming cache hit replays as chat stream.""" + from litellm import aresponses + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.types.utils import CallTypes + + caching_handler = LLMCachingHandler( + original_function=aresponses, request_kwargs={}, start_time=datetime.now() + ) + cached_result = { + "id": "chatcmpl-bridge-cache-test", + "object": "chat.completion", + "created": int(time.time()), + "model": "gpt-5.4", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hi!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}, + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.aresponses.value, + kwargs={ + "model": "gpt-5.4", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + logging_obj=_build_logging_obj(CallTypes.aresponses.value, stream=True), + model="gpt-5.4", + args=(), + ) + + assert isinstance(result, CustomStreamWrapper) + + +def test_convert_cached_responses_bridge_chat_completion_nonstream(): + """openai/responses chat-completions bridge: non-streaming cache hit replays as ModelResponse.""" + from litellm import responses + from litellm.types.utils import CallTypes, ModelResponse + + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + cached_result = { + "id": "chatcmpl-bridge-nonstream", + "object": "chat.completion", + "created": int(time.time()), + "model": "gpt-5.4", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hi!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}, + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={ + "model": "gpt-5.4", + "stream": False, + "messages": [{"role": "user", "content": "hi"}], + }, + logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False), + model="gpt-5.4", + args=(), + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hi!" + + +def test_convert_cached_responses_legacy_nonstream_path(): + """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) falls through legacy path.""" + from litellm import responses + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import CallTypes + + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + cached_result = { + "id": "resp_legacy_nonstream", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-4o", + "object": "response", + "output": [ + { + "type": "message", + "id": "msg_legacy", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "legacy response", + "annotations": [], + } + ], + } + ], + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={"model": "gpt-4o", "input": "hi", "stream": False}, + logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False), + model="gpt-4o", + args=(), + ) + + assert isinstance(result, ResponsesAPIResponse) + assert result.id == "resp_legacy_nonstream" + + +def test_convert_cached_responses_legacy_stream_path(): + """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) on stream falls through legacy path.""" + from litellm import responses + from litellm.responses.streaming_iterator import ( + CachedResponsesAPIStreamingIterator, + ) + from litellm.types.utils import CallTypes + + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + cached_result = { + "id": "resp_legacy_stream", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-4o", + "object": "response", + "output": [ + { + "type": "message", + "id": "msg_legacy_stream", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "legacy stream", + "annotations": [], + } + ], + } + ], + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={"model": "gpt-4o", "input": "hi", "stream": True}, + logging_obj=_build_logging_obj(CallTypes.responses.value, stream=True), + model="gpt-4o", + args=(), + ) + + assert isinstance(result, CachedResponsesAPIStreamingIterator) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py new file mode 100644 index 00000000000..734033ed6be --- /dev/null +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py @@ -0,0 +1,150 @@ +import os +import sys +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.completion_extras.litellm_responses_transformation.handler import ( + ResponsesToCompletionBridgeHandler, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.utils import ModelResponse + + +def test_is_preformatted_cached_chat_stream_true(): + stream = MagicMock(spec=CustomStreamWrapper) + stream.custom_llm_provider = "cached_response" + assert ( + ResponsesToCompletionBridgeHandler._is_preformatted_cached_chat_stream(stream) + is True + ) + + +def test_is_preformatted_cached_chat_stream_false_wrong_provider(): + stream = MagicMock(spec=CustomStreamWrapper) + stream.custom_llm_provider = "openai" + assert ( + ResponsesToCompletionBridgeHandler._is_preformatted_cached_chat_stream(stream) + is False + ) + + +def test_is_preformatted_cached_chat_stream_false_wrong_type(): + assert ( + ResponsesToCompletionBridgeHandler._is_preformatted_cached_chat_stream( + {"object": "chat.completion.chunk"} + ) + is False + ) + + +def _bridge_kwargs(stream: bool): + logging_obj = LiteLLMLogging( + litellm_call_id="test-call", + call_type="completion", + model="gpt-5.4", + messages=[{"role": "user", "content": "hi"}], + function_id="fn-id", + stream=stream, + start_time=datetime.now(), + ) + return { + "model": "gpt-5.4", + "custom_llm_provider": "openai", + "messages": [{"role": "user", "content": "hi"}], + "optional_params": {"stream": stream}, + "litellm_params": {}, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": logging_obj, + } + + +def test_completion_returns_cached_model_response_directly(): + """Non-streaming bridge cache hit: responses() returns a ModelResponse -> bridge returns it as-is.""" + cached = ModelResponse(id="chatcmpl-cached-nonstream", model="gpt-5.4") + bridge = ResponsesToCompletionBridgeHandler() + + with ( + patch.object( + bridge.transformation_handler, + "transform_request", + return_value={"model": "gpt-5.4", "input": "hi"}, + ), + patch("litellm.responses", return_value=cached), + ): + result = bridge.completion(**_bridge_kwargs(stream=False)) + + assert result is cached + + +@pytest.mark.asyncio +async def test_acompletion_returns_cached_model_response_directly(): + cached = ModelResponse(id="chatcmpl-cached-nonstream-async", model="gpt-5.4") + bridge = ResponsesToCompletionBridgeHandler() + + with ( + patch.object( + bridge.transformation_handler, + "transform_request", + return_value={"model": "gpt-5.4", "input": "hi"}, + ), + patch("litellm.aresponses", new=AsyncMock(return_value=cached)), + ): + result = await bridge.acompletion(**_bridge_kwargs(stream=False)) + + assert result is cached + + +def test_completion_skips_rewrapping_preformatted_cached_chat_stream(): + """Streaming bridge cache hit returning CustomStreamWrapper(cached_response) -> bridge skips re-wrapping.""" + stream = MagicMock(spec=CustomStreamWrapper) + stream.custom_llm_provider = "cached_response" + bridge = ResponsesToCompletionBridgeHandler() + + with ( + patch.object( + bridge.transformation_handler, + "transform_request", + return_value={"model": "gpt-5.4", "input": "hi"}, + ), + patch("litellm.responses", return_value=stream), + patch.object( + bridge, + "_apply_post_stream_processing", + side_effect=lambda s, *a, **kw: s, + ) as post, + ): + result = bridge.completion(**_bridge_kwargs(stream=True)) + + post.assert_called_once() + assert result is stream + + +@pytest.mark.asyncio +async def test_acompletion_skips_rewrapping_preformatted_cached_chat_stream(): + stream = MagicMock(spec=CustomStreamWrapper) + stream.custom_llm_provider = "cached_response" + bridge = ResponsesToCompletionBridgeHandler() + + with ( + patch.object( + bridge.transformation_handler, + "transform_request", + return_value={"model": "gpt-5.4", "input": "hi"}, + ), + patch("litellm.aresponses", new=AsyncMock(return_value=stream)), + patch.object( + bridge, + "_apply_post_stream_processing", + side_effect=lambda s, *a, **kw: s, + ) as post, + ): + result = await bridge.acompletion(**_bridge_kwargs(stream=True)) + + post.assert_called_once() + assert result is stream diff --git a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py index 009f432fca1..05bdc40112c 100644 --- a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py @@ -230,3 +230,30 @@ def test_transform_request_drops_user_metadata_with_additional_drop_params(): assert "metadata" not in result assert result["litellm_metadata"]["internal_key"] == "secret" + + +def test_translate_responses_chunk_passthrough_chat_completion_chunk(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + chat_chunk = { + "id": "chatcmpl-cache-passthrough", + "object": "chat.completion.chunk", + "created": 1779104834, + "model": "gpt-5.4", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Hi! How can I help?"}, + "finish_reason": None, + } + ], + } + + result = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( + chat_chunk + ) + + assert result.choices[0].delta.content == "Hi! How can I help?" + assert result.choices[0].finish_reason is None From 36c494fdd29c6c9a93d057451154fcb17cce05b5 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 19 May 2026 04:57:44 +0530 Subject: [PATCH 07/11] Litellm oss staging (#28161) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(opentelemetry): JSON-serialize dict metadata fields for OTEL span attributes (#27451) (#27455) Squash-merged by litellm-agent from Anai-Guo's PR. * feat(dashscope): add embeddings and reranks(qwen3-rerank) support via OpenAI-compatible endpoint (#27508) Squash-merged by litellm-agent from yimao's PR. * fix(vertex_ai/gemini): raise BadRequestError when image_url or url fi… (#24550) Squash-merged by litellm-agent from krisxia0506's PR. * fix(vertex_ai): raise error on mid-stream 429/error chunks instead of silently swallowing (#23711) Squash-merged by litellm-agent from krisxia0506's PR. * fix: raise BadRequestError for file content blocks missing 'file' sub… (#24503) Squash-merged by litellm-agent from krisxia0506's PR. * Fix Gemini MIME detection for extensionless GCS URIs (#27278) Squash-merged by litellm-agent from krisxia0506's PR. * fix(vertex_ai/partner_models): drop unused vertexai SDK gate from count_tokens (closes #28084) (#28107) Squash-merged by litellm-agent from voidborne-d's PR. * feat(chart): add support for autoscaling behavior in HPA (#27990) Squash-merged by litellm-agent from FabrizioCafolla's PR. * feat(proxy): add blocked flag to models for pause/resume from the UI (#27927) Squash-merged by litellm-agent from Cyberfilo's PR. * fix: pass socket timeouts to Redis cluster clients (#27920) Squash-merged by litellm-agent from tomdee's PR. * Fix/cache token (#28009) Squash-merged by litellm-agent from escon1004's PR. * fix(deepseek): forward reasoning_content in multi-turn thinking mode conversations (#28080) Squash-merged by litellm-agent from Divyansh8321's PR. * fix(guardrails): return HTTP 400 instead of 500 for blocked requests (#27617) * fix: reset org and tag budgets (#27326) * reset org budgets * reset tag budgets --------- Co-authored-by: Michael Riad Zaky * fix(ui): omit allowed_routes from key edit save when unchanged (#27553) * fix(ui): omit allowed_routes from key edit save when unchanged When a team admin opens Edit Settings on a key with key_type=AI APIs and saves without changing anything, the UI re-sends the existing allowed_routes value, which the backend's _check_allowed_routes_caller_permission gate rejects for non-proxy-admins (LIT-2681). Strip allowed_routes from the patch in handleSubmit when it deep-equals the original keyData.allowed_routes. The backend treats absence as "leave alone," so no-op saves now succeed for non-admins. Admins explicitly editing the field still send the new value. * fix(ui): order-insensitive allowed_routes diff + cover null-original case Address Greptile review: - Switch the "is allowed_routes unchanged" check to a Set-based comparison so a server-side reorder of the array doesn't register as a user edit and re-trigger LIT-2681. - Add two regression tests: (1) keyData.allowed_routes is null and the form is untouched — patch should strip the field; (2) server returned routes in a different order than the user originally entered — patch should still recognize the value as unchanged. * chore(ui): strip ticket refs and tighten comments in key edit fix - Remove internal-tracker references from in-code comments - Tighten the WHY comment in handleSubmit to two lines - Drop redundant test-block comments — test names already describe the case * fix(ui): annotate Set generic in allowed_routes diff to fix tsc * fix(guardrails): return HTTP 400 instead of 500 for guardrail-blocked requests GuardrailRaisedException and BlockedPiiEntityError both lacked a status_code attribute. When these exceptions reached the proxy exception handler (getattr(e, 'status_code', 500)), the fallback defaulted to HTTP 500 — making intentional guardrail blocks indistinguishable from server errors and causing unnecessary client retries. Changes: - Add status_code=400 (keyword-only) to GuardrailRaisedException - Add status_code=400 (keyword-only) to BlockedPiiEntityError - Update _is_guardrail_intervention() to recognize both exceptions so downstream loggers record 'guardrail_intervened' instead of 'guardrail_failed_to_respond' - Add 6 unit tests for default/custom status codes and getattr pattern - Strengthen existing blocked-action test with status_code assertion Fixes #24348 --------- Co-authored-by: Michael-RZ-Berri Co-authored-by: Michael Riad Zaky Co-authored-by: ryan-crabbe-berri Co-authored-by: Krrish Dholakia * fix(router/proxy): address Greptile P1+P2 review comments on PR #28161 - router: raise ServiceUnavailableError (503) instead of RouterRateLimitErrorBasic (429) when a specifically-addressed deployment is administratively blocked; 429 misleads retry-enabled clients into spinning forever against a paused model - proxy_server: compute get_fully_blocked_model_names() once before both branches in model_list() instead of duplicating the call in each branch - deepseek: upgrade silent debug log to warning when injecting placeholder reasoning_content so callers are clearly notified of degraded multi-turn quality - tests: update two blocked-deployment assertions to expect ServiceUnavailableError Co-authored-by: Cursor * fix: address bug detection findings (cache token order, mutable defaults) Co-authored-by: Yassin Kortam * fix: address bugs in async pass-through, anthropic cache token detection, rerank tests - async_get_available_deployment_for_pass_through: enforce blocked check on specific deployments - cost_calculator: detect anthropic-style usage by attribute presence (not truthiness) to avoid mixing OpenAI cached_tokens into anthropic normalization when read=0 - dashscope rerank tests: pass request to httpx.Response constructions for consistency Co-authored-by: Yassin Kortam * fix code qa * fix(vertex_ai/gemini): strip MIME parameters from GCS contentType GCS object metadata's contentType field can include parameters such as 'text/html; charset=utf-8'. Strip them in _apply_gemini_mime_type_aliases so downstream get_file_extension_from_mime_type sees a bare MIME type. Co-authored-by: Yassin Kortam * fix(vertex_ai/gemini): clarify mime-type error message string concatenation Co-authored-by: Yassin Kortam --------- Co-authored-by: Tai An Co-authored-by: Vincent Co-authored-by: Kris Xia Co-authored-by: d 🔹 Co-authored-by: Fabrizio Cafolla Co-authored-by: Filippo Menghi <113345637+Cyberfilo@users.noreply.github.com> Co-authored-by: Tom Denham Co-authored-by: escon1004 <70471150+escon1004@users.noreply.github.com> Co-authored-by: Divyansh Singhal <97736786+Divyansh8321@users.noreply.github.com> Co-authored-by: robin-fiddler Co-authored-by: Michael-RZ-Berri Co-authored-by: Michael Riad Zaky Co-authored-by: ryan-crabbe-berri Co-authored-by: Krrish Dholakia Co-authored-by: Cursor Co-authored-by: Yassin Kortam --- README.md | 2 +- deploy/charts/litellm-helm/templates/hpa.yaml | 4 + .../charts/litellm-helm/tests/hpa_tests.yaml | 36 ++ deploy/charts/litellm-helm/values.yaml | 1 + .../migration.sql | 4 + .../litellm_proxy_extras/schema.prisma | 5 +- litellm/__init__.py | 6 + litellm/_redis.py | 2 + litellm/cost_calculator.py | 82 ++- litellm/exceptions.py | 4 + litellm/integrations/custom_guardrail.py | 11 +- litellm/integrations/opentelemetry.py | 9 +- .../prompt_templates/common_utils.py | 14 +- .../prompt_templates/factory.py | 29 +- litellm/llms/dashscope/common_utils.py | 28 + litellm/llms/dashscope/embed/__init__.py | 7 + .../llms/dashscope/embed/transformation.py | 191 +++++++ litellm/llms/dashscope/rerank/__init__.py | 7 + .../llms/dashscope/rerank/transformation.py | 241 +++++++++ litellm/llms/deepseek/chat/transformation.py | 106 +++- litellm/llms/gemini/chat/transformation.py | 24 +- .../llms/openai/chat/gpt_transformation.py | 8 +- .../llms/vertex_ai/gemini/transformation.py | 503 +++++++++++++++++- .../vertex_and_google_ai_studio_gemini.py | 39 +- .../vertex_ai_partner_models/main.py | 25 +- litellm/main.py | 27 + litellm/proxy/_types.py | 1 + litellm/proxy/db/db_spend_update_writer.py | 37 +- .../model_management_endpoints.py | 17 + litellm/proxy/proxy_server.py | 14 + litellm/proxy/route_llm_request.py | 6 +- litellm/proxy/schema.prisma | 5 +- litellm/router.py | 120 ++++- litellm/types/router.py | 4 + litellm/types/utils.py | 10 +- litellm/utils.py | 12 + schema.prisma | 5 +- .../test_deepseek_completion.py | 110 ++++ ...test_dashscope_embedding_transformation.py | 141 +++++ .../test_dashscope_rerank_transformation.py | 328 ++++++++++++ .../llms/test_file_content_block.py | 433 +++++++++++++++ .../test_gemini_image_url_missing_field.py | 52 ++ ...test_vertex_and_google_ai_studio_gemini.py | 441 +++++++++++++++ .../llms/vertex_ai/test_vertex.py | 27 + .../test_vertex_gemini_gcs_uri_mime.py | 466 ++++++++++++++++ .../test_count_tokens_no_vertexai_sdk.py | 121 +++++ .../test_generic_guardrail_api.py | 1 + .../test_model_management_endpoints.py | 172 ++++++ tests/test_litellm/test_cost_calculator.py | 314 +++++++++++ .../test_guardrail_exception_status_codes.py | 66 +++ tests/test_litellm/test_redis.py | 8 +- tests/test_litellm/test_router.py | 170 ++++++ .../UsagePage/components/UsagePageView.tsx | 7 +- .../add_model/advanced_settings.tsx | 20 + .../add_model/handle_add_model_submit.tsx | 41 +- .../src/components/model_info_view.tsx | 90 ++++ 56 files changed, 4547 insertions(+), 107 deletions(-) create mode 100644 deploy/charts/litellm-helm/tests/hpa_tests.yaml create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260514120000_add_blocked_to_proxy_model_table/migration.sql create mode 100644 litellm/llms/dashscope/common_utils.py create mode 100644 litellm/llms/dashscope/embed/__init__.py create mode 100644 litellm/llms/dashscope/embed/transformation.py create mode 100644 litellm/llms/dashscope/rerank/__init__.py create mode 100644 litellm/llms/dashscope/rerank/transformation.py create mode 100644 tests/test_litellm/llms/dashscope/test_dashscope_embedding_transformation.py create mode 100644 tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py create mode 100644 tests/test_litellm/llms/test_file_content_block.py create mode 100644 tests/test_litellm/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py create mode 100644 tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py create mode 100644 tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py create mode 100644 tests/test_litellm/test_guardrail_exception_status_codes.py diff --git a/README.md b/README.md index 72fd43925c9..8df351e9303 100644 --- a/README.md +++ b/README.md @@ -292,7 +292,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ | [CompactifAI (`compactifai`)](https://docs.litellm.ai/docs/providers/compactifai) | ✅ | ✅ | ✅ | | | | | | | | | [Custom (`custom`)](https://docs.litellm.ai/docs/providers/custom_llm_server) | ✅ | ✅ | ✅ | | | | | | | | | [Custom OpenAI (`custom_openai`)](https://docs.litellm.ai/docs/providers/openai_compatible) | ✅ | ✅ | ✅ | | | ✅ | ✅ | ✅ | ✅ | | -| [Dashscope (`dashscope`)](https://docs.litellm.ai/docs/providers/dashscope) | ✅ | ✅ | ✅ | | | | | | | | +| [Dashscope (`dashscope`)](https://docs.litellm.ai/docs/providers/dashscope) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ | | [Databricks (`databricks`)](https://docs.litellm.ai/docs/providers/databricks) | ✅ | ✅ | ✅ | | | | | | | | | [DataRobot (`datarobot`)](https://docs.litellm.ai/docs/providers/datarobot) | ✅ | ✅ | ✅ | | | | | | | | | [Deepgram (`deepgram`)](https://docs.litellm.ai/docs/providers/deepgram) | ✅ | ✅ | ✅ | | | ✅ | | | | | diff --git a/deploy/charts/litellm-helm/templates/hpa.yaml b/deploy/charts/litellm-helm/templates/hpa.yaml index 71e199c5aeb..fec4d1f5c5e 100644 --- a/deploy/charts/litellm-helm/templates/hpa.yaml +++ b/deploy/charts/litellm-helm/templates/hpa.yaml @@ -12,6 +12,10 @@ spec: name: {{ include "litellm.fullname" . }} minReplicas: {{ .Values.autoscaling.minReplicas }} maxReplicas: {{ .Values.autoscaling.maxReplicas }} + {{- if .Values.autoscaling.behavior }} + behavior: + {{- toYaml .Values.autoscaling.behavior | nindent 4 }} + {{- end }} metrics: {{- if .Values.autoscaling.targetCPUUtilizationPercentage }} - type: Resource diff --git a/deploy/charts/litellm-helm/tests/hpa_tests.yaml b/deploy/charts/litellm-helm/tests/hpa_tests.yaml new file mode 100644 index 00000000000..ec18c3591d3 --- /dev/null +++ b/deploy/charts/litellm-helm/tests/hpa_tests.yaml @@ -0,0 +1,36 @@ +suite: "hpa with behavior" +templates: + - hpa.yaml +tests: + - it: "renders behavior when set" + set: + autoscaling.enabled: true + autoscaling.behavior: + scaleUp: + stabilizationWindowSeconds: 60 + policies: + - type: Pods + value: 2 + periodSeconds: 60 + scaleDown: + stabilizationWindowSeconds: 90 + policies: + - type: Pods + value: 1 + periodSeconds: 60 + asserts: + - isKind: { of: HorizontalPodAutoscaler } + - equal: { path: spec.behavior.scaleUp.stabilizationWindowSeconds, value: 60 } + - equal: { path: spec.behavior.scaleDown.stabilizationWindowSeconds, value: 90 } + +--- +suite: "hpa without behavior" +templates: + - hpa.yaml +tests: + - it: "does not render behavior when not set" + set: + autoscaling.enabled: true + asserts: + - isKind: { of: HorizontalPodAutoscaler } + - isNull: { path: spec.behavior } diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index 9c7c013341b..81558ed5b29 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -184,6 +184,7 @@ autoscaling: maxReplicas: 100 targetCPUUtilizationPercentage: 80 # targetMemoryUtilizationPercentage: 80 + # behavior: {} # Autoscaling with keda is mutually exclusive with hpa keda: diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260514120000_add_blocked_to_proxy_model_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260514120000_add_blocked_to_proxy_model_table/migration.sql new file mode 100644 index 00000000000..3253b63a884 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260514120000_add_blocked_to_proxy_model_table/migration.sql @@ -0,0 +1,4 @@ +-- AlterTable +-- Adds the admin-toggleable pause flag used by the router's blocked filter and the +-- credential lookup helpers; defaults to false so existing rows behave unchanged. +ALTER TABLE "LiteLLM_ProxyModelTable" ADD COLUMN IF NOT EXISTS "blocked" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b53507abe6a..78143fe0411 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable { // Models on proxy model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) - model_name String + model_name String litellm_params Json - model_info Json? + model_info Json? + blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/__init__.py b/litellm/__init__.py index e1b367fb234..b9da0524095 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1880,6 +1880,12 @@ if TYPE_CHECKING: from .llms.dashscope.chat.transformation import ( DashScopeChatConfig as DashScopeChatConfig, ) + from .llms.dashscope.embed.transformation import ( + DashScopeEmbeddingConfig as DashScopeEmbeddingConfig, + ) + from .llms.dashscope.rerank.transformation import ( + DashScopeRerankConfig as DashScopeRerankConfig, + ) from .llms.moonshot.chat.transformation import ( MoonshotChatConfig as MoonshotChatConfig, ) diff --git a/litellm/_redis.py b/litellm/_redis.py index 65284162663..5ab551453bb 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -100,6 +100,8 @@ def _get_redis_cluster_kwargs(client=None): "azure_tenant_id", "azure_client_secret", "max_connections", + "socket_timeout", + "socket_connect_timeout", } return available_args diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index c0ec148d03a..2257861aff6 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -173,17 +173,45 @@ def _cost_per_token_custom_pricing_helper( prompt_tokens: float = 0, completion_tokens: float = 0, response_time_ms: Optional[float] = 0.0, + cached_tokens: float = 0, + cache_creation_tokens: float = 0, ### CUSTOM PRICING ### custom_cost_per_token: Optional[CostPerToken] = None, custom_cost_per_second: Optional[float] = None, ) -> Optional[Tuple[float, float]]: - """Internal helper function for calculating cost, if custom pricing given""" + """Internal helper function for calculating cost, if custom pricing given. + + prompt_tokens is assumed to include both cached_tokens and cache_creation_tokens + (OpenAI-compatible convention). Anthropic-style usage where prompt_tokens excludes + cache tokens is handled at the caller (cost_per_token) before invoking this helper. + """ if custom_cost_per_token is None and custom_cost_per_second is None: return None if custom_cost_per_token is not None: - input_cost = custom_cost_per_token["input_cost_per_token"] * prompt_tokens - output_cost = custom_cost_per_token["output_cost_per_token"] * completion_tokens + input_cost_per_token = custom_cost_per_token["input_cost_per_token"] + output_cost_per_token = custom_cost_per_token["output_cost_per_token"] + + cache_read_input_token_cost = custom_cost_per_token.get( + "cache_read_input_token_cost", + input_cost_per_token, + ) + cache_creation_input_token_cost = custom_cost_per_token.get( + "cache_creation_input_token_cost", + input_cost_per_token, + ) + + regular_prompt_tokens = max( + prompt_tokens - cached_tokens - cache_creation_tokens, + 0, + ) + + input_cost = ( + regular_prompt_tokens * input_cost_per_token + + cached_tokens * cache_read_input_token_cost + + cache_creation_tokens * cache_creation_input_token_cost + ) + output_cost = completion_tokens * output_cost_per_token return input_cost, output_cost elif custom_cost_per_second is not None: output_cost = custom_cost_per_second * response_time_ms / 1000 # type: ignore @@ -323,10 +351,56 @@ def cost_per_token( # noqa: PLR0915 ) ## CUSTOM PRICING ## + # Normalize cache token counts across providers: + # - OpenAI-compatible: usage.prompt_tokens_details.cached_tokens + # (prompt_tokens already INCLUDES cached_tokens) + # - Anthropic: usage.cache_read_input_tokens / cache_creation_input_tokens + # (prompt_tokens does NOT include these — adjust before calling helper) + _cache_read_tokens: float = 0 + _cache_creation_tokens: float = 0 + _is_anthropic_style = False + + if usage_object is not None: + _pt_details = getattr(usage_object, "prompt_tokens_details", None) + if _pt_details is not None: + _cache_read_tokens = float(getattr(_pt_details, "cached_tokens", 0) or 0) + # OpenAI-compatible providers report cache-write tokens under + # either `cache_write_tokens` (kimi-k2) or `cache_creation_tokens`. + # Mirror db_spend_update_writer to stay symmetric. + _cache_creation_tokens = float( + getattr(_pt_details, "cache_write_tokens", 0) + or getattr(_pt_details, "cache_creation_tokens", 0) + or 0 + ) + + _anthropic_read = getattr(usage_object, "cache_read_input_tokens", None) + _anthropic_create = getattr(usage_object, "cache_creation_input_tokens", None) + if _anthropic_read is not None or _anthropic_create is not None: + _is_anthropic_style = True + if _anthropic_read is not None: + _cache_read_tokens = float(_anthropic_read) + if _anthropic_create is not None: + _cache_creation_tokens = float(_anthropic_create) + + if not _cache_read_tokens and cache_read_input_tokens: + _cache_read_tokens = float(cache_read_input_tokens) + _is_anthropic_style = True + if not _cache_creation_tokens and cache_creation_input_tokens: + _cache_creation_tokens = float(cache_creation_input_tokens) + _is_anthropic_style = True + + # Anthropic reports prompt_tokens as input_tokens (excluding cache tokens). + # Adjust so the helper's "prompt_tokens includes cache tokens" invariant holds. + _normalized_prompt_tokens = float(prompt_tokens) + if _is_anthropic_style: + _normalized_prompt_tokens += _cache_read_tokens + _cache_creation_tokens + response_cost = _cost_per_token_custom_pricing_helper( - prompt_tokens=prompt_tokens, + prompt_tokens=_normalized_prompt_tokens, completion_tokens=completion_tokens, response_time_ms=response_time_ms, + cached_tokens=_cache_read_tokens, + cache_creation_tokens=_cache_creation_tokens, custom_cost_per_second=custom_cost_per_second, custom_cost_per_token=custom_cost_per_token, ) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 8b005291556..17f5b43c273 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -918,9 +918,11 @@ class GuardrailRaisedException(Exception): guardrail_name: Optional[str] = None, message: str = "", should_wrap_with_default_message: bool = True, + status_code: int = 400, ): default_message = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}" self.guardrail_name = guardrail_name + self.status_code = status_code self.message = default_message if should_wrap_with_default_message else message super().__init__(self.message) @@ -930,12 +932,14 @@ class BlockedPiiEntityError(Exception): self, entity_type: str, guardrail_name: Optional[str] = None, + status_code: int = 400, ): """ Raised when a blocked entity is detected by a guardrail. """ self.entity_type = entity_type self.guardrail_name = guardrail_name + self.status_code = status_code self.message = f"Blocked entity detected: {entity_type} by Guardrail: {guardrail_name}. This entity is not allowed to be used in this request." super().__init__(self.message) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index c937ad0a7bf..82a35f2eedd 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -43,7 +43,11 @@ if TYPE_CHECKING: dc = DualCache() -from litellm.exceptions import ModifyResponseException as ModifyResponseException +from litellm.exceptions import ( + BlockedPiiEntityError, + GuardrailRaisedException, + ModifyResponseException, +) class CustomGuardrail(CustomLogger): @@ -737,12 +741,15 @@ class CustomGuardrail(CustomLogger): (this was logged previously as an API failure - guardrail_failed_to_respond). Guardrails signal intentional blocks by raising: + - GuardrailRaisedException (generic guardrail API, tool permission) + - BlockedPiiEntityError (Presidio PII detection) - HTTPException with status 400 (content policy violation) - ModifyResponseException (passthrough mode violation) """ - if isinstance(e, ModifyResponseException): return True + if isinstance(e, (GuardrailRaisedException, BlockedPiiEntityError)): + return True if ( HTTPException is not None and isinstance(e, HTTPException) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 2de8e8bf7f1..ced15a01660 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1107,8 +1107,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): "mcp_tool_call_metadata", "vector_store_request_metadata", ]: - if md.get(key) is not None: - common_attrs[f"metadata.{key}"] = str(md[key]) + value = md.get(key) + if value is None: + continue + if isinstance(value, (dict, list)): + common_attrs[f"metadata.{key}"] = safe_dumps(value) + else: + common_attrs[f"metadata.{key}"] = str(value) # get hidden params hidden_params = getattr(std_log, "hidden_params", None) or (std_log or {}).get( diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index f417b4a5f61..3ee56dfc5ca 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -20,6 +20,7 @@ from typing import ( cast, ) +import litellm from litellm import verbose_logger from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( @@ -1170,9 +1171,16 @@ def migrate_file_to_image_url( ChatCompletionImageUrlObject, ) - file_id = message["file"].get("file_id") - file_data = message["file"].get("file_data") - format = message["file"].get("format") + file_sub = message.get("file") + if file_sub is None: + raise litellm.BadRequestError( + message="Content block has type='file' but is missing the required 'file' field", + model=None, + llm_provider=None, + ) + file_id = file_sub.get("file_id") + file_data = file_sub.get("file_data") + format = file_sub.get("format") if not file_id and not file_data: raise ValueError("file_id and file_data are both None") image_url_object = ChatCompletionImageObject( diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index d40ca4e3597..79d527d1eb8 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2057,9 +2057,16 @@ def anthropic_process_openai_file_message( AnthropicMessagesContainerUploadParam, ]: file_message = cast(ChatCompletionFileObject, message) - file_data = file_message["file"].get("file_data") - file_id = file_message["file"].get("file_id") - format = file_message["file"].get("format") + file_sub = file_message.get("file") + if file_sub is None: + raise litellm.BadRequestError( + message="Content block has type='file' but is missing the required 'file' field", + model=None, + llm_provider="anthropic", + ) + file_data = file_sub.get("file_data") + file_id = file_sub.get("file_id") + format = file_sub.get("format") if file_data: image_chunk = convert_to_anthropic_image_obj( openai_image_url=file_data, @@ -4879,7 +4886,13 @@ class BedrockConverseMessagesProcessor: @staticmethod def _process_file_message(message: ChatCompletionFileObject) -> BedrockContentBlock: - file_message = message["file"] + file_message = message.get("file") + if file_message is None: + raise litellm.BadRequestError( + message="Content block has type='file' but is missing the required 'file' field", + model=None, + llm_provider="bedrock", + ) file_data = file_message.get("file_data") file_id = file_message.get("file_id") @@ -4900,7 +4913,13 @@ class BedrockConverseMessagesProcessor: async def _async_process_file_message( message: ChatCompletionFileObject, ) -> BedrockContentBlock: - file_message = message["file"] + file_message = message.get("file") + if file_message is None: + raise litellm.BadRequestError( + message="Content block has type='file' but is missing the required 'file' field", + model=None, + llm_provider="bedrock", + ) file_data = file_message.get("file_data") file_id = file_message.get("file_id") format = file_message.get("format") diff --git a/litellm/llms/dashscope/common_utils.py b/litellm/llms/dashscope/common_utils.py new file mode 100644 index 00000000000..b3b89cbbebf --- /dev/null +++ b/litellm/llms/dashscope/common_utils.py @@ -0,0 +1,28 @@ +""" +Common utilities for the DashScope LLM provider. +""" + +from typing import Optional + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException + + +class DashScopeError(BaseLLMException): + """Exception class for DashScope provider errors.""" + + def __init__( + self, + status_code: int, + message: str, + headers: Optional[httpx.Headers] = None, + ): + self.status_code = status_code + self.message = message + self.headers = headers or httpx.Headers() + super().__init__( + status_code=status_code, + message=message, + headers=dict(self.headers), + ) diff --git a/litellm/llms/dashscope/embed/__init__.py b/litellm/llms/dashscope/embed/__init__.py new file mode 100644 index 00000000000..4962b1f3251 --- /dev/null +++ b/litellm/llms/dashscope/embed/__init__.py @@ -0,0 +1,7 @@ +""" +DashScope Embedding Module +""" + +from .transformation import DashScopeEmbeddingConfig + +__all__ = ["DashScopeEmbeddingConfig"] diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py new file mode 100644 index 00000000000..5bc0e5ca817 --- /dev/null +++ b/litellm/llms/dashscope/embed/transformation.py @@ -0,0 +1,191 @@ +""" +Transformation logic from OpenAI /v1/embeddings format to DashScope's /v1/embeddings format. + +Supports +- text-embedding-v4 +- text-embedding-v3 + +Endpoint +- https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings + +Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api +""" + +from typing import List, Optional, Union + +import httpx + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues +from litellm.types.utils import EmbeddingResponse, Usage + +from ..common_utils import DashScopeError + +DEFAULT_API_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1" + + +class DashScopeEmbeddingConfig(BaseEmbeddingConfig): + """ + Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api + + DashScope exposes an OpenAI-compatible /v1/embeddings endpoint, so the + request and response shapes are nearly identical to OpenAI's. + """ + + def __init__(self) -> None: + pass + + def get_supported_openai_params(self, model: str) -> List[str]: + # DashScope's compatible-mode embeddings API accepts the same params as OpenAI. + # `dimensions` / `encoding_format` are only honored by text-embedding-v3 / v4; + # earlier versions silently ignore them server-side. + return ["dimensions", "encoding_format", "user"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool = False, + ) -> dict: + supported = self.get_supported_openai_params(model) + for k, v in non_default_params.items(): + if v is None: + continue + if k in supported: + optional_params[k] = v + # unsupported params are dropped when drop_params=True; + # the upstream _check_valid_arg already raised UnsupportedParamsError + # for drop_params=False before this method is called. + return optional_params + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + if api_key is None: + api_key = get_secret_str("DASHSCOPE_API_KEY") + if api_key is None: + raise ValueError( + "DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly." + ) + default_headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + return {**default_headers, **headers} + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + base = api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE + base = base.rstrip("/") + if base.endswith("/embeddings"): + return base + return f"{base}/embeddings" + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict, + headers: dict, + ) -> dict: + data: dict = { + "model": model, + "input": input, + } + for key in ("dimensions", "encoding_format", "user"): + value = optional_params.get(key) + if value is not None: + data[key] = value + return data + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + request_data: dict, + optional_params: dict, + litellm_params: dict, + ) -> EmbeddingResponse: + try: + response_json = raw_response.json() + except Exception as e: + raise DashScopeError( + status_code=raw_response.status_code, + message=f"Failed to parse DashScope response as JSON: {str(e)}", + ) + + logging_obj.post_call( + input=request_data.get("input"), + api_key=api_key, + additional_args={"complete_input_dict": request_data}, + original_response=response_json, + ) + + if "error" in response_json: + error = response_json["error"] + message = ( + error.get("message", str(error)) + if isinstance(error, dict) + else str(error) + ) + raise DashScopeError( + status_code=raw_response.status_code, + message=message, + ) + + model_response.object = "list" + model_response.data = response_json.get("data", []) + model_response.model = response_json.get("model", model) + + usage = response_json.get("usage") or {} + prompt_tokens = usage.get("prompt_tokens", 0) + total_tokens = usage.get("total_tokens", prompt_tokens) + setattr( + model_response, + "usage", + Usage( + prompt_tokens=prompt_tokens, + completion_tokens=0, + total_tokens=total_tokens, + ), + ) + + if "id" in response_json: + setattr(model_response, "id", response_json["id"]) + + return model_response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: Union[dict, httpx.Headers], + ) -> BaseLLMException: + if isinstance(headers, dict): + headers = httpx.Headers(headers) + return DashScopeError( + status_code=status_code, + message=error_message, + headers=headers, + ) diff --git a/litellm/llms/dashscope/rerank/__init__.py b/litellm/llms/dashscope/rerank/__init__.py new file mode 100644 index 00000000000..2a1401f6dc0 --- /dev/null +++ b/litellm/llms/dashscope/rerank/__init__.py @@ -0,0 +1,7 @@ +""" +DashScope Rerank Module +""" + +from .transformation import DashScopeRerankConfig + +__all__ = ["DashScopeRerankConfig"] diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py new file mode 100644 index 00000000000..629f3cf4af7 --- /dev/null +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -0,0 +1,241 @@ +""" +Transformation logic for DashScope's OpenAI-compatible /v1/reranks API. + +Supports +- qwen3-rerank + +(Other DashScope rerankers — gte-rerank-v2 / qwen3-vl-rerank — share the same +endpoint but have not been validated against this transformer. Behavior with +those models is undefined.) + +Endpoint +- https://dashscope.aliyuncs.com/compatible-api/v1/reranks + +Note: chat/embed live under `/compatible-mode/v1/`, but DashScope's rerank +route is exposed under `/compatible-api/v1/reranks` per the docs. Override +with `DASHSCOPE_API_BASE_RERANK` to point at a different host or path. + +Empirically, qwen3-rerank accepts `return_documents=true` and echoes +`results[].document.text` back, even though the public docs list the flag +as supported only for gte-rerank-v2 / qwen3-vl-rerank. + +Docs - https://help.aliyun.com/zh/model-studio/text-rerank-api +""" + +from typing import Any, Dict, List, Optional, Union + +import httpx + +from litellm._uuid import uuid +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.rerank import ( + OptionalRerankParams, + RerankBilledUnits, + RerankResponse, + RerankResponseMeta, + RerankTokens, +) + +from ..common_utils import DashScopeError + +DEFAULT_RERANK_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" + + +class DashScopeRerankConfig(BaseRerankConfig): + """ + Reference: https://help.aliyun.com/zh/model-studio/text-rerank-api + + Targets DashScope's qwen3-rerank model. Request fields: model, query, + documents, top_n, return_documents. Response: results[].index, + results[].relevance_score, optionally results[].document.text (when + return_documents=true), plus a top-level usage.total_tokens counter. + """ + + def __init__(self) -> None: + pass + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: + if api_base is None: + api_base = get_secret_str("DASHSCOPE_API_BASE_RERANK") or DEFAULT_RERANK_URL + + if api_base == DEFAULT_RERANK_URL: + return DEFAULT_RERANK_URL + + cleaned = api_base.rstrip("/") + if cleaned.endswith("/reranks") or cleaned.endswith("/rerank"): + return cleaned + + if cleaned.endswith("/v1"): + return f"{cleaned}/reranks" + + # Unknown base: append /reranks rather than silently ignoring the caller's api_base. + return f"{cleaned}/reranks" + + def validate_environment( + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + optional_params: Optional[dict] = None, + ) -> dict: + if api_key is None: + api_key = get_secret_str("DASHSCOPE_API_KEY") + if api_key is None: + raise ValueError( + "DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly." + ) + + default_headers = { + "Authorization": f"Bearer {api_key}", + "accept": "application/json", + "content-type": "application/json", + } + return {**default_headers, **headers} + + def get_supported_cohere_rerank_params(self, model: str) -> list: + return ["query", "documents", "top_n", "return_documents"] + + def map_cohere_rerank_params( + self, + non_default_params: Optional[dict], + model: str, + drop_params: bool, + query: str, + documents: List[Union[str, Dict[str, Any]]], + custom_llm_provider: Optional[str] = None, + top_n: Optional[int] = None, + rank_fields: Optional[List[str]] = None, + return_documents: Optional[bool] = True, + max_chunks_per_doc: Optional[int] = None, + max_tokens_per_doc: Optional[int] = None, + ) -> Dict: + # qwen3-rerank accepts query/documents/top_n/return_documents. The + # rest (rank_fields, max_*_per_doc) are silently dropped. + params: OptionalRerankParams = OptionalRerankParams( + query=query, + documents=documents, + ) + if top_n is not None: + params["top_n"] = top_n + if return_documents is not None: + params["return_documents"] = return_documents + return dict(params) + + def transform_rerank_request( + self, + model: str, + optional_rerank_params: Dict, + headers: dict, + litellm_params: Optional[dict] = None, + ) -> dict: + if "query" not in optional_rerank_params: + raise ValueError("query is required for DashScope rerank") + if "documents" not in optional_rerank_params: + raise ValueError("documents is required for DashScope rerank") + + request: Dict[str, Any] = { + "model": model, + "query": optional_rerank_params["query"], + "documents": optional_rerank_params["documents"], + } + if optional_rerank_params.get("top_n") is not None: + request["top_n"] = optional_rerank_params["top_n"] + if optional_rerank_params.get("return_documents") is not None: + request["return_documents"] = optional_rerank_params["return_documents"] + return request + + def transform_rerank_response( + self, + model: str, + raw_response: httpx.Response, + model_response: RerankResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str] = None, + request_data: Optional[dict] = None, + optional_params: Optional[dict] = None, + litellm_params: Optional[dict] = None, + ) -> RerankResponse: + request_data = request_data or {} + optional_params = optional_params or {} + litellm_params = litellm_params or {} + try: + response_json = raw_response.json() + except Exception: + raise DashScopeError( + status_code=raw_response.status_code, + message=raw_response.text, + ) + + logging_obj.post_call( + input=request_data.get("query"), + api_key=api_key, + additional_args={"complete_input_dict": request_data}, + original_response=response_json, + ) + + # DashScope error envelope: {"code": "...", "message": "...", "request_id": "..."} + if "code" in response_json and "results" not in response_json: + raise DashScopeError( + status_code=raw_response.status_code, + message=response_json.get("message", str(response_json)), + ) + + results = response_json.get("results") + if results is None: + raise DashScopeError( + status_code=raw_response.status_code, + message=f"No results in DashScope rerank response: {response_json}", + ) + + # qwen3-rerank returns: + # {"index": int, "relevance_score": float} + # plus, when return_documents=true was sent: + # "document": {"text": "..."} + # which already matches LiteLLM's RerankResponseDocument shape. + transformed_results: List[dict] = [] + for r in results: + item: Dict[str, Any] = { + "index": r["index"], + "relevance_score": r["relevance_score"], + } + doc = r.get("document") + if isinstance(doc, dict): + item["document"] = doc + elif isinstance(doc, str): + # Defensive: spec says dict, but normalize string-shaped echoes. + item["document"] = {"text": doc} + transformed_results.append(item) + + usage = response_json.get("usage") or {} + total_tokens = usage.get("total_tokens") + billed_units = RerankBilledUnits(total_tokens=total_tokens) + tokens = RerankTokens(input_tokens=total_tokens) + meta = RerankResponseMeta(billed_units=billed_units, tokens=tokens) + + return RerankResponse( + id=response_json.get("id") or str(uuid.uuid4()), + results=transformed_results, # type: ignore + meta=meta, + ) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: Union[dict, httpx.Headers], + ) -> BaseLLMException: + if isinstance(headers, dict): + headers = httpx.Headers(headers) + return DashScopeError( + status_code=status_code, + message=error_message, + headers=headers, + ) diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index 5cd8d119542..7ed3e484535 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -2,13 +2,15 @@ Translates from OpenAI's `/v1/chat/completions` to DeepSeek's `/v1/chat/completions` """ -from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, overload +from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast, overload +import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +from litellm.utils import supports_reasoning from ...openai.chat.gpt_transformation import OpenAIGPTConfig @@ -62,6 +64,48 @@ class DeepSeekChatConfig(OpenAIGPTConfig): return optional_params + def _fill_reasoning_content( + self, messages: List[AllMessageValues] + ) -> List[AllMessageValues]: + """ + DeepSeek thinking mode requires `reasoning_content` to be passed back on + every assistant message in multi-turn conversations. If it is missing, + the API returns: + "The reasoning_content in the thinking mode must be passed back to the API." + + For each assistant message that is missing `reasoning_content`: + 1. Promote it from `provider_specific_fields["reasoning_content"]` if present + (LiteLLM stores provider-specific response fields there). + 2. Otherwise inject a single space — the minimum value the API accepts. + """ + result: List[AllMessageValues] = [] + for msg in messages: + if msg.get("role") == "assistant" and not msg.get("reasoning_content"): + patched = dict(cast(dict, msg)) + provider_fields = patched.get("provider_specific_fields") or {} + stored = provider_fields.get("reasoning_content") + if stored: + patched["reasoning_content"] = stored + cleaned = dict(provider_fields) + cleaned.pop("reasoning_content", None) + patched["provider_specific_fields"] = cleaned + else: + litellm.verbose_logger.warning( + "DeepSeek thinking mode: assistant message is missing " + "`reasoning_content` and none was saved in " + "`provider_specific_fields`. A single-space placeholder " + "is being injected to satisfy API validation, but the " + "model will receive a blank reasoning chain for this turn, " + "which may silently degrade multi-turn response quality. " + "Preserve `reasoning_content` from the original assistant " + "response when building multi-turn conversation history." + ) + patched["reasoning_content"] = " " + result.append(cast(AllMessageValues, patched)) + else: + result.append(msg) + return result + @overload def _transform_messages( self, messages: List[AllMessageValues], model: str, is_async: Literal[True] @@ -91,6 +135,66 @@ class DeepSeekChatConfig(OpenAIGPTConfig): messages=messages, model=model, is_async=False ) + def _thinking_mode_active(self, model: str, optional_params: dict) -> bool: + """ + Returns True only when thinking mode is actually active for this request: + - model supports reasoning (capability check) + - user explicitly passed thinking={"type": "enabled"} (opt-in check) + """ + return ( + supports_reasoning(model=model, custom_llm_provider="deepseek") + and (optional_params.get("thinking") or {}).get("type") == "enabled" + ) + + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Ensures `reasoning_content` is forwarded on assistant messages for + multi-turn thinking-mode conversations (issue #28045). + + Only runs when thinking mode is actually active - guarded by both + supports_reasoning() (model capability) and optional_params["thinking"] + (user explicitly enabled it), preventing spurious injection on models + like deepseek-v3.2 that support thinking as opt-in but not always-on. + """ + if self._thinking_mode_active(model=model, optional_params=optional_params): + messages = self._fill_reasoning_content(messages) + return super().transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + + async def async_transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Async equivalent of transform_request — applies the same reasoning_content + fix for multi-turn thinking-mode conversations. + """ + if self._thinking_mode_active(model=model, optional_params=optional_params): + messages = self._fill_reasoning_content(messages) + return await super().async_transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 72569e5c6cd..16e17dcc876 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -1,5 +1,7 @@ from typing import List, Optional, cast +import litellm + from litellm.litellm_core_utils.prompt_templates.factory import ( convert_generic_image_chunk_to_openai_image_obj, convert_to_anthropic_image_obj, @@ -101,7 +103,10 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): return supported_params def _transform_messages( - self, messages: List[AllMessageValues], model: Optional[str] = None + self, + messages: List[AllMessageValues], + model: Optional[str] = None, + litellm_params: Optional[dict] = None, ) -> List[ContentType]: """ Google AI Studio Gemini does not support HTTP/HTTPS URLs for files. @@ -141,14 +146,23 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): img_element["image_url"] = converted_image_url # type: ignore elif element.get("type") == "file": file_element = cast(ChatCompletionFileObject, element) - file_id = file_element["file"].get("file_id") + _file_field = file_element.get("file") + if _file_field is None: + raise litellm.BadRequestError( + message="Content block has type='file' but is missing the required 'file' field", + model=model, + llm_provider="gemini", + ) + file_id = _file_field.get("file_id") if file_id and ("http://" in file_id or "https://" in file_id): # Convert HTTP/HTTPS file URL to base64 data try: base64_data = convert_url_to_base64(file_id) - file_element["file"]["file_data"] = base64_data # type: ignore - file_element["file"].pop("file_id", None) # type: ignore + _file_field["file_data"] = base64_data # type: ignore + _file_field.pop("file_id", None) # type: ignore except Exception: # If conversion fails, leave as is and let the API handle it pass - return _gemini_convert_messages_with_history(messages=messages, model=model) + return _gemini_convert_messages_with_history( + messages=messages, model=model, litellm_params=litellm_params + ) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 6b7ec4dfb1c..5464b5bb7ee 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -287,7 +287,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): content_item["image_url"] = new_image_url_obj elif content_item.get("type") == "file": content_item = cast(ChatCompletionFileObject, content_item) - file_obj = content_item["file"] + file_obj = content_item.get("file") + if file_obj is None: + raise litellm.BadRequestError( + message="Content block has type='file' but is missing the required 'file' field", + model=None, + llm_provider="openai", + ) new_file_obj = ChatCompletionFileObjectFile( **{ # type: ignore k: v diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 9afa5dec465..f56992a2502 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -6,13 +6,16 @@ Why separate file? Make it easy to see how transformation works import json import os -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, cast +import re +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast +from urllib.parse import quote import httpx from pydantic import BaseModel import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( _get_image_mime_type_from_url, ) @@ -57,6 +60,45 @@ from ..common_utils import ( get_supports_system_message, ) +# Typed as Any to avoid introducing a module-load-time cyclic import to +# vertex_llm_base. The instance is lazily constructed by _get_vertex_base() +# the first time GCS metadata needs to be fetched. +_GCS_METADATA_VERTEX_BASE: Optional[Any] = None +# Shared sync client for GCS JSON API metadata reads so proxy/SSL settings +# from litellm's HTTP stack apply (see Greptile review on PR #27278). +_GCS_METADATA_HTTP_HANDLER: Optional[HTTPHandler] = None +_GEMINI_MIME_TYPE_ALIASES: Dict[str, str] = { + "image/jpg": "image/jpeg", +} + + +def _apply_gemini_mime_type_aliases(mime_type: str) -> str: + """Normalize known MIME aliases only; does not consult the file-type registry. + + Also strips MIME parameters (e.g. ``; charset=utf-8``) so that values + sourced from GCS object metadata (``contentType``) validate correctly. + """ + normalized = mime_type.split(";", 1)[0].strip().lower() + return _GEMINI_MIME_TYPE_ALIASES.get(normalized, normalized) + + +def _get_vertex_base() -> Any: + """Lazily return the shared VertexBase instance to avoid a module-load-time cyclic import.""" + global _GCS_METADATA_VERTEX_BASE + if _GCS_METADATA_VERTEX_BASE is None: + from ..vertex_llm_base import VertexBase + + _GCS_METADATA_VERTEX_BASE = VertexBase() + return _GCS_METADATA_VERTEX_BASE + + +def _get_gcs_metadata_http_handler() -> HTTPHandler: + global _GCS_METADATA_HTTP_HANDLER + if _GCS_METADATA_HTTP_HANDLER is None: + _GCS_METADATA_HTTP_HANDLER = HTTPHandler(timeout=5.0) + return _GCS_METADATA_HTTP_HANDLER + + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -171,12 +213,299 @@ def _apply_gemini_metadata( return cast(PartType, part_dict) +def _parse_gs_uri(gs_uri: str) -> Tuple[str, str]: + if not gs_uri.startswith("gs://"): + raise ValueError(f"Invalid gs URI: {gs_uri}") + uri_without_scheme = gs_uri[5:] # drop gs:// + uri_parts = uri_without_scheme.split("/", 1) + if len(uri_parts) != 2 or not uri_parts[0] or not uri_parts[1]: + raise ValueError(f"Invalid gs URI: {gs_uri}") + return uri_parts[0], uri_parts[1] + + +def _is_valid_gcs_bucket_name(bucket: str) -> bool: + """ + Validate bucket name against core GCS naming constraints. + """ + bucket_length = len(bucket) + max_bucket_length = 222 if "." in bucket else 63 + if bucket_length < 3 or bucket_length > max_bucket_length: + return False + if "." in bucket and any( + len(label) == 0 or len(label) > 63 for label in bucket.split(".") + ): + return False + if not re.fullmatch(r"[a-z0-9][a-z0-9._-]*[a-z0-9]", bucket): + return False + if ".." in bucket: + return False + if re.fullmatch(r"\d+\.\d+\.\d+\.\d+", bucket): + return False + return True + + +def _gs_uri_requires_content_type_metadata(url: str) -> bool: + """ + True when _process_gemini_media would call _get_gcs_object_content_type + (extension-less gs:// and no explicit format passed into that helper). + """ + if "gs://" not in url: + return False + extension_with_dot = os.path.splitext(url)[-1] + extension = extension_with_dot[1:] if extension_with_dot else "" + return len(extension) == 0 + + +def _image_url_payload_may_need_sync_gcs_metadata_fetch( + raw_image_url: Any, +) -> bool: + """ + True when this image_url value (content-part image_url or assistant ``images[]`` + entry) can trigger a blocking GCS metadata read for MIME resolution. + """ + fmt: Optional[str] = None + url: Optional[str] = None + if isinstance(raw_image_url, dict): + url = raw_image_url.get("url") # type: ignore[assignment] + if not isinstance(url, str): + return False + fmt = ( + raw_image_url.get("format") + or raw_image_url.get("mime_type") + or raw_image_url.get("content_type") + ) + elif isinstance(raw_image_url, str): + url = raw_image_url + else: + return False + if "gs://" not in url or fmt: + return False + return _gs_uri_requires_content_type_metadata(url) + + +def _openai_messages_may_need_sync_gcs_metadata_fetch( + messages: List[AllMessageValues], +) -> bool: + """ + Heuristic: True if any message part can trigger a blocking GCS JSON + metadata read inside _transform_request_body (extension-less gs:// without + explicit MIME hints). Covers user/system ``content`` parts and assistant + ``images`` (same paths as ``_gemini_convert_messages_with_history``). Used + to decide whether ``async_transform_request_body`` should offload the sync + transform via ``asyncify``. + """ + for raw in messages: + msg: Any = raw + if not isinstance(msg, dict) and hasattr(msg, "model_dump"): + msg = msg.model_dump(exclude_none=False) + if not isinstance(msg, dict): + continue + images_field = msg.get("images") + if isinstance(images_field, list): + for image_item in images_field: + if not isinstance(image_item, dict): + continue + if _image_url_payload_may_need_sync_gcs_metadata_fetch( + image_item.get("image_url") + ): + return True + + content = msg.get("content") + if not isinstance(content, list): + continue + for item in content: + if not isinstance(item, dict): + continue + itype = item.get("type") + if itype == "image_url": + if _image_url_payload_may_need_sync_gcs_metadata_fetch( + item.get("image_url") + ): + return True + elif itype == "file": + file_obj = item.get("file") + if not isinstance(file_obj, dict): + continue + fmt = ( + file_obj.get("format") + or file_obj.get("mime_type") + or file_obj.get("content_type") + ) + passed = file_obj.get("file_id") or file_obj.get("file_data") + if ( + isinstance(passed, str) + and "gs://" in passed + and not fmt + and _gs_uri_requires_content_type_metadata(passed) + ): + return True + return False + + +def _get_gcs_object_content_type( + image_url: str, + vertex_project: Optional[str] = None, + vertex_credentials: Optional[Any] = None, +) -> Optional[str]: + """ + Resolve content type from GCS object metadata. + + Only attaches a Bearer token when the caller explicitly supplies Vertex + credentials, to avoid using the server's default Google credentials on + the Gemini API-key (Google AI Studio) path and being used as an oracle + for private GCS object metadata. Without explicit credentials we only + issue an anonymous request, which only succeeds for publicly-readable + objects. + """ + try: + bucket, object_name = _parse_gs_uri(image_url) + except ValueError: + return None + if not _is_valid_gcs_bucket_name(bucket): + return None + + headers: Dict[str, str] = {} + explicit_vertex_auth_provided = ( + vertex_project is not None or vertex_credentials is not None + ) + if explicit_vertex_auth_provided: + try: + access_token, _ = _get_vertex_base().get_access_token( + credentials=vertex_credentials, + project_id=vertex_project, + ) + headers["Authorization"] = f"Bearer {access_token}" + except Exception as e: + raise litellm.BadRequestError( + message=( + "Unable to fetch GCS metadata with provided Vertex credentials/project. " + f"Original error: {str(e)}" + ), + model=None, + llm_provider="vertex_ai", + ) + + # Build the URL via httpx.URL with a fixed scheme/host and URL-encode both + # bucket and object so CodeQL does not flag the interpolation as a + # potential SSRF that could resolve to an arbitrary host. + encoded_bucket = quote(bucket, safe="") + encoded_object = quote(object_name, safe="") + metadata_url = httpx.URL( + scheme="https", + host="storage.googleapis.com", + path=f"/storage/v1/b/{encoded_bucket}/o/{encoded_object}", + params={"fields": "contentType"}, + ) + try: + response = _get_gcs_metadata_http_handler().get( + url=str(metadata_url), + headers=headers or None, + ) + except httpx.RequestError as e: + if explicit_vertex_auth_provided: + raise litellm.BadRequestError( + message=( + "Unable to reach GCS JSON API for object metadata with provided " + f"Vertex credentials. {type(e).__name__}: {e}" + ), + model=None, + llm_provider="vertex_ai", + ) from e + return None + + if response.is_error: + if explicit_vertex_auth_provided: + preview = (response.text or "")[:1024] + raise litellm.BadRequestError( + message=( + "Unable to read GCS object metadata with provided Vertex credentials. " + f"HTTP {response.status_code}. Response body (truncated): {preview!r}" + ), + model=None, + llm_provider="vertex_ai", + ) + return None + + try: + payload = response.json() + except ValueError as e: + if explicit_vertex_auth_provided: + raise litellm.BadRequestError( + message=( + "GCS metadata response was not valid JSON when using provided " + f"Vertex credentials (HTTP {response.status_code}). Error: {e}" + ), + model=None, + llm_provider="vertex_ai", + ) from e + return None + + if not isinstance(payload, dict): + if explicit_vertex_auth_provided: + raise litellm.BadRequestError( + message=( + "GCS metadata response was not a JSON object when using provided " + f"Vertex credentials (HTTP {response.status_code})." + ), + model=None, + llm_provider="vertex_ai", + ) + return None + + content_type = payload.get("contentType") + if isinstance(content_type, str) and len(content_type) > 0: + return content_type + + if explicit_vertex_auth_provided: + preview = (response.text or "")[:1024] + raise litellm.BadRequestError( + message=( + "GCS metadata JSON did not include a non-empty contentType field when " + f"using provided Vertex credentials (HTTP {response.status_code}). " + f"Body (truncated): {preview!r}" + ), + model=None, + llm_provider="vertex_ai", + ) + return None + + +def _normalize_and_validate_gemini_mime_type( + mime_type: str, model: Optional[str] +) -> str: + # Import lazily to avoid a module-level cyclic-import alert with + # litellm.types.files. + from litellm.types.files import get_file_extension_from_mime_type + + normalized_mime_type = _apply_gemini_mime_type_aliases(mime_type) + try: + file_extension = get_file_extension_from_mime_type(normalized_mime_type) + file_type = get_file_type_from_extension(file_extension) + except ValueError: + raise litellm.BadRequestError( + message=f"File type not supported by gemini - {normalized_mime_type}", + model=model, + llm_provider="vertex_ai", + ) + + if not is_gemini_1_5_accepted_file_type(file_type): + raise litellm.BadRequestError( + message=f"File type not supported by gemini - {file_type}", + model=model, + llm_provider="vertex_ai", + ) + + return get_file_mime_type_for_file_type(file_type) + + def _process_gemini_media( image_url: str, format: Optional[str] = None, media_resolution_enum: Optional[Dict[str, str]] = None, model: Optional[str] = None, video_metadata: Optional[Dict[str, Any]] = None, + vertex_project: Optional[str] = None, + vertex_credentials: Optional[Any] = None, ) -> PartType: """ Given a media URL (image, audio, or video), return the appropriate PartType for Gemini @@ -193,20 +522,63 @@ def _process_gemini_media( try: # GCS URIs if "gs://" in image_url: - # Figure out file type extension_with_dot = os.path.splitext(image_url)[-1] # Ex: ".png" extension = extension_with_dot[1:] # Ex: "png" + explicit_gcs_format = False if not format: - file_type = get_file_type_from_extension(extension) + mime_type: Optional[str] = None + # For extension-less gs:// URIs, we cannot infer from path. + # If callers pass `format`/`mime_type`, this branch is skipped. + if extension: + file_type = get_file_type_from_extension(extension) - # Validate the file type is supported by Gemini - if not is_gemini_1_5_accepted_file_type(file_type): - raise Exception(f"File type not supported by gemini - {file_type}") + # Validate the file type is supported by Gemini + if not is_gemini_1_5_accepted_file_type(file_type): + raise litellm.BadRequestError( + message=f"File type not supported by gemini - {file_type}", + model=model, + llm_provider="vertex_ai", + ) - mime_type = get_file_mime_type_for_file_type(file_type) + mime_type = get_file_mime_type_for_file_type(file_type) + else: + mime_type = _get_gcs_object_content_type( + image_url=image_url, + vertex_project=vertex_project, + vertex_credentials=vertex_credentials, + ) + if mime_type is None: + raise litellm.BadRequestError( + message=( + f"Unable to determine mime type for gs URI: {image_url}. " + "This gs:// URI has no file extension and GCS metadata " + "lookup failed. Set it explicitly using image_url.format " + "(or image_url.mime_type/content_type) or " + "message.content[].file.format." + ), + model=model, + llm_provider="vertex_ai", + ) else: mime_type = format + explicit_gcs_format = True + if mime_type is None: + raise litellm.BadRequestError( + message=f"File type not supported by gemini - {image_url}", + model=model, + llm_provider="vertex_ai", + ) + if explicit_gcs_format: + # Callers who pass format/mime_type explicitly for gs:// URIs + # rely on pass-through to Gemini (pre-PR behavior). Only apply + # known MIME aliases; skip litellm's file-type registry. + mime_type = _apply_gemini_mime_type_aliases(mime_type) + else: + mime_type = _normalize_and_validate_gemini_mime_type( + mime_type=mime_type, + model=model, + ) file_data = FileDataType(mime_type=mime_type, file_uri=image_url) part: PartType = {"file_data": file_data} return _apply_gemini_metadata( @@ -258,8 +630,6 @@ def _snake_to_camel(snake_str: str) -> str: def _camel_to_snake(camel_str: str) -> str: """Convert camelCase to snake_case""" - import re - return re.sub(r"(? List[ContentType]: """ Converts given messages from OpenAI format to Gemini format @@ -326,6 +697,16 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 msg_i = 0 tool_call_responses = [] + vertex_project = None + vertex_credentials = None + if litellm_params: + vertex_project = litellm_params.get("vertex_project") or litellm_params.get( + "vertex_ai_project" + ) + vertex_credentials = litellm_params.get( + "vertex_credentials" + ) or litellm_params.get("vertex_ai_credentials") + try: while msg_i < len(messages): user_content: List[PartType] = [] @@ -351,20 +732,42 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 img_element = element format: Optional[str] = None media_resolution_enum: Optional[Dict[str, str]] = None - if isinstance(img_element["image_url"], dict): - image_url = img_element["image_url"]["url"] - format = img_element["image_url"].get("format") - detail = img_element["image_url"].get("detail") + raw_image_url = img_element.get("image_url") + if raw_image_url is None: + raise litellm.BadRequestError( + message="Invalid message content: element type is 'image_url' but 'image_url' field is missing ", + model=model, + llm_provider="vertex_ai", + ) + if isinstance(raw_image_url, dict): + image_url = raw_image_url.get("url") + if image_url is None: + raise litellm.BadRequestError( + message="Invalid message content: element type is 'image_url' but 'url' field is missing inside 'image_url' ", + model=model, + llm_provider="vertex_ai", + ) + # TypedDict does not declare mime_type/content_type; + # read via Dict[str, Any] for caller-provided MIME fields. + image_url_dict = cast(Dict[str, Any], raw_image_url) + format = ( + image_url_dict.get("format") + or image_url_dict.get("mime_type") + or image_url_dict.get("content_type") + ) + detail = image_url_dict.get("detail") media_resolution_enum = ( _convert_detail_to_media_resolution_enum(detail) ) else: - image_url = img_element["image_url"] + image_url = raw_image_url _part = _process_gemini_media( image_url=image_url, format=format, media_resolution_enum=media_resolution_enum, model=model, + vertex_project=vertex_project, + vertex_credentials=vertex_credentials, ) _parts.append(_part) elif element["type"] == "input_audio": @@ -390,15 +793,31 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 image_url=openai_image_str, format=audio_format_modified, model=model, + vertex_project=vertex_project, + vertex_credentials=vertex_credentials, ) _parts.append(_part) elif element["type"] == "file": file_element = cast(ChatCompletionFileObject, element) - file_id = file_element["file"].get("file_id") - format = file_element["file"].get("format") - file_data = file_element["file"].get("file_data") - detail = file_element["file"].get("detail") - video_metadata = file_element["file"].get("video_metadata") + _file_field = file_element.get("file") + if _file_field is None: + raise litellm.BadRequestError( + message="Content block has type='file' but is missing the required 'file' field", + model=model, + llm_provider="vertex_ai", + ) + # TypedDict does not declare mime_type/content_type; + # read via Dict[str, Any] for caller-provided MIME fields. + file_dict = cast(Dict[str, Any], _file_field) + file_id = file_dict.get("file_id") + format = ( + file_dict.get("format") + or file_dict.get("mime_type") + or file_dict.get("content_type") + ) + file_data = file_dict.get("file_data") + detail = file_dict.get("detail") + video_metadata = file_dict.get("video_metadata") passed_file = file_id or file_data if passed_file is None: raise Exception( @@ -417,13 +836,23 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 model=model, media_resolution_enum=media_resolution_enum, video_metadata=video_metadata, + vertex_project=vertex_project, + vertex_credentials=vertex_credentials, ) _parts.append(_part) - except Exception: - raise Exception( - "Unable to determine mime type for file_id: {}, set this explicitly using message[{}].content[{}].file.format".format( - file_id, msg_i, element_idx - ) + except litellm.BadRequestError: + raise + except Exception as e: + raise litellm.BadRequestError( + message=( + f"Unable to determine mime type for file: " + f"{file_id or 'provided data'}, set this explicitly " + f"using message[{msg_i}].content[{element_idx}].file.format " + f"(or file.mime_type/content_type). " + f"Original error: {str(e)}" + ), + model=model, + llm_provider="vertex_ai", ) user_content.extend(_parts) elif _message_content is not None and isinstance(_message_content, str): @@ -528,7 +957,11 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 image_url_obj = image_item.get("image_url") if isinstance(image_url_obj, dict): assistant_image_url = image_url_obj.get("url") - format = image_url_obj.get("format") + format = ( + image_url_obj.get("format") + or image_url_obj.get("mime_type") + or image_url_obj.get("content_type") + ) detail = image_url_obj.get("detail") media_resolution_enum = ( _convert_detail_to_media_resolution_enum(detail) @@ -539,6 +972,8 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 format=format, media_resolution_enum=media_resolution_enum, model=model, + vertex_project=vertex_project, + vertex_credentials=vertex_credentials, ) assistant_content.append(_part) @@ -713,11 +1148,11 @@ def _transform_request_body( # noqa: PLR0915 try: if custom_llm_provider == "gemini": content = litellm.GoogleAIStudioGeminiConfig()._transform_messages( - messages=messages, model=model + messages=messages, model=model, litellm_params=litellm_params ) else: content = litellm.VertexGeminiConfig()._transform_messages( - messages=messages, model=model + messages=messages, model=model, litellm_params=litellm_params ) tools: Optional[Tools] = optional_params.pop("tools", None) tool_choice: Optional[ToolConfig] = optional_params.pop("tool_choice", None) @@ -893,6 +1328,20 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) + if _openai_messages_may_need_sync_gcs_metadata_fetch(messages): + # _transform_request_body may issue a sync httpx.get (up to 5s timeout) + # via _get_gcs_object_content_type to fetch GCS object metadata. Run the + # whole sync transformation on a worker thread so it does not block the + # async event loop. + return await asyncify(_transform_request_body)( + messages=messages, + model=model, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + cached_content=cached_content, + optional_params=optional_params, + ) + return _transform_request_body( messages=messages, model=model, 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 6278de662f8..49c1c335467 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 @@ -2533,9 +2533,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return model_response def _transform_messages( - self, messages: List[AllMessageValues], model: Optional[str] = None + self, + messages: List[AllMessageValues], + model: Optional[str] = None, + litellm_params: Optional[dict] = None, ) -> List[ContentType]: - return _gemini_convert_messages_with_history(messages=messages, model=model) + return _gemini_convert_messages_with_history( + messages=messages, model=model, litellm_params=litellm_params + ) def get_error_class( self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers] @@ -3139,6 +3144,31 @@ class ModelResponseIterator: self.cumulative_tool_call_index: int = 0 self.has_seen_tool_calls: bool = False + @staticmethod + def _check_streaming_error(chunk: dict) -> None: + """Detect embedded errors (e.g. 429 RESOURCE_EXHAUSTED) in streaming chunks and raise VertexAIError.""" + if "error" not in chunk: + return + error_data = chunk["error"] + if not isinstance(error_data, dict): + raise VertexAIError( + status_code=500, + message=f"Unexpected error format in mid-stream chunk: {error_data}", + ) + raw_code = error_data.get("code", 500) + if raw_code is None: + raw_code = 500 + try: + error_code = int(raw_code) + except (TypeError, ValueError): + error_code = 500 + error_message = error_data.get("message", "Unknown error") + error_status = error_data.get("status", "UNKNOWN") + raise VertexAIError( + status_code=error_code, + message=f"{error_status} - {error_message}", + ) + def _apply_stream_candidates( self, _candidates: List[Candidates], @@ -3256,6 +3286,11 @@ class ModelResponseIterator: def chunk_parser(self, chunk: dict) -> Optional["ModelResponseStream"]: try: verbose_logger.debug(f"RAW GEMINI CHUNK: {chunk}") + + # Detect mid-stream error chunks (e.g. 429 RESOURCE_EXHAUSTED). + # Vertex AI can return errors as HTTP 200 but with an "error" field in the SSE body. + self._check_streaming_error(chunk) + from litellm.types.utils import ModelResponseStream processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore 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 123d925f7c1..eb67e3aa828 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -292,22 +292,15 @@ class VertexAIPartnerModels(VertexBase): Returns: Dict containing token count information """ - try: - import vertexai - except Exception as e: - raise VertexAIError( - status_code=400, - message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""", - ) - - if not ( - hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models") - ): - raise VertexAIError( - status_code=400, - message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", - ) - + # Note: we intentionally do not import `vertexai` (the Gemini SDK shipped + # by `google-cloud-aiplatform`) on this path. Partner models such as + # Claude on Vertex use the Anthropic Messages API protocol directly via + # `:rawPredict`, and `VertexAIPartnerModelsTokenCounter` reaches that + # endpoint with an authenticated httpx client — it never touches the + # Gemini SDK. Requiring `google-cloud-aiplatform>=1.38` here turned a + # SDK-free Anthropic-protocol call into a hard dependency on the Gemini + # SDK (see #28084), breaking `/v1/messages/count_tokens` for Claude-on- + # Vertex on any LiteLLM install without that extra. Stay SDK-free. try: from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import ( VertexAIPartnerModelsTokenCounter, diff --git a/litellm/main.py b/litellm/main.py index c3d1c2e05b0..b5364f8ba17 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5720,6 +5720,33 @@ def embedding( # noqa: PLR0915 aembedding=aembedding, headers=headers, ) + elif custom_llm_provider == "dashscope": + dashscope_key = ( + api_key or litellm.api_key or get_secret_str("DASHSCOPE_API_KEY") + ) + if dashscope_key is None: + raise ValueError( + "Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter." + ) + if extra_headers is not None and isinstance(extra_headers, dict): + headers = extra_headers + else: + headers = {} + response = base_llm_http_handler.embedding( + model=model, + input=input, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + logging_obj=logging, + api_base=api_base, + optional_params=optional_params, + litellm_params={}, + model_response=EmbeddingResponse(), + api_key=dashscope_key, + client=client, + aembedding=aembedding, + headers=headers, + ) elif custom_llm_provider == "ovhcloud": api_key = api_key or litellm.api_key or get_secret_str("OVHCLOUD_API_KEY") api_base = ( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d4fa497698a..c989f5dff13 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4549,6 +4549,7 @@ class PrismaCompatibleUpdateDBModel(TypedDict, total=False): model_name: str litellm_params: str model_info: str + blocked: bool updated_at: str updated_by: str diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 7611c9c9692..e7f14df5294 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -69,6 +69,35 @@ else: ProxyLogging = Any +def _extract_cache_read_tokens(usage_obj: dict) -> int: + """ + Anthropic: top-level cache_read_input_tokens field. + OpenAI-compatible (moonshotai, openai, deepseek, etc.): prompt_tokens_details.cached_tokens. + """ + explicit = usage_obj.get("cache_read_input_tokens", 0) or 0 + if explicit: + return int(explicit) + details = usage_obj.get("prompt_tokens_details") or {} + return int(details.get("cached_tokens", 0) or 0) + + +def _extract_cache_creation_tokens(usage_obj: dict) -> int: + """ + Anthropic: top-level cache_creation_input_tokens field. + OpenAI-compatible (kimi-k2 etc.): prompt_tokens_details.cache_write_tokens + or prompt_tokens_details.cache_creation_tokens. + """ + explicit = usage_obj.get("cache_creation_input_tokens", 0) or 0 + if explicit: + return int(explicit) + details = usage_obj.get("prompt_tokens_details") or {} + return int( + details.get("cache_write_tokens", 0) + or details.get("cache_creation_tokens", 0) + or 0 + ) + + class DBSpendUpdateWriter: """ Module responsible for @@ -1992,12 +2021,8 @@ class DBSpendUpdateWriter: api_requests=1, successful_requests=1 if request_status == "success" else 0, failed_requests=1 if request_status != "success" else 0, - cache_read_input_tokens=usage_obj.get("cache_read_input_tokens", 0) - or 0, - cache_creation_input_tokens=usage_obj.get( - "cache_creation_input_tokens", 0 - ) - or 0, + cache_read_input_tokens=_extract_cache_read_tokens(usage_obj), + cache_creation_input_tokens=_extract_cache_creation_tokens(usage_obj), ) return daily_transaction except Exception as e: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index af84bc123ff..472306eb818 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -150,6 +150,9 @@ def update_db_model( model_info[key] = value.isoformat() prisma_compatible_model_dict["model_info"] = json.dumps(model_info) + if updated_patch.blocked is not None: + prisma_compatible_model_dict["blocked"] = updated_patch.blocked + return prisma_compatible_model_dict @@ -230,6 +233,20 @@ async def patch_model( premium_user=premium_user, ) + # Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins + # passed the auth check above for team-scoped models, but they must not + # be able to unblock (or block) a model their proxy admin has paused. + if ( + patch_data.blocked is not None + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN + ): + raise ProxyException( + message="Only proxy admins can change a model's blocked flag.", + type=ProxyErrorTypes.auth_error.value, + code=status.HTTP_403_FORBIDDEN, + param="blocked", + ) + # Handle team model updates with proper alias management update_data = await _update_team_model_in_db( db_model=db_model, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 97252b4c977..5d89d3fa9c5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4727,6 +4727,7 @@ class ProxyConfig: if _id is not None: model.model_info["id"] = _id model.model_info["db_model"] = True + model.model_info["blocked"] = bool(getattr(model, "blocked", False)) if premium_user is True: # seeing "created_at", "updated_at", "created_by", "updated_by" is a LiteLLM Enterprise Feature @@ -8089,6 +8090,11 @@ async def model_list( proxy_logging_obj=proxy_logging_obj, ) + # Compute once — used in both branches below to hide paused models from the listing. + blocked_names = ( + llm_router.get_fully_blocked_model_names() if llm_router is not None else set() + ) + # If scope=expand and user has admin privileges, return all proxy models if should_expand_scope: # Get all proxy models as if user is a proxy admin @@ -8121,6 +8127,10 @@ async def model_list( only_model_access_groups=only_model_access_groups or False, ) + # Hide paused models from the public listing (admins manage them via /model/info) + if blocked_names: + all_models = [m for m in all_models if m not in blocked_names] + # Build response data with all proxy models model_data = [] for model in all_models: @@ -8154,6 +8164,10 @@ async def model_list( user_api_key_cache=user_api_key_cache, ) + # Hide paused models from the public listing (admins manage them via /model/info) + if blocked_names: + all_models = [m for m in all_models if m not in blocked_names] + # Build response data model_data = [] for model in all_models: diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index f86cde87401..06fd35448f1 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -430,7 +430,11 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin deployment = llm_router.get_deployment_by_model_group_name( model_group_name=model ) - if deployment and deployment.litellm_params: + if ( + deployment + and deployment.litellm_params + and not llm_router._is_deployment_blocked(deployment) + ): deployment_creds = deployment.litellm_params.model_dump( exclude_none=True ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b53507abe6a..78143fe0411 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable { // Models on proxy model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) - model_name String + model_name String litellm_params Json - model_info Json? + model_info Json? + blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/router.py b/litellm/router.py index fac48b45fb3..420c9b8a816 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -30,6 +30,7 @@ from typing import ( List, Literal, Optional, + Set, Tuple, Union, cast, @@ -6957,12 +6958,11 @@ class Router: unhealthy_deployments = _get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) - healthy_deployments: list = [] - for deployment in _all_deployments: - if deployment["model_info"]["id"] in unhealthy_deployments: - continue - else: - healthy_deployments.append(deployment) + unhealthy_set = set(unhealthy_deployments) + healthy_deployments: list = [ + d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_set + ] + healthy_deployments = self._filter_blocked_deployments(healthy_deployments) return healthy_deployments, _all_deployments @@ -6990,10 +6990,12 @@ class Router: ) # Convert to set for O(1) lookup instead of O(n) unhealthy_deployments_set = set(unhealthy_deployments) - healthy_deployments: list = [] - for deployment in _all_deployments: - if deployment["model_info"]["id"] not in unhealthy_deployments_set: - healthy_deployments.append(deployment) + healthy_deployments: list = [ + d + for d in _all_deployments + if d["model_info"]["id"] not in unhealthy_deployments_set + ] + healthy_deployments = self._filter_blocked_deployments(healthy_deployments) return healthy_deployments, _all_deployments def routing_strategy_pre_call_checks(self, deployment: dict): @@ -8142,10 +8144,14 @@ class Router: def get_deployment_credentials(self, model_id: str) -> Optional[dict]: """ - Returns -> dict of credentials for a given model id + Returns -> dict of credentials for a given model id. + + Returns None if the deployment is paused via `LiteLLM_ProxyModelTable.blocked`, + so file/batch/passthrough callers that resolve credentials directly cannot keep + using a paused deployment. """ deployment = self.get_deployment(model_id=model_id) - if deployment is None: + if deployment is None or self._is_deployment_blocked(deployment): return None return CredentialLiteLLMParams( **deployment.litellm_params.model_dump(exclude_none=True) @@ -8190,7 +8196,9 @@ class Router: Returns: Dictionary containing api_key, api_base, custom_llm_provider, etc. - Returns None if model not found. + Returns None if model not found, or if the resolved deployment is + paused via `LiteLLM_ProxyModelTable.blocked` (so passthrough callers + cannot bypass an admin pause by resolving credentials directly). Example: credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm") @@ -8216,7 +8224,7 @@ class Router: elif isinstance(deployment_dict, Deployment): deployment = deployment_dict - if deployment is None: + if deployment is None or self._is_deployment_blocked(deployment): return None # Get basic credentials @@ -9243,6 +9251,29 @@ class Router: return model_names + def get_fully_blocked_model_names(self) -> Set[str]: + """ + Returns the set of model_names where every backing deployment has `blocked=True`. + + Used by `/v1/models` to hide paused models from client listings while still + surfacing them on admin endpoints (e.g. `/model/info`). A model with at least + one non-blocked deployment is still serviceable and remains visible. + """ + deployments = self.get_model_list() or [] + blocked_by_name: Dict[str, bool] = {} + for deployment in deployments: + name = deployment.get("model_name") or "" + if not name: + continue + is_blocked = (deployment.get("model_info") or {}).get("blocked") is True + if name in blocked_by_name: + blocked_by_name[name] = blocked_by_name[name] and is_blocked + else: + blocked_by_name[name] = is_blocked + return { + name for name, fully_blocked in blocked_by_name.items() if fully_blocked + } + def _get_team_specific_model( self, deployment: DeploymentTypedDict, team_id: Optional[str] = None ) -> Optional[str]: @@ -10131,6 +10162,12 @@ class Router: ) if isinstance(healthy_deployments, dict): + if (healthy_deployments.get("model_info") or {}).get("blocked") is True: + raise litellm.ServiceUnavailableError( + message=f"Model '{model}' is administratively paused. Contact your proxy admin to unblock it.", + model=model, + llm_provider="", + ) return healthy_deployments # Health-check-based filtering (before cooldown) @@ -10164,6 +10201,8 @@ class Router: ) healthy_deployments = _pre_cooldown_deployments + healthy_deployments = self._filter_blocked_deployments(healthy_deployments) + healthy_deployments = await self.async_callback_filter_deployments( model=model, healthy_deployments=healthy_deployments, @@ -10387,6 +10426,12 @@ class Router: # 3. If specific deployment returned, verify if it supports pass-through if isinstance(healthy_deployments, dict): + if (healthy_deployments.get("model_info") or {}).get("blocked") is True: + raise litellm.ServiceUnavailableError( + message=f"Model '{model}' is administratively paused. Contact your proxy admin to unblock it.", + model=model, + llm_provider="", + ) litellm_params = healthy_deployments.get("litellm_params", {}) if litellm_params.get("use_in_pass_through"): return healthy_deployments @@ -10555,6 +10600,12 @@ class Router: ) if isinstance(healthy_deployments, dict): + if (healthy_deployments.get("model_info") or {}).get("blocked") is True: + raise litellm.ServiceUnavailableError( + message=f"Model '{model}' is administratively paused. Contact your proxy admin to unblock it.", + model=model, + llm_provider="", + ) return healthy_deployments parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs( @@ -10585,6 +10636,8 @@ class Router: ) healthy_deployments = _pre_cooldown_deployments + healthy_deployments = self._filter_blocked_deployments(healthy_deployments) + # filter pre-call checks if self.enable_pre_call_checks and messages is not None: healthy_deployments = self._pre_call_checks( @@ -10704,6 +10757,12 @@ class Router: # 2. If the returned is a specific deployment (Dict), verify and return directly if isinstance(healthy_deployments, dict): + if (healthy_deployments.get("model_info") or {}).get("blocked") is True: + raise litellm.ServiceUnavailableError( + message=f"Model '{model}' is administratively paused. Contact your proxy admin to unblock it.", + model=model, + llm_provider="", + ) litellm_params = healthy_deployments.get("litellm_params", {}) if litellm_params.get("use_in_pass_through"): return healthy_deployments @@ -10743,6 +10802,9 @@ class Router: healthy_deployments=pass_through_deployments, cooldown_deployments=cooldown_deployments, ) + pass_through_deployments = self._filter_blocked_deployments( + pass_through_deployments + ) # 5. Apply pre-call checks (if enabled) if self.enable_pre_call_checks and messages is not None: @@ -10832,6 +10894,36 @@ class Router: if deployment["model_info"]["id"] not in cooldown_set ] + def _filter_blocked_deployments( + self, healthy_deployments: List[Dict] + ) -> List[Dict]: + """ + Filters out deployments that an admin has paused via `LiteLLM_ProxyModelTable.blocked`. + + Applied alongside the cooldown filter on every routing entry point that calls + `_common_checks_available_deployment` directly — the primary sync/async path, + the sync pass-through path, and the retry / health-check helpers — so paused + deployments never serve a request. The async pass-through path inherits this + filter through its delegation to `async_get_healthy_deployments`. + """ + return [ + deployment + for deployment in healthy_deployments + if (deployment.get("model_info") or {}).get("blocked") is not True + ] + + @staticmethod + def _is_deployment_blocked(deployment: "Deployment") -> bool: + """ + Returns True when a `Deployment` Pydantic instance carries the admin-paused + flag. Used by credential-lookup helpers so passthrough file / batch endpoints + cannot bypass the pause by resolving credentials directly. + """ + model_info = getattr(deployment, "model_info", None) + if model_info is None: + return False + return getattr(model_info, "blocked", None) is True + async def _async_filter_health_check_unhealthy_deployments( self, healthy_deployments: List[Dict], diff --git a/litellm/types/router.py b/litellm/types/router.py index 926815ba317..6601f552b52 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -133,6 +133,9 @@ class ModelInfo(BaseModel): # the model_name that can be used by the team when making LLM calls team_public_model_name: Optional[str] = None + # admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked + blocked: Optional[bool] = None + def __init__(self, id: Optional[Union[str, int]] = None, **params): if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided @@ -323,6 +326,7 @@ class updateDeployment(BaseModel): model_name: Optional[str] = None litellm_params: Optional[updateLiteLLMParams] = None model_info: Optional[ModelInfo] = None + blocked: Optional[bool] = None model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 400edcac889..6084f14e2df 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -107,9 +107,13 @@ class LiteLLMCommonStrings(Enum): SupportedCacheControls = ["ttl", "s-maxage", "no-cache", "no-store"] -class CostPerToken(TypedDict): - input_cost_per_token: float - output_cost_per_token: float +class CostPerToken(TypedDict, total=False): + # Required base rates — kept under total=False so we can mark them + # Required individually while leaving the cache rates NotRequired. + input_cost_per_token: Required[float] + output_cost_per_token: Required[float] + cache_read_input_token_cost: float + cache_creation_input_token_cost: float class ProviderField(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index cefd348078b..54cea313b0c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8402,6 +8402,12 @@ class ProviderConfigManager: ) return VolcEngineEmbeddingConfig() + elif litellm.LlmProviders.DASHSCOPE == provider: + from litellm.llms.dashscope.embed.transformation import ( + DashScopeEmbeddingConfig, + ) + + return DashScopeEmbeddingConfig() elif litellm.LlmProviders.OVHCLOUD == provider: return litellm.OVHCloudEmbeddingConfig() elif litellm.LlmProviders.SNOWFLAKE == provider: @@ -8481,6 +8487,12 @@ class ProviderConfigManager: return litellm.VoyageRerankConfig() elif litellm.LlmProviders.WATSONX == provider: return litellm.IBMWatsonXRerankConfig() + elif litellm.LlmProviders.DASHSCOPE == provider: + from litellm.llms.dashscope.rerank.transformation import ( + DashScopeRerankConfig, + ) + + return DashScopeRerankConfig() return litellm.CohereRerankConfig() @staticmethod diff --git a/schema.prisma b/schema.prisma index b53507abe6a..78143fe0411 100644 --- a/schema.prisma +++ b/schema.prisma @@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable { // Models on proxy model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) - model_name String + model_name String litellm_params Json - model_info Json? + model_info Json? + blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/tests/llm_translation/test_deepseek_completion.py b/tests/llm_translation/test_deepseek_completion.py index da402a51b68..2ede5d3f3f8 100644 --- a/tests/llm_translation/test_deepseek_completion.py +++ b/tests/llm_translation/test_deepseek_completion.py @@ -176,3 +176,113 @@ def test_completion_cost_deepseek(): pass except Exception as e: pytest.fail(f"Error occurred: {e}") + + +def test_deepseek_fill_reasoning_content_multiturn(): + """ + Unit test for _fill_reasoning_content. + Reproduces issue #28045: DeepSeek thinking mode fails in multi-turn conversations + because reasoning_content is not passed back to the API. + """ + from litellm.llms.deepseek.chat.transformation import DeepSeekChatConfig + + config = DeepSeekChatConfig() + + # Case 1: assistant message already has reasoning_content — should be left as-is + messages_with_rc = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi", "reasoning_content": "I thought about it"}, + {"role": "user", "content": "Follow up"}, + ] + result = config._fill_reasoning_content(messages_with_rc) + assert result[1]["reasoning_content"] == "I thought about it" + + # Case 2: assistant message has reasoning_content in provider_specific_fields — should be promoted + messages_with_psf = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": "Hi", + "provider_specific_fields": {"reasoning_content": "stored thinking"}, + }, + {"role": "user", "content": "Follow up"}, + ] + result = config._fill_reasoning_content(messages_with_psf) + assert result[1]["reasoning_content"] == "stored thinking" + # Should be removed from provider_specific_fields to avoid duplication + assert "reasoning_content" not in result[1].get("provider_specific_fields", {}) + + # Case 3: assistant message has no reasoning_content anywhere — should inject placeholder + messages_no_rc = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + {"role": "user", "content": "Follow up"}, + ] + result = config._fill_reasoning_content(messages_no_rc) + assert result[1]["reasoning_content"] == " " + + # Case 4: non-assistant messages should never be touched + messages_user_only = [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "You are helpful"}, + ] + result = config._fill_reasoning_content(messages_user_only) + assert "reasoning_content" not in result[0] + assert "reasoning_content" not in result[1] + + +def test_deepseek_fill_reasoning_content_guard_in_transform_request(): + """ + _fill_reasoning_content must only run when BOTH conditions are true: + 1. supports_reasoning() is True for the model + 2. thinking mode is explicitly enabled in optional_params ({"type": "enabled"}) + + This prevents spurious injection on models like deepseek-v3.2 that support + thinking as opt-in but not always-on. Addresses oss-pr-review-agent feedback + on PR #28057. + """ + from litellm.llms.deepseek.chat.transformation import DeepSeekChatConfig + + config = DeepSeekChatConfig() + + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + {"role": "user", "content": "Follow up"}, + ] + + # Case 1: reasoning model + thinking enabled -> injection should happen + result = config.transform_request( + model="deepseek-reasoner", + messages=messages, + optional_params={"thinking": {"type": "enabled"}}, + litellm_params={}, + headers={}, + ) + assert result["messages"][1].get("reasoning_content") == " ", ( + "reasoning_content should be injected when thinking is enabled" + ) + + # Case 2: reasoning model + thinking NOT in optional_params -> no injection + result = config.transform_request( + model="deepseek-reasoner", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "reasoning_content" not in result["messages"][1], ( + "reasoning_content should not be injected when thinking is not enabled" + ) + + # Case 3: non-reasoning model + thinking enabled -> no injection + result = config.transform_request( + model="deepseek-chat", + messages=messages, + optional_params={"thinking": {"type": "enabled"}}, + litellm_params={}, + headers={}, + ) + assert "reasoning_content" not in result["messages"][1], ( + "reasoning_content should not be injected for non-reasoning models" + ) diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_embedding_transformation.py b/tests/test_litellm/llms/dashscope/test_dashscope_embedding_transformation.py new file mode 100644 index 00000000000..5e4d0177e8d --- /dev/null +++ b/tests/test_litellm/llms/dashscope/test_dashscope_embedding_transformation.py @@ -0,0 +1,141 @@ +""" +Unit tests for DashScope embedding transformation. +""" + +import json +import os +import sys +from unittest.mock import MagicMock + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.dashscope.common_utils import DashScopeError +from litellm.llms.dashscope.embed.transformation import ( + DEFAULT_API_BASE, + DashScopeEmbeddingConfig, +) +from litellm.types.utils import EmbeddingResponse + + +def test_validate_environment_and_url(): + config = DashScopeEmbeddingConfig() + headers = config.validate_environment( + headers={}, + model="text-embedding-v4", + messages=[], + optional_params={}, + litellm_params={}, + api_key="sk-test", + ) + assert headers["Authorization"] == "Bearer sk-test" + + url = config.get_complete_url( + api_base=None, + api_key="sk-test", + model="text-embedding-v4", + optional_params={}, + litellm_params={}, + ) + assert url == f"{DEFAULT_API_BASE}/embeddings" + + +def test_transform_embedding_request(): + config = DashScopeEmbeddingConfig() + data = config.transform_embedding_request( + model="text-embedding-v4", + input=["风急天高猿啸哀"], + optional_params={"dimensions": 1024, "encoding_format": "float"}, + headers={}, + ) + assert data == { + "model": "text-embedding-v4", + "input": ["风急天高猿啸哀"], + "dimensions": 1024, + "encoding_format": "float", + } + + +def test_transform_embedding_response_success(): + config = DashScopeEmbeddingConfig() + payload = { + "data": [ + {"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}, + ], + "model": "text-embedding-v4", + "object": "list", + "usage": {"prompt_tokens": 5, "total_tokens": 5}, + "id": "73591b79-xxxx", + } + raw = httpx.Response( + status_code=200, + content=json.dumps(payload).encode("utf-8"), + request=httpx.Request("POST", "https://example.com"), + ) + result = config.transform_embedding_response( + model="text-embedding-v4", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="sk-x", + request_data={"input": ["a"]}, + optional_params={}, + litellm_params={}, + ) + assert result.model == "text-embedding-v4" + assert len(result.data) == 1 + assert result.usage.prompt_tokens == 5 + + +def test_transform_embedding_request_user_param(): + config = DashScopeEmbeddingConfig() + data = config.transform_embedding_request( + model="text-embedding-v4", + input=["hello"], + optional_params={"user": "user-123"}, + headers={}, + ) + assert data["user"] == "user-123" + + +def test_map_openai_params_drops_unsupported_with_drop_params(): + config = DashScopeEmbeddingConfig() + result = config.map_openai_params( + non_default_params={"dimensions": 512, "unknown_param": "value"}, + optional_params={}, + model="text-embedding-v4", + drop_params=True, + ) + assert result == {"dimensions": 512} + assert "unknown_param" not in result + + +def test_transform_embedding_response_error(): + config = DashScopeEmbeddingConfig() + payload = { + "error": { + "message": "Incorrect API key provided.", + "type": "invalid_request_error", + "code": "invalid_api_key", + } + } + raw = httpx.Response( + status_code=401, + content=json.dumps(payload).encode("utf-8"), + request=httpx.Request("POST", "https://example.com"), + ) + with pytest.raises(DashScopeError) as exc: + config.transform_embedding_response( + model="text-embedding-v4", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="sk-bad", + request_data={"input": ["a"]}, + optional_params={}, + litellm_params={}, + ) + assert exc.value.status_code == 401 + assert "Incorrect API key" in exc.value.message diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py b/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py new file mode 100644 index 00000000000..0e8d58b6530 --- /dev/null +++ b/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py @@ -0,0 +1,328 @@ +""" +Unit tests for DashScope rerank transformation. +""" + +import json +import os +import sys +from unittest.mock import MagicMock + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.dashscope.common_utils import DashScopeError +from litellm.llms.dashscope.rerank.transformation import ( + DEFAULT_RERANK_URL, + DashScopeRerankConfig, +) +from litellm.types.rerank import RerankResponse + + +class TestDashScopeRerankURL: + def setup_method(self): + self.config = DashScopeRerankConfig() + + def test_default_url(self): + url = self.config.get_complete_url(api_base=None, model="qwen3-rerank") + assert url == DEFAULT_RERANK_URL + + def test_explicit_v1_base_appends_reranks(self): + url = self.config.get_complete_url( + api_base="https://dashscope.aliyuncs.com/compatible-mode/v1", + model="qwen3-rerank", + ) + assert url == "https://dashscope.aliyuncs.com/compatible-mode/v1/reranks" + + def test_intl_v1_base_appends_reranks(self): + url = self.config.get_complete_url( + api_base="https://dashscope-intl.aliyuncs.com/compatible-mode/v1", + model="qwen3-rerank", + ) + assert url == "https://dashscope-intl.aliyuncs.com/compatible-mode/v1/reranks" + + def test_already_complete_url_passthrough(self): + full = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" + assert self.config.get_complete_url(api_base=full, model="qwen3-rerank") == full + + def test_trailing_slash_stripped(self): + full = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks/" + assert self.config.get_complete_url( + api_base=full, model="qwen3-rerank" + ) == full.rstrip("/") + + def test_custom_v1_base_appends_reranks(self): + url = self.config.get_complete_url( + api_base="https://my-proxy.example.com/v1", model="qwen3-rerank" + ) + assert url == "https://my-proxy.example.com/v1/reranks" + + +class TestDashScopeRerankRequest: + def setup_method(self): + self.config = DashScopeRerankConfig() + + def test_validate_environment_with_explicit_key(self): + headers = self.config.validate_environment( + headers={}, model="qwen3-rerank", api_key="sk-test" + ) + assert headers["Authorization"] == "Bearer sk-test" + assert headers["content-type"] == "application/json" + + def test_validate_environment_missing_key(self, monkeypatch): + monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False) + with pytest.raises(ValueError, match="DASHSCOPE_API_KEY"): + self.config.validate_environment( + headers={}, model="qwen3-rerank", api_key=None + ) + + def test_validate_environment_falls_back_to_env(self, monkeypatch): + monkeypatch.setenv("DASHSCOPE_API_KEY", "env-key") + headers = self.config.validate_environment( + headers={}, model="qwen3-rerank", api_key=None + ) + assert headers["Authorization"] == "Bearer env-key" + + def test_supported_params(self): + assert self.config.get_supported_cohere_rerank_params("qwen3-rerank") == [ + "query", + "documents", + "top_n", + "return_documents", + ] + + def test_map_params_drops_unsupported(self): + # qwen3-rerank accepts query/documents/top_n/return_documents. + # rank_fields and max_*_per_doc are silently dropped. + params = self.config.map_cohere_rerank_params( + non_default_params={}, + model="qwen3-rerank", + drop_params=False, + query="什么是文本排序模型", + documents=["d1", "d2"], + top_n=2, + rank_fields=["title"], + return_documents=True, + max_chunks_per_doc=5, + max_tokens_per_doc=100, + ) + assert params == { + "query": "什么是文本排序模型", + "documents": ["d1", "d2"], + "top_n": 2, + "return_documents": True, + } + + def test_transform_request_full(self): + body = self.config.transform_rerank_request( + model="qwen3-rerank", + optional_rerank_params={ + "query": "如何制作美味的苹果派?", + "documents": ["a", "b"], + "top_n": 5, + "return_documents": True, + }, + headers={}, + ) + assert body == { + "model": "qwen3-rerank", + "query": "如何制作美味的苹果派?", + "documents": ["a", "b"], + "top_n": 5, + "return_documents": True, + } + + def test_transform_request_omits_unset_optional(self): + body = self.config.transform_rerank_request( + model="qwen3-rerank", + optional_rerank_params={"query": "q", "documents": ["a"]}, + headers={}, + ) + assert "top_n" not in body + assert "return_documents" not in body + + def test_transform_request_requires_query(self): + with pytest.raises(ValueError, match="query"): + self.config.transform_rerank_request( + model="qwen3-rerank", + optional_rerank_params={"documents": ["a"]}, + headers={}, + ) + + def test_transform_request_requires_documents(self): + with pytest.raises(ValueError, match="documents"): + self.config.transform_rerank_request( + model="qwen3-rerank", + optional_rerank_params={"query": "q"}, + headers={}, + ) + + +class TestDashScopeRerankResponse: + def setup_method(self): + self.config = DashScopeRerankConfig() + self.logging = MagicMock() + + def _resp(self, body, status_code=200): + return httpx.Response( + status_code=status_code, + content=json.dumps(body).encode(), + request=httpx.Request("POST", "https://example.com"), + ) + + def test_success_response(self): + body = { + "object": "list", + "results": [ + {"index": 0, "relevance_score": 0.93}, + {"index": 2, "relevance_score": 0.34}, + ], + "model": "qwen3-rerank", + "id": "85ba5752", + "usage": {"total_tokens": 79}, + } + out = self.config.transform_rerank_response( + model="qwen3-rerank", + raw_response=self._resp(body), + model_response=RerankResponse(), + logging_obj=self.logging, + api_key="sk", + request_data={"query": "q"}, + ) + assert out.id == "85ba5752" + assert out.results == [ + {"index": 0, "relevance_score": 0.93}, + {"index": 2, "relevance_score": 0.34}, + ] + assert out.meta == { + "billed_units": {"total_tokens": 79}, + "tokens": {"input_tokens": 79}, + } + + def test_response_with_return_documents_real_payload(self): + # Verbatim sample from a real qwen3-rerank call with return_documents=true. + body = { + "object": "list", + "results": [ + { + "document": { + "text": "苹果派的制作步骤包括准备面团、切苹果、调制馅料、组装和烘烤。" + }, + "index": 1, + "relevance_score": 0.8304247466067356, + }, + { + "document": { + "text": "制作苹果派时，预先煮软苹果可以缩短烘烤时间。" + }, + "index": 3, + "relevance_score": 0.7142660211908354, + }, + ], + "model": "qwen3-rerank", + "id": "e191b077-97c4-9929-b121-c2fbd2c7b0af", + "usage": {"total_tokens": 192}, + } + out = self.config.transform_rerank_response( + model="qwen3-rerank", + raw_response=self._resp(body), + model_response=RerankResponse(), + logging_obj=self.logging, + request_data={"query": "如何制作美味的苹果派?"}, + ) + assert out.id == "e191b077-97c4-9929-b121-c2fbd2c7b0af" + assert out.results == [ + { + "index": 1, + "relevance_score": 0.8304247466067356, + "document": { + "text": "苹果派的制作步骤包括准备面团、切苹果、调制馅料、组装和烘烤。" + }, + }, + { + "index": 3, + "relevance_score": 0.7142660211908354, + "document": {"text": "制作苹果派时，预先煮软苹果可以缩短烘烤时间。"}, + }, + ] + assert out.meta == { + "billed_units": {"total_tokens": 192}, + "tokens": {"input_tokens": 192}, + } + + def test_response_string_document_normalized(self): + # Defensive path: if a future API revision returns a bare string, + # normalize to {"text": ...} so downstream code stays consistent. + body = { + "results": [{"index": 0, "relevance_score": 0.9, "document": "hello"}], + "model": "qwen3-rerank", + "usage": {"total_tokens": 5}, + } + out = self.config.transform_rerank_response( + model="qwen3-rerank", + raw_response=self._resp(body), + model_response=RerankResponse(), + logging_obj=self.logging, + ) + assert out.results[0]["document"] == {"text": "hello"} + + def test_missing_id_generates_uuid(self): + body = {"results": [{"index": 0, "relevance_score": 0.5}], "usage": {}} + out = self.config.transform_rerank_response( + model="qwen3-rerank", + raw_response=self._resp(body), + model_response=RerankResponse(), + logging_obj=self.logging, + ) + assert out.id is not None and len(out.id) > 0 + + def test_error_envelope_raises(self): + body = { + "code": "InvalidApiKey", + "message": "Invalid API-key provided.", + "request_id": "fb53", + } + with pytest.raises(DashScopeError) as exc_info: + self.config.transform_rerank_response( + model="qwen3-rerank", + raw_response=self._resp(body, status_code=401), + model_response=RerankResponse(), + logging_obj=self.logging, + ) + assert "Invalid API-key provided." in str(exc_info.value) + + def test_non_json_response_raises(self): + bad = httpx.Response( + status_code=500, + content=b"bad gateway", + request=httpx.Request("POST", "https://example.com"), + ) + with pytest.raises(DashScopeError): + self.config.transform_rerank_response( + model="qwen3-rerank", + raw_response=bad, + model_response=RerankResponse(), + logging_obj=self.logging, + ) + + def test_get_error_class(self): + err = self.config.get_error_class( + error_message="boom", status_code=500, headers={} + ) + assert isinstance(err, DashScopeError) + assert err.status_code == 500 + + +class TestProviderConfigManagerDispatch: + def test_dashscope_returns_rerank_config(self): + import litellm + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_rerank_config( + model="qwen3-rerank", + provider=litellm.LlmProviders.DASHSCOPE, + api_base=None, + present_version_params=[], + ) + assert isinstance(cfg, DashScopeRerankConfig) diff --git a/tests/test_litellm/llms/test_file_content_block.py b/tests/test_litellm/llms/test_file_content_block.py new file mode 100644 index 00000000000..5552c1a4d68 --- /dev/null +++ b/tests/test_litellm/llms/test_file_content_block.py @@ -0,0 +1,433 @@ +""" +Tests for handling malformed or invalid 'file' content blocks (missing or null +`file` sub-field, HTTP file_id URLs for Google AI Studio). + +Regression tests for: +- litellm/llms/vertex_ai/gemini/transformation.py +- litellm/llms/gemini/chat/transformation.py +- litellm/litellm_core_utils/prompt_templates/common_utils.py + (migrate_file_to_image_url raises on missing `file`; file-id helpers skip non-OpenAI shapes) +- litellm/litellm_core_utils/prompt_templates/factory.py (Bedrock + Anthropic) +- litellm/llms/openai/chat/gpt_transformation.py +""" + +import asyncio +import copy +from typing import List, cast + +import pytest + +import litellm +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + get_file_ids_from_messages, + migrate_file_to_image_url, + update_messages_with_model_file_ids, +) +from litellm.litellm_core_utils.prompt_templates.factory import ( + BedrockConverseMessagesProcessor, + anthropic_process_openai_file_message, +) +from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, +) +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionFileObject, + OpenAIMessageContentListBlock, +) + +_MALFORMED_MESSAGES_RAW = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "file"}, # Missing required "file" sub-field + ], + } +] + +_WELL_FORMED_MESSAGES_RAW = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + { + "type": "file", + "file": {"file_id": "file-abc123", "format": "pdf"}, + }, + ], + } +] + +MALFORMED_FILE_OBJECT: ChatCompletionFileObject = cast( + ChatCompletionFileObject, {"type": "file"} +) + +EXPLICIT_NULL_FILE_OBJECT: ChatCompletionFileObject = cast( + ChatCompletionFileObject, + {"type": "file", "file": None}, +) + + +def _malformed() -> List[AllMessageValues]: + return copy.deepcopy(cast(List[AllMessageValues], _MALFORMED_MESSAGES_RAW)) + + +def _well_formed() -> List[AllMessageValues]: + return copy.deepcopy(cast(List[AllMessageValues], _WELL_FORMED_MESSAGES_RAW)) + + +def _explicit_null_file_in_content() -> List[AllMessageValues]: + return copy.deepcopy( + cast( + List[AllMessageValues], + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "file", "file": None}, + ], + } + ], + ) + ) + + +# --------------------------------------------------------------------------- +# vertex_ai/gemini/transformation.py +# --------------------------------------------------------------------------- + + +def test_gemini_convert_messages_malformed_file_raises_bad_request(): + """_gemini_convert_messages_with_history should raise BadRequestError (not KeyError) + when a content block has type='file' but no 'file' sub-field.""" + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + _gemini_convert_messages_with_history( + messages=_malformed(), + model="gemini-2.0-flash", + ) + + +def test_gemini_convert_messages_explicit_null_file_field_raises_bad_request(): + """Explicit JSON null for `file` must be rejected like a missing `file` key.""" + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + _gemini_convert_messages_with_history( + messages=_explicit_null_file_in_content(), + model="gemini-2.0-flash", + ) + + +# --------------------------------------------------------------------------- +# gemini/chat/transformation.py - GoogleAIStudioGeminiConfig +# --------------------------------------------------------------------------- + + +def test_google_ai_studio_transform_messages_malformed_file_raises_bad_request(): + """GoogleAIStudioGeminiConfig._transform_messages should raise BadRequestError + when a content block has type='file' but no 'file' sub-field.""" + config = GoogleAIStudioGeminiConfig() + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + config._transform_messages(messages=_malformed(), model="gemini-2.0-flash") + + +def test_google_ai_studio_transform_messages_explicit_null_file_field_raises_bad_request(): + """Explicit JSON null for `file` must be rejected like a missing `file` key.""" + config = GoogleAIStudioGeminiConfig() + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + config._transform_messages( + messages=_explicit_null_file_in_content(), model="gemini-2.0-flash" + ) + + +def test_google_ai_studio_transform_messages_http_file_id_converts_to_base64(monkeypatch): + """Google AI Studio rejects raw HTTP(S) file URLs; _transform_messages should + fetch and replace them with base64 `file_data` before conversion.""" + # Data URL shape so downstream Gemini media parsing accepts the inlined bytes + # (mirrors real `convert_url_to_base64` output from `_process_image_response`). + fake_file_data = "data:application/pdf;base64,aGVsbG8=" + + def _fake_convert_url_to_base64(url: str) -> str: + assert url == "https://example.com/doc.pdf" + return fake_file_data + + monkeypatch.setattr( + "litellm.llms.gemini.chat.transformation.convert_url_to_base64", + _fake_convert_url_to_base64, + ) + messages = cast( + List[AllMessageValues], + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/doc.pdf", + "format": "pdf", + }, + }, + ], + } + ], + ) + config = GoogleAIStudioGeminiConfig() + config._transform_messages(messages=messages, model="gemini-2.0-flash") + content = messages[0].get("content") + assert isinstance(content, list) + file_block = next(c for c in content if isinstance(c, dict) and c.get("type") == "file") + file_field = file_block.get("file") + assert isinstance(file_field, dict) + assert file_field.get("file_data") == fake_file_data + assert "file_id" not in file_field + + +def test_google_ai_studio_transform_messages_http_file_id_convert_failure_leaves_file_unchanged( + monkeypatch, +): + """If convert_url_to_base64 fails, the Studio prep step must not mutate the block + (see try/except in GoogleAIStudioGeminiConfig._transform_messages).""" + https_id = "https://example.com/missing.pdf" + + def _raise(_url: str) -> str: + raise litellm.ImageFetchError("simulated fetch failure") + + monkeypatch.setattr( + "litellm.llms.gemini.chat.transformation.convert_url_to_base64", + _raise, + ) + messages = cast( + List[AllMessageValues], + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + { + "type": "file", + "file": { + "file_id": https_id, + "format": "application/pdf", + }, + }, + ], + } + ], + ) + config = GoogleAIStudioGeminiConfig() + config._transform_messages(messages=messages, model="gemini-2.0-flash") + content = messages[0].get("content") + assert isinstance(content, list) + file_block = next(c for c in content if isinstance(c, dict) and c.get("type") == "file") + file_field = file_block.get("file") + assert isinstance(file_field, dict) + assert file_field.get("file_id") == https_id + assert file_field.get("format") == "application/pdf" + assert "file_data" not in file_field + + +# --------------------------------------------------------------------------- +# common_utils.py - update_messages_with_model_file_ids +# --------------------------------------------------------------------------- + + +def test_update_messages_with_model_file_ids_malformed_skips_non_openai_file_block(): + """Non-OpenAI file blocks (e.g. missing nested `file` dict) are skipped so callers + relying on LangChain v1 / provider-native shapes are not rejected here.""" + messages = _malformed() + result = update_messages_with_model_file_ids( + messages=messages, + model_id="some-model", + model_file_id_mapping={}, + ) + assert result == messages + content = result[0].get("content") + assert isinstance(content, list) + file_block = next(c for c in content if isinstance(c, dict) and c.get("type") == "file") + assert "file" not in file_block + + +def test_update_messages_with_model_file_ids_well_formed_updates(): + """update_messages_with_model_file_ids should update file_id for well-formed blocks.""" + mapping = {"file-abc123": {"some-model": "provider-file-xyz"}} + result = update_messages_with_model_file_ids( + messages=_well_formed(), + model_id="some-model", + model_file_id_mapping=mapping, + ) + content = result[0].get("content") + assert isinstance(content, list) + file_block = next(c for c in content if c.get("type") == "file") + assert file_block.get("file", {}).get("file_id") == "provider-file-xyz" + + +# --------------------------------------------------------------------------- +# common_utils.py - get_file_ids_from_messages +# --------------------------------------------------------------------------- + + +def test_get_file_ids_from_messages_malformed_skips_non_openai_file_block(): + """Blocks with type='file' but no OpenAI `file` sub-dict yield no extracted ids.""" + assert get_file_ids_from_messages(messages=_malformed()) == [] + + +def test_get_file_ids_from_messages_well_formed_returns_ids(): + """get_file_ids_from_messages should extract file_id from well-formed blocks.""" + messages: List[AllMessageValues] = cast( + List[AllMessageValues], + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "file", "file": {"file_id": "file-abc123", "format": "pdf"}}, + ], + } + ], + ) + result = get_file_ids_from_messages(messages=messages) + assert result == ["file-abc123"] + + +# --------------------------------------------------------------------------- +# factory.py - BedrockConverseMessagesProcessor (sync + async) +# --------------------------------------------------------------------------- + + +def test_bedrock_process_file_message_malformed_raises_bad_request(): + """_process_file_message should raise BadRequestError (not KeyError) + when the file object is missing the 'file' sub-field.""" + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + BedrockConverseMessagesProcessor._process_file_message(MALFORMED_FILE_OBJECT) + + +def test_bedrock_process_file_message_explicit_null_file_field_raises_bad_request(): + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + BedrockConverseMessagesProcessor._process_file_message(EXPLICIT_NULL_FILE_OBJECT) + + +def test_bedrock_async_process_file_message_malformed_raises_bad_request(): + """_async_process_file_message should raise BadRequestError (not KeyError) + when the file object is missing the 'file' sub-field.""" + + async def _run() -> None: + with pytest.raises( + litellm.BadRequestError, match="missing the required 'file' field" + ): + await BedrockConverseMessagesProcessor._async_process_file_message( + MALFORMED_FILE_OBJECT + ) + + asyncio.run(_run()) + + +def test_bedrock_async_process_file_message_explicit_null_file_field_raises_bad_request(): + async def _run() -> None: + with pytest.raises( + litellm.BadRequestError, match="missing the required 'file' field" + ): + await BedrockConverseMessagesProcessor._async_process_file_message( + EXPLICIT_NULL_FILE_OBJECT + ) + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# openai/chat/gpt_transformation.py +# --------------------------------------------------------------------------- + + +def test_openai_apply_common_transform_malformed_file_raises_bad_request(): + """_apply_common_transform_content_item should raise BadRequestError (not KeyError) + when a content block has type='file' but no 'file' sub-field.""" + config = OpenAIGPTConfig() + malformed_block: OpenAIMessageContentListBlock = cast( + OpenAIMessageContentListBlock, {"type": "file"} + ) + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + config._apply_common_transform_content_item(malformed_block) + + +def test_openai_apply_common_transform_explicit_null_file_field_raises_bad_request(): + config = OpenAIGPTConfig() + explicit_null_block: OpenAIMessageContentListBlock = cast( + OpenAIMessageContentListBlock, + {"type": "file", "file": None}, + ) + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + config._apply_common_transform_content_item(explicit_null_block) + + +def test_openai_apply_common_transform_well_formed_file_does_not_raise(): + """_apply_common_transform_content_item should not raise for well-formed file blocks.""" + config = OpenAIGPTConfig() + well_formed_block: OpenAIMessageContentListBlock = cast( + OpenAIMessageContentListBlock, + {"type": "file", "file": {"file_id": "file-abc123"}}, + ) + result = config._apply_common_transform_content_item(well_formed_block) + assert result.get("type") == "file" + file_field = cast(ChatCompletionFileObject, result).get("file", {}) + assert file_field.get("file_id") == "file-abc123" + + +# --------------------------------------------------------------------------- +# factory.py - anthropic_process_openai_file_message +# --------------------------------------------------------------------------- + + +def test_anthropic_process_openai_file_message_malformed_raises_bad_request(): + """anthropic_process_openai_file_message should raise BadRequestError (not KeyError) + when the file object is missing the 'file' sub-field.""" + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + anthropic_process_openai_file_message(MALFORMED_FILE_OBJECT) + + +def test_anthropic_process_openai_file_message_explicit_null_file_field_raises_bad_request(): + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + anthropic_process_openai_file_message(EXPLICIT_NULL_FILE_OBJECT) + + +def test_anthropic_process_openai_file_message_well_formed_file_id_does_not_raise(): + """anthropic_process_openai_file_message should not raise for a well-formed file_id block.""" + well_formed: ChatCompletionFileObject = cast( + ChatCompletionFileObject, + {"type": "file", "file": {"file_id": "file-abc123"}}, + ) + result = anthropic_process_openai_file_message(well_formed) + assert result.get("type") in ("document", "image", "container_upload") + + +# --------------------------------------------------------------------------- +# common_utils.py - migrate_file_to_image_url +# --------------------------------------------------------------------------- + + +def test_migrate_file_to_image_url_malformed_raises_bad_request(): + """migrate_file_to_image_url should raise BadRequestError (not KeyError) + when the file object is missing the 'file' sub-field.""" + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + migrate_file_to_image_url(MALFORMED_FILE_OBJECT) + + +def test_migrate_file_to_image_url_explicit_null_file_field_raises_bad_request(): + with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"): + migrate_file_to_image_url(EXPLICIT_NULL_FILE_OBJECT) + + +def test_migrate_file_to_image_url_well_formed_returns_image_url(): + """migrate_file_to_image_url should return an image_url block for a well-formed file.""" + well_formed: ChatCompletionFileObject = cast( + ChatCompletionFileObject, + {"type": "file", "file": {"file_id": "file-abc123", "format": "png"}}, + ) + result = migrate_file_to_image_url(well_formed) + assert result.get("type") == "image_url" + image_url = result.get("image_url", {}) + assert isinstance(image_url, dict) + assert image_url.get("url") == "file-abc123" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py b/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py new file mode 100644 index 00000000000..10fc68ecaad --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py @@ -0,0 +1,52 @@ +import pytest +from typing import List, cast + +import litellm +from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, +) +from litellm.types.llms.openai import AllMessageValues + + +def test_missing_image_url_field_raises_bad_request_error(): + """When element type is 'image_url' but 'image_url' field is missing, a BadRequestError is raised.""" + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": [{"type": "image_url"}]}], + ) + with pytest.raises(litellm.BadRequestError) as exc_info: + _gemini_convert_messages_with_history(messages, model="gemini-1.5-pro") + assert "'image_url' field is missing" in str(exc_info.value) + + +def test_missing_url_inside_image_url_dict_raises_bad_request_error(): + """When image_url is a dict but 'url' key is absent, a BadRequestError is raised.""" + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": [{"type": "image_url", "image_url": {"detail": "high"}}]}], + ) + with pytest.raises(litellm.BadRequestError) as exc_info: + _gemini_convert_messages_with_history(messages, model="gemini-1.5-pro") + assert "'url' field is missing inside" in str(exc_info.value) + + +def test_explicit_null_image_url_raises_bad_request_error(): + """When image_url key is present but explicitly null, a BadRequestError is raised.""" + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": [{"type": "image_url", "image_url": None}]}], + ) + with pytest.raises(litellm.BadRequestError) as exc_info: + _gemini_convert_messages_with_history(messages, model="gemini-1.5-pro") + assert "'image_url' field is missing" in str(exc_info.value) + + +def test_empty_dict_image_url_raises_bad_request_error(): + """When image_url is an empty dict (no url), a BadRequestError is raised.""" + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": [{"type": "image_url", "image_url": {}}]}], + ) + with pytest.raises(litellm.BadRequestError) as exc_info: + _gemini_convert_messages_with_history(messages, model="gemini-1.5-pro") + assert "'url' field is missing inside" in str(exc_info.value) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 353d19b0198..1e0ad04c3c2 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -4290,3 +4290,444 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): msg = str(exc_info.value) assert "secret content" not in msg assert "Error converting to valid response block" in msg + + +def test_chunk_parser_raises_on_429_error_chunk(): + """Test chunk_parser raises VertexAIError on 429 RESOURCE_EXHAUSTED error chunk""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": 429, + "message": "Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/error-code-429 for more details.", + "status": "RESOURCE_EXHAUSTED", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 429 + assert "RESOURCE_EXHAUSTED" in exc_info.value.message + assert "Resource exhausted" in exc_info.value.message + + +def test_chunk_parser_raises_on_500_error_chunk(): + """Test chunk_parser raises VertexAIError on 500 INTERNAL error chunk""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": 500, + "message": "Internal error encountered.", + "status": "INTERNAL", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "INTERNAL" in exc_info.value.message + + +def test_chunk_parser_raises_on_error_chunk_with_minimal_fields(): + """Test chunk_parser handles error chunks with missing optional fields""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": 429, + "message": "Resource exhausted.", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 429 + + +def test_chunk_parser_normal_chunk_unaffected_by_error_check(): + """Test that normal streaming chunks still work correctly after error check addition""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + normal_chunk = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Hello"}], + }, + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 1, + "totalTokenCount": 6, + }, + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + result = streaming_obj.chunk_parser(normal_chunk) + assert result is not None + assert len(result.choices) > 0 + assert result.choices[0].delta.content == "Hello" + + +def test_chunk_parser_raises_on_non_dict_error(): + """Test chunk_parser raises VertexAIError when chunk['error'] is not a dict""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = {"error": "something went wrong"} + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Unexpected error format" in exc_info.value.message + + +def test_chunk_parser_raises_on_string_error_code(): + """Test chunk_parser correctly converts string error code to int""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # code field is a string "429" rather than an int + error_chunk = { + "error": { + "code": "429", + "message": "Resource exhausted.", + "status": "RESOURCE_EXHAUSTED", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 429 + assert isinstance(exc_info.value.status_code, int) + + +def test_chunk_parser_error_chunk_explicit_null_code_uses_500(): + """JSON null for code must not call int(None); status defaults to 500.""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": None, + "message": "Something went wrong.", + "status": "UNKNOWN", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Something went wrong" in exc_info.value.message + + +def test_chunk_parser_error_chunk_non_numeric_code_defaults_to_500(): + """Non-numeric code must not become ValueError -> RuntimeError in __next__.""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": "NOT_A_NUMBER", + "message": "Malformed.", + "status": "INVALID", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Malformed" in exc_info.value.message + + +def test_chunk_parser_error_chunk_empty_dict_defaults_to_500(): + """Empty error object {} uses default code 500 and default message/status strings.""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = {"error": {}} + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "UNKNOWN" in exc_info.value.message + assert "Unknown error" in exc_info.value.message + + +def test_chunk_parser_error_chunk_non_dict_int_value(): + """Non-dict error payloads (e.g. bare JSON number) must raise with status 500, not TypeError.""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = {"error": 503} + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Unexpected error format" in exc_info.value.message + assert "503" in exc_info.value.message + + +def test_chunk_parser_error_chunk_non_dict_null_value(): + """JSON null for error must hit the non-dict branch (same as int/string).""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = {"error": None} + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Unexpected error format" in exc_info.value.message + + +def test_mid_stream_429_error_raises_during_iteration(): + """ + Simulate a full streaming scenario: normal thinking chunks arrive first, + then a 429 RESOURCE_EXHAUSTED error chunk arrives mid-stream. + Verify that ModelResponseIterator raises VertexAIError during iteration. + """ + import json + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # Simulate Vertex AI SSE stream: normal chunks followed by a 429 error chunk + normal_chunk_1 = json.dumps( + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Let me think about this...", "thought": True}], + }, + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + "modelVersion": "gemini-3.1-flash-image-preview", + } + ) + + normal_chunk_2 = json.dumps( + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "I'll generate the image now.", "thought": True}], + }, + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 12, + "totalTokenCount": 22, + }, + } + ) + + error_chunk = json.dumps( + { + "error": { + "code": 429, + "message": "Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/error-code-429 for more details.", + "status": "RESOURCE_EXHAUSTED", + } + } + ) + + # Build a mock SSE stream (lines returned by iter_lines) + sse_lines = iter([normal_chunk_1, normal_chunk_2, error_chunk]) + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=sse_lines, + sync_stream=True, + logging_obj=logging_obj, + ) + + # Iterate the stream: first chunks should succeed, then 429 error should be raised + results = [] + with pytest.raises(VertexAIError) as exc_info: + for chunk in streaming_obj: + if chunk is not None: + results.append(chunk) + + # Verify: received normal chunks before the error + assert ( + len(results) >= 1 + ), "Should have received at least 1 normal chunk before the error" + + # Verify: 429 error is properly raised + assert exc_info.value.status_code == 429 + assert "RESOURCE_EXHAUSTED" in str(exc_info.value.message) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 2e9629f95de..be0e59e8b7d 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -1219,6 +1219,32 @@ def test_process_gemini_media(): mime_type="image/jpeg", file_uri="gs://bucket/image" ) + # Test gs url without extension using mime_type from image_url object + image_message = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "gs://bucket/image-without-extension", + "mime_type": "image/png", + }, + } + ], + } + ] + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + converted = _gemini_convert_messages_with_history( + messages=image_message, model="gemini-2.5-flash" + ) + assert converted[0]["parts"][0]["file_data"] == FileDataType( + mime_type="image/png", file_uri="gs://bucket/image-without-extension" + ) + # Test HTTPS JPG URL https_result = _process_gemini_media("https://example.com/image.jpg") print("https_result JPG", https_result) @@ -1256,6 +1282,7 @@ def test_process_gemini_media(): assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..." + def test_get_image_mime_type_from_url(): """Test the _get_image_mime_type_from_url function for different image URLs""" from litellm.llms.vertex_ai.gemini.transformation import ( diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py b/tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py new file mode 100644 index 00000000000..e0eccad80e2 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py @@ -0,0 +1,466 @@ +"""Vertex Gemini: extensionless gs:// MIME + GCS metadata tests. + +Split from test_vertex.py to satisfy CI per-file size limits. +""" +import asyncio +import os +import sys +import time + +from dotenv import load_dotenv + +load_dotenv() + +import pytest + +import litellm +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + +def test_process_gemini_media_gcs_explicit_format_octet_stream_and_alias(): + """Explicit format bypasses registry; image/jpg alias still applies.""" + from litellm.types.llms.vertex_ai import FileDataType + + r1 = _process_gemini_media( + "gs://bucket/object-no-ext", + format="application/octet-stream", + ) + assert r1["file_data"] == FileDataType( + mime_type="application/octet-stream", + file_uri="gs://bucket/object-no-ext", + ) + r2 = _process_gemini_media("gs://bucket/object-no-ext", format="image/jpg") + assert r2["file_data"] == FileDataType( + mime_type="image/jpeg", + file_uri="gs://bucket/object-no-ext", + ) + + +def test_process_gemini_media_gcs_without_extension_errors_and_metadata_mock(): + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value=None, + ): + with pytest.raises(litellm.BadRequestError) as exc: + _process_gemini_media("gs://bucket/image-without-extension") + assert "Unable to determine mime type for gs URI" in str(exc.value) + + from litellm.types.llms.vertex_ai import FileDataType + + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/jpeg", + ) as m: + r = _process_gemini_media("gs://bucket/image-without-extension") + assert r["file_data"] == FileDataType( + mime_type="image/jpeg", file_uri="gs://bucket/image-without-extension" + ) + m.assert_called() + + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/jpg", + ): + r_alias = _process_gemini_media("gs://bucket/image-without-extension") + assert r_alias["file_data"]["mime_type"] == "image/jpeg" + + +def test_process_gemini_media_rejects_gcs_metadata_mime_not_supported_by_gemini(): + """Non-empty GCS contentType that fails _normalize_and_validate_gemini_mime_type.""" + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="application/x-litellm-unit-test-unknown-mime", + ): + with pytest.raises( + litellm.BadRequestError, + match="File type not supported by gemini", + ): + _process_gemini_media("gs://bucket/object-without-extension") + + +def test_file_block_uses_mime_type_alias_for_extensionless_gcs(): + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + from litellm.types.llms.vertex_ai import FileDataType + + messages = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_id": "gs://bucket/no-extension-object", + "mime_type": "application/pdf", + }, + } + ], + } + ] + converted = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-flash" + ) + assert converted[0]["parts"][0]["file_data"] == FileDataType( + mime_type="application/pdf", file_uri="gs://bucket/no-extension-object" + ) + + +@pytest.mark.parametrize( + "bucket,expected", + [ + (("a." * 110) + "aa", True), + ("ab", False), + ("a" * 64, False), + ("ab..cd", False), + ("1.2.3.4", False), + ("192.168.0.1", False), + ("Bucket-Upper", False), + ("bucket@name", False), + ("bucket name", False), + ("-mybucket", False), + ("mybucket-", False), + (".mybucket", False), + ("mybucket.", False), + ], +) +def test_is_valid_gcs_bucket_name_matrix(bucket, expected): + from litellm.llms.vertex_ai.gemini.transformation import _is_valid_gcs_bucket_name + + assert _is_valid_gcs_bucket_name(bucket) is expected + + +def test_get_gcs_object_content_type_explicit_vertex_success_and_token_failure(): + from litellm.llms.vertex_ai.gemini import transformation as gt + + mock_v = MagicMock() + mock_v.get_access_token.return_value = ("test-token", "test-project") + resp = MagicMock() + resp.is_error = False + resp.status_code = 200 + resp.json.return_value = {"contentType": "image/png"} + http = MagicMock() + http.get.return_value = resp + + with ( + patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=http, + ), + ): + assert ( + gt._get_gcs_object_content_type( + image_url="gs://my-bucket/path/to/image-without-extension", + vertex_project="project-123", + vertex_credentials="credential-json", + ) + == "image/png" + ) + mock_v.get_access_token.assert_called_once_with( + credentials="credential-json", + project_id="project-123", + ) + + mock_v2 = MagicMock() + mock_v2.get_access_token.side_effect = Exception("token failure") + with patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v2): + with pytest.raises( + litellm.BadRequestError, + match="Unable to fetch GCS metadata with provided Vertex credentials/project", + ): + gt._get_gcs_object_content_type( + image_url="gs://my-bucket/path/to/image-without-extension", + vertex_project="project-123", + vertex_credentials="credential-json", + ) + + +def test_get_gcs_object_content_type_http_error_explicit_vs_anonymous(): + from litellm.llms.vertex_ai.gemini import transformation as gt + + mock_v = MagicMock() + mock_v.get_access_token.return_value = ("t", "p") + err_resp = MagicMock() + err_resp.is_error = True + err_resp.status_code = 403 + err_resp.text = '{"error":{"message":"Permission denied"}}' + http = MagicMock() + http.get.return_value = err_resp + + with ( + patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=http, + ), + ): + with pytest.raises(litellm.BadRequestError, match="HTTP 403") as ei: + gt._get_gcs_object_content_type( + image_url="gs://my-bucket/path/to/obj", + vertex_project="project-123", + vertex_credentials="credential-json", + ) + assert "Permission denied" in str(ei.value) + + mock_v2 = MagicMock() + anon_err = MagicMock() + anon_err.is_error = True + anon_err.status_code = 403 + anon_err.text = "Forbidden" + http2 = MagicMock() + http2.get.return_value = anon_err + with ( + patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v2), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=http2, + ), + ): + assert ( + gt._get_gcs_object_content_type(image_url="gs://public-bucket/public-object") + is None + ) + mock_v2.get_access_token.assert_not_called() + + +def test_get_gcs_object_content_type_anonymous_success_no_auth_header(): + from litellm.llms.vertex_ai.gemini import transformation as gt + + mock_v = MagicMock() + ok = MagicMock() + ok.is_error = False + ok.status_code = 200 + ok.json.return_value = {"contentType": "image/jpeg"} + http = MagicMock() + http.get.return_value = ok + + with ( + patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=http, + ), + ): + assert ( + gt._get_gcs_object_content_type(image_url="gs://public-bucket/public-object") + == "image/jpeg" + ) + mock_v.get_access_token.assert_not_called() + hdrs = http.get.call_args.kwargs.get("headers") + assert hdrs is None or "Authorization" not in hdrs + + +def test_async_transform_request_body_offloads_extensionless_gs_not_plain_text(): + from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation + + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "gs://bucket/image-without-extension"}, + } + ], + } + ] + + def slow_http_get(*args, **kwargs): + time.sleep(0.5) + response = MagicMock() + response.is_error = False + response.status_code = 200 + response.raise_for_status.return_value = None + response.json.return_value = {"contentType": "image/png"} + return response + + async def fake_check_and_create_cache(self, **kwargs): + return kwargs["messages"], kwargs["optional_params"], None + + mock_v = MagicMock() + mock_v.get_access_token.return_value = ("token", "project") + mock_http = MagicMock() + mock_http.get.side_effect = slow_http_get + + async def run_scenario() -> float: + async def concurrent_sleep() -> float: + start = time.monotonic() + await asyncio.sleep(0.05) + return time.monotonic() - start + + task = asyncio.create_task( + gemini_transformation.async_transform_request_body( + gemini_api_key=None, + messages=messages, + api_base=None, + model="gemini-2.5-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + ) + elapsed = await concurrent_sleep() + await task + return elapsed + + with ( + patch.object(gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_v), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=mock_http, + ), + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching." + "ContextCachingEndpoints.async_check_and_create_cache", + new=fake_check_and_create_cache, + ), + ): + sleep_elapsed = asyncio.run(run_scenario()) + + assert sleep_elapsed < 0.4, ( + f"Event loop blocked for {sleep_elapsed:.3f}s; " + "async_transform_request_body did not offload sync GCS metadata" + ) + + async def fake_cache2(self, **kwargs): + return kwargs["messages"], kwargs["optional_params"], None + + async def run_plain(): + with patch( + "litellm.llms.vertex_ai.gemini.transformation.asyncify", + side_effect=AssertionError("asyncify must not run without extensionless gs://"), + ): + return await gemini_transformation.async_transform_request_body( + gemini_api_key=None, + messages=[{"role": "user", "content": "hello"}], + api_base=None, + model="gemini-2.5-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching." + "ContextCachingEndpoints.async_check_and_create_cache", + new=fake_cache2, + ): + body = asyncio.run(run_plain()) + assert body is not None and "contents" in body + + +@pytest.mark.parametrize( + "messages,expected", + [ + ([{"role": "user", "content": "hello"}], False), + ( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "gs://bucket/image-without-extension"}, + } + ], + } + ], + True, + ), + ( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "gs://bucket/image.png"}, + } + ], + } + ], + False, + ), + ( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "gs://bucket/image-without-extension", + "mime_type": "image/png", + }, + } + ], + } + ], + False, + ), + ( + [ + { + "role": "assistant", + "content": [], + "images": [ + {"image_url": {"url": "gs://bucket/gen-without-extension"}}, + ], + } + ], + True, + ), + ( + [ + { + "role": "assistant", + "content": [], + "images": [{"image_url": {"url": "gs://bucket/gen.png"}}], + } + ], + False, + ), + ( + [ + { + "role": "assistant", + "content": [], + "images": [ + { + "image_url": { + "url": "gs://bucket/gen-no-ext", + "mime_type": "image/png", + }, + } + ], + } + ], + False, + ), + ], +) +def test_openai_messages_may_need_sync_gcs_metadata_fetch_matrix(messages, expected): + from litellm.llms.vertex_ai.gemini.transformation import ( + _openai_messages_may_need_sync_gcs_metadata_fetch, + ) + + assert _openai_messages_may_need_sync_gcs_metadata_fetch(messages) is expected diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py new file mode 100644 index 00000000000..b483a75a939 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py @@ -0,0 +1,121 @@ +""" +Regression tests for #28084: + +`VertexAIPartnerModels.count_tokens` (for Claude / Mistral / Llama on Vertex) +used to gate on `import vertexai` even though the actual count-tokens path goes +through `VertexAIPartnerModelsTokenCounter.handle_count_tokens_request`, which +talks to the publisher's `:rawPredict` endpoint over plain httpx and never +touches the Gemini SDK. The unused gate broke `/v1/messages/count_tokens` for +any LiteLLM install that did not pull in `google-cloud-aiplatform` (which is +not in the default `proxy` / `proxy-dev` extras). + +These tests pin the absence of that gate by: + +1. simulating `vertexai` being unimportable and verifying the partner-model + path does not raise the historical "vertexai import failed" error before + reaching the network/auth layer, and +2. asserting that import of the partner-model count-tokens handler module by + itself does not pull `vertexai` into `sys.modules`. +""" + +import sys + +import pytest + +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import ( + VertexAIPartnerModelsTokenCounter, +) +from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels + + +@pytest.mark.asyncio +async def test_count_tokens_does_not_require_vertexai_sdk(monkeypatch): + """Even when `import vertexai` would fail, count_tokens must not raise the + historical "vertexai import failed" gate. The downstream handler talks to + `:rawPredict` over httpx with an access token — no Gemini SDK needed.""" + + # Simulate `vertexai` being unimportable, regardless of what is actually on + # the test environment's sys.path. + monkeypatch.setitem(sys.modules, "vertexai", None) + monkeypatch.setitem(sys.modules, "vertexai.preview", None) + + captured = {} + + async def fake_ensure_access_token( + self, credentials, project_id, custom_llm_provider + ): + return "fake-token", "fake-project" + + def fake_build_endpoint(self, model, project_id, vertex_location, api_base=None): + captured["model_to_endpoint"] = model + return "https://fake-endpoint" + + monkeypatch.setattr( + VertexAIPartnerModelsTokenCounter, + "_ensure_access_token_async", + fake_ensure_access_token, + ) + monkeypatch.setattr( + VertexAIPartnerModelsTokenCounter, + "_build_count_tokens_endpoint", + fake_build_endpoint, + ) + + class FakeResponse: + status_code = 200 + + def json(self): + return {"input_tokens": 9} + + class FakeClient: + async def post(self, url, headers=None, json=None, **kwargs): + captured["url"] = url + captured["headers"] = headers + captured["json"] = json + return FakeResponse() + + import litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler as handler_mod + + monkeypatch.setattr( + handler_mod, "get_async_httpx_client", lambda **kwargs: FakeClient() + ) + + result = await VertexAIPartnerModels().count_tokens( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + litellm_params={"vertex_location": "us-east5"}, + vertex_project="test-project", + vertex_location="us-east5", + vertex_credentials=None, + ) + + # We should reach the publisher endpoint and parse its response, not raise + # the vertexai-import gate. + assert result == { + "input_tokens": 9, + "tokenizer_used": "vertex_ai_partner_models", + } + assert captured["headers"] == {"Authorization": "Bearer fake-token"} + assert captured["model_to_endpoint"] == "claude-sonnet-4-6" + + +def test_handler_module_does_not_import_vertexai_sdk(): + """Importing the partner-model count-tokens handler must not load the + Gemini SDK into sys.modules. Operators who only need Claude-on-Vertex + token counting should not pay for `google-cloud-aiplatform`.""" + + # Force-evict any prior load so this assertion measures what THIS module + # pulls in, not what an unrelated earlier test did. + for mod in list(sys.modules): + if mod == "vertexai" or mod.startswith("vertexai."): + sys.modules.pop(mod, None) + + # Re-import the handler module to verify it stays SDK-free. + import importlib + + import litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler as handler_mod + + importlib.reload(handler_mod) + + leaked = [m for m in sys.modules if m == "vertexai" or m.startswith("vertexai.")] + assert leaked == [], f"unexpected vertexai SDK imports: {leaked}" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index e01038cd35f..6ec793a1bb0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -553,6 +553,7 @@ class TestGuardrailActions: # Verify the exception has the clean error message (no wrapper) assert str(exc_info.value) == "Content contains harmful instructions" assert exc_info.value.guardrail_name == "generic_guardrail_api" + assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_action_intervened_modifies_content( diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index f0bf4578636..b65f6305b77 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1446,3 +1446,175 @@ class TestGetTeamDeployments: result = await _get_team_deployments(team_id, prisma_client) assert len(result) == 1 assert result[0] is dep1 + + +def _build_db_model_for_blocked_test(): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + return Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="dep-0"), + ) + + +class TestUpdateDBModelBlocked: + """`update_db_model` must thread `blocked` through to the Prisma payload only + when the caller explicitly set it — PATCH semantics: an absent field means + "leave the stored value untouched".""" + + def test_update_db_model_passes_blocked_true_to_db(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_for_blocked_test(), + updated_patch=updateDeployment(blocked=True), + ) + assert result["blocked"] is True + + def test_update_db_model_passes_blocked_false_to_db(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_for_blocked_test(), + updated_patch=updateDeployment(blocked=False), + ) + assert result["blocked"] is False + + def test_update_db_model_omits_blocked_when_patch_is_none(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_for_blocked_test(), + updated_patch=updateDeployment(), + ) + assert "blocked" not in result + + +class TestGetModelInfoWithIdBlocked: + """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` + column into the in-memory `model_info` dict so the router filter can read it.""" + + def test_get_model_info_with_id_propagates_blocked_true(self): + from litellm.proxy.proxy_server import ProxyConfig + + model = MagicMock() + model.model_id = "dep-1" + model.model_info = {} + model.blocked = True + info = ProxyConfig().get_model_info_with_id(model=model, db_model=True) + assert info.id == "dep-1" + assert getattr(info, "blocked") is True + + def test_get_model_info_with_id_defaults_blocked_to_false_when_missing(self): + from litellm.proxy.proxy_server import ProxyConfig + + model = MagicMock(spec=["model_id", "model_info"]) + model.model_id = "dep-2" + model.model_info = {} + info = ProxyConfig().get_model_info_with_id(model=model, db_model=True) + assert getattr(info, "blocked") is False + + +class TestPatchModelBlockedAuthGate: + """Only proxy admins may flip `blocked` — team admins authorized for + team-scoped models via `can_user_make_model_call` must still be rejected + when they attempt to toggle the pause flag.""" + + @pytest.mark.asyncio + async def test_team_admin_cannot_toggle_blocked(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + + non_admin = UserAPIKeyAuth( + user_id="team_admin", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + existing_row = MagicMock() + existing_row.litellm_params = {"model": "openai/gpt-4o-mini"} + existing_row.model_dump.return_value = { + "model_name": "gpt-4o-mini", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": "m1"}, + } + existing_row.model_dump_json.return_value = "{}" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=existing_row + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(Exception) as exc_info: + await patch_model( + model_id="m1", + patch_data=updateDeployment(blocked=True), + user_api_key_dict=non_admin, + ) + err = exc_info.value + assert getattr(err, "param", "") == "blocked" + assert "proxy admin" in getattr(err, "message", "").lower() + + @pytest.mark.asyncio + async def test_proxy_admin_can_toggle_blocked(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + existing_row = MagicMock() + existing_row.litellm_params = {"model": "openai/gpt-4o-mini"} + existing_row.model_dump.return_value = { + "model_name": "gpt-4o-mini", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": "m1"}, + } + existing_row.model_dump_json.return_value = "{}" + updated_row = MagicMock() + updated_row.model_dump_json.return_value = "{}" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=existing_row + ) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock( + return_value=updated_row + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=None), + ), + ): + result = await patch_model( + model_id="m1", + patch_data=updateDeployment(blocked=True), + user_api_key_dict=admin, + ) + assert result is updated_row + mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index d5c1132c5fd..1be4abbec6e 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -2057,3 +2057,317 @@ def test_openrouter_gemini_3_1_flash_lite_preview_pricing(): assert model_info["output_cost_per_token"] == 1.5e-06 assert model_info["max_input_tokens"] == 1048576 assert model_info["max_output_tokens"] == 65536 + + +def test_custom_pricing_applies_cache_read_input_cost(): + """ + Bug 1 reproduction: custom_cost_per_token with cache_read_input_token_cost + should bill cached prompt tokens at the cache rate, not the full input rate. + """ + usage = Usage( + prompt_tokens=6074, + completion_tokens=285, + total_tokens=6359, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=3456, + audio_tokens=0, + ), + ) + + response = ModelResponse( + id="test-id", + created=1234567890, + model="openai/gpt-5.4", + object="chat.completion", + choices=[], + usage=usage, + ) + + cost = litellm.completion_cost( + completion_response=response, + model="openai/gpt-5.4", + custom_llm_provider="openai", + custom_cost_per_token={ + "input_cost_per_token": 0.0000025, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.00000025, + }, + ) + + expected = (6074 - 3456) * 0.0000025 + 3456 * 0.00000025 + 285 * 0.000015 + + assert cost == pytest.approx(expected) + + +def test_custom_pricing_applies_cache_creation_input_cost_via_prompt_details(): + """ + OpenAI-compatible providers report cache-write tokens under + prompt_tokens_details.cache_creation_tokens. The custom-pricing helper must + bill those at cache_creation_input_token_cost, not the full input rate. + """ + pt_details = PromptTokensDetailsWrapper(cached_tokens=1000, audio_tokens=0) + pt_details.cache_creation_tokens = 500 + + usage = Usage( + prompt_tokens=4000, + completion_tokens=100, + total_tokens=4100, + prompt_tokens_details=pt_details, + ) + + response = ModelResponse( + id="test-id", + created=1234567890, + model="openai/gpt-5.4", + object="chat.completion", + choices=[], + usage=usage, + ) + + cost = litellm.completion_cost( + completion_response=response, + model="openai/gpt-5.4", + custom_llm_provider="openai", + custom_cost_per_token={ + "input_cost_per_token": 0.0000025, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.00000025, + "cache_creation_input_token_cost": 0.000003125, + }, + ) + + expected = ( + (4000 - 1000 - 500) * 0.0000025 + + 1000 * 0.00000025 + + 500 * 0.000003125 + + 100 * 0.000015 + ) + + assert cost == pytest.approx(expected) + + +def test_custom_pricing_applies_cache_creation_input_cost_via_cache_write_tokens_alias(): + """ + Some OpenAI-compatible providers (e.g. kimi-k2) emit cache-write tokens as + `cache_write_tokens` rather than `cache_creation_tokens`. The cost + calculator must mirror db_spend_update_writer and accept either name — + otherwise daily aggregation counts the tokens but the per-request cost + bills them at the full input rate. + + Drives `cost_per_token` directly with a SimpleNamespace usage stub so the + `cache_write_tokens` alias survives the call (Pydantic's Usage init + rebuilds prompt_tokens_details and drops dynamic attributes). + """ + from types import SimpleNamespace + + from litellm.cost_calculator import cost_per_token + + pt_details = SimpleNamespace(cached_tokens=1000, cache_write_tokens=500) + usage_stub = SimpleNamespace( + prompt_tokens=4000, + completion_tokens=100, + total_tokens=4100, + prompt_tokens_details=pt_details, + cache_read_input_tokens=None, + cache_creation_input_tokens=None, + ) + + prompt_cost, completion_cost = cost_per_token( + model="moonshotai/kimi-k2", + prompt_tokens=4000, + completion_tokens=100, + custom_llm_provider="openai", + usage_object=usage_stub, + custom_cost_per_token={ + "input_cost_per_token": 0.0000025, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.00000025, + "cache_creation_input_token_cost": 0.000003125, + }, + ) + + expected_prompt = ( + (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 + ) + expected_completion = 100 * 0.000015 + + assert prompt_cost == pytest.approx(expected_prompt) + assert completion_cost == pytest.approx(expected_completion) + + +# --------------------------------------------------------------------------- +# Bug 2 — db_spend_update_writer cache token extraction helpers. +# --------------------------------------------------------------------------- + + +def test_extract_cache_read_tokens_anthropic_top_level(): + from litellm.proxy.db.db_spend_update_writer import _extract_cache_read_tokens + + usage_obj = { + "prompt_tokens": 100, + "cache_read_input_tokens": 80, + "prompt_tokens_details": {"cached_tokens": 80}, + } + # Anthropic top-level value should win over prompt_tokens_details fallback. + assert _extract_cache_read_tokens(usage_obj) == 80 + + +def test_extract_cache_read_tokens_openai_compatible_fallback(): + from litellm.proxy.db.db_spend_update_writer import _extract_cache_read_tokens + + # Anthropic field absent — fall back to prompt_tokens_details.cached_tokens. + usage_obj = { + "prompt_tokens": 22583, + "prompt_tokens_details": {"cached_tokens": 22016}, + } + assert _extract_cache_read_tokens(usage_obj) == 22016 + + +def test_extract_cache_read_tokens_zero_when_missing(): + from litellm.proxy.db.db_spend_update_writer import _extract_cache_read_tokens + + assert _extract_cache_read_tokens({}) == 0 + assert _extract_cache_read_tokens({"cache_read_input_tokens": None}) == 0 + assert ( + _extract_cache_read_tokens({"prompt_tokens_details": {"cached_tokens": None}}) + == 0 + ) + + +def test_extract_cache_creation_tokens_anthropic_top_level(): + from litellm.proxy.db.db_spend_update_writer import ( + _extract_cache_creation_tokens, + ) + + usage_obj = { + "prompt_tokens": 100, + "cache_creation_input_tokens": 50, + "prompt_tokens_details": {"cache_write_tokens": 50}, + } + # Anthropic top-level should short-circuit the fallback. + assert _extract_cache_creation_tokens(usage_obj) == 50 + + +def test_extract_cache_creation_tokens_openai_cache_write_alias(): + from litellm.proxy.db.db_spend_update_writer import ( + _extract_cache_creation_tokens, + ) + + # kimi-k2 emits cache_write_tokens. + usage_obj = { + "prompt_tokens": 1000, + "prompt_tokens_details": {"cache_write_tokens": 200}, + } + assert _extract_cache_creation_tokens(usage_obj) == 200 + + +def test_extract_cache_creation_tokens_openai_cache_creation_alias(): + from litellm.proxy.db.db_spend_update_writer import ( + _extract_cache_creation_tokens, + ) + + # Other OpenAI-compatible providers emit cache_creation_tokens. + usage_obj = { + "prompt_tokens": 1000, + "prompt_tokens_details": {"cache_creation_tokens": 300}, + } + assert _extract_cache_creation_tokens(usage_obj) == 300 + + +def test_extract_cache_creation_tokens_zero_when_missing(): + from litellm.proxy.db.db_spend_update_writer import ( + _extract_cache_creation_tokens, + ) + + assert _extract_cache_creation_tokens({}) == 0 + assert _extract_cache_creation_tokens({"cache_creation_input_tokens": None}) == 0 + assert ( + _extract_cache_creation_tokens( + {"prompt_tokens_details": {"cache_write_tokens": None}} + ) + == 0 + ) + + +def test_custom_pricing_anthropic_style_cache_tokens_not_double_counted(): + """ + Anthropic providers report cache tokens at the top level of Usage, and + `prompt_tokens` EXCLUDES them. The helper expects `prompt_tokens` to + include cache tokens, so cost_per_token must adjust before invoking it — + otherwise regular_prompt_tokens goes negative and clamps to 0. + """ + usage = Usage( + prompt_tokens=2000, + completion_tokens=100, + total_tokens=2100, + cache_read_input_tokens=1500, + cache_creation_input_tokens=300, + ) + + response = ModelResponse( + id="test-id", + created=1234567890, + model="anthropic/claude-3-5-sonnet", + object="chat.completion", + choices=[], + usage=usage, + ) + + cost = litellm.completion_cost( + completion_response=response, + model="anthropic/claude-3-5-sonnet", + custom_llm_provider="anthropic", + custom_cost_per_token={ + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.0000003, + "cache_creation_input_token_cost": 0.00000375, + }, + ) + + # Anthropic prompt_tokens=2000 excludes cache. After normalization the + # helper sees 2000 + 1500 + 300 = 3800, of which 2000 are uncached. + expected = 2000 * 0.000003 + 1500 * 0.0000003 + 300 * 0.00000375 + 100 * 0.000015 + + assert cost == pytest.approx(expected) + + +def test_custom_pricing_without_cache_keys_preserves_legacy_behavior(): + """ + Backward compatibility: when custom_cost_per_token omits both cache rates, + cached tokens must be billed at input_cost_per_token (matching the pre-fix + behavior) so existing callers see no change. + """ + usage = Usage( + prompt_tokens=1000, + completion_tokens=100, + total_tokens=1100, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=400, + audio_tokens=0, + ), + ) + + response = ModelResponse( + id="test-id", + created=1234567890, + model="openai/gpt-5.4", + object="chat.completion", + choices=[], + usage=usage, + ) + + cost = litellm.completion_cost( + completion_response=response, + model="openai/gpt-5.4", + custom_llm_provider="openai", + custom_cost_per_token={ + "input_cost_per_token": 0.0000025, + "output_cost_per_token": 0.000015, + }, + ) + + # All 1000 prompt tokens billed at input rate, regardless of cached_tokens. + expected = 1000 * 0.0000025 + 100 * 0.000015 + + assert cost == pytest.approx(expected) diff --git a/tests/test_litellm/test_guardrail_exception_status_codes.py b/tests/test_litellm/test_guardrail_exception_status_codes.py new file mode 100644 index 00000000000..c4df1295580 --- /dev/null +++ b/tests/test_litellm/test_guardrail_exception_status_codes.py @@ -0,0 +1,66 @@ +""" +Tests for guardrail exception status codes. + +GuardrailRaisedException and BlockedPiiEntityError must carry +``status_code = 400`` so the proxy exception handler +(``getattr(e, "status_code", 500)``) returns HTTP 400 instead of 500 +for intentional guardrail blocks. +""" + +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException + + +class TestGuardrailRaisedExceptionStatusCode: + """GuardrailRaisedException should default to status_code=400.""" + + def test_default_status_code(self): + exc = GuardrailRaisedException( + guardrail_name="test_guardrail", + message="blocked", + ) + assert exc.status_code == 400 + + def test_custom_status_code(self): + exc = GuardrailRaisedException( + guardrail_name="test_guardrail", + message="rate limited", + status_code=429, + ) + assert exc.status_code == 429 + + def test_getattr_fallback_resolves_to_400(self): + """The proxy uses ``getattr(e, 'status_code', 500)`` — verify it + resolves to 400, not the 500 default.""" + exc = GuardrailRaisedException( + guardrail_name="test_guardrail", + message="blocked", + ) + assert getattr(exc, "status_code", 500) == 400 + + +class TestBlockedPiiEntityErrorStatusCode: + """BlockedPiiEntityError should default to status_code=400.""" + + def test_default_status_code(self): + exc = BlockedPiiEntityError( + entity_type="CREDIT_CARD", + guardrail_name="presidio", + ) + assert exc.status_code == 400 + + def test_custom_status_code(self): + exc = BlockedPiiEntityError( + entity_type="SSN", + guardrail_name="presidio", + status_code=403, + ) + assert exc.status_code == 403 + + def test_getattr_fallback_resolves_to_400(self): + """The proxy uses ``getattr(e, 'status_code', 500)`` — verify it + resolves to 400, not the 500 default.""" + exc = BlockedPiiEntityError( + entity_type="PHONE_NUMBER", + guardrail_name="presidio", + ) + assert getattr(exc, "status_code", 500) == 400 diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 282c9d72d51..a89e30a0e06 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -1,5 +1,4 @@ import json -import os from unittest.mock import MagicMock, patch import pytest @@ -165,6 +164,13 @@ def test_max_connections_in_cluster_kwargs(): ), "max_connections should be in available Redis cluster kwargs" +def test_socket_timeouts_in_cluster_kwargs(): + """Test that Redis cluster clients can receive socket timeout configuration""" + kwargs = _get_redis_cluster_kwargs() + assert "socket_timeout" in kwargs + assert "socket_connect_timeout" in kwargs + + def test_get_redis_async_client_with_connection_pool(): """Test that connection_pool parameter is properly passed to Redis client""" # Create a mock connection pool diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 48facace528..d8be527689e 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3697,3 +3697,173 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" ) + + +def _router_with_two_deployments(blocked_flags): + import litellm + + model_list = [] + for idx, blocked in enumerate(blocked_flags): + model_list.append( + { + "model_name": "gpt-4o", + "litellm_params": {"model": f"openai/gpt-4o-{idx}"}, + "model_info": {"id": f"dep-{idx}", "blocked": blocked}, + } + ) + return litellm.Router(model_list=model_list) + + +def test_get_fully_blocked_model_names_marks_name_when_all_deployments_blocked(): + router = _router_with_two_deployments([True, True]) + assert router.get_fully_blocked_model_names() == {"gpt-4o"} + + +def test_get_fully_blocked_model_names_keeps_name_when_partial_blocked(): + router = _router_with_two_deployments([True, False]) + assert router.get_fully_blocked_model_names() == set() + + +def test_get_fully_blocked_model_names_treats_missing_key_as_unblocked(): + import litellm + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "dep-0"}, + } + ] + ) + assert router.get_fully_blocked_model_names() == set() + + +@pytest.mark.asyncio +async def test_async_get_healthy_deployments_skips_blocked_deployment(): + router = _router_with_two_deployments([True, False]) + healthy, all_dep = await router._async_get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) + healthy_ids = [d["model_info"]["id"] for d in healthy] + assert "dep-0" not in healthy_ids + assert "dep-1" in healthy_ids + assert len(all_dep) == 2 + + +def test_get_healthy_deployments_sync_skips_blocked_deployment(): + router = _router_with_two_deployments([False, True]) + healthy, all_dep = router._get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) + healthy_ids = [d["model_info"]["id"] for d in healthy] + assert "dep-0" in healthy_ids + assert "dep-1" not in healthy_ids + assert len(all_dep) == 2 + + +def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): + router = _router_with_two_deployments([True, False]) + filtered = router._filter_blocked_deployments(router.get_model_list() or []) + ids = [d["model_info"]["id"] for d in filtered] + assert ids == ["dep-1"] + + +@pytest.mark.asyncio +async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): + router = _router_with_two_deployments([True, False]) + deployments = await router.async_get_healthy_deployments( + model="gpt-4o", request_kwargs={} + ) + assert isinstance(deployments, list) + ids = [d["model_info"]["id"] for d in deployments] + assert "dep-0" not in ids + assert "dep-1" in ids + + +def test_public_get_available_deployment_skips_blocked_on_primary_path(): + router = _router_with_two_deployments([True, False]) + deployment = router.get_available_deployment(model="gpt-4o", request_kwargs={}) + assert deployment["model_info"]["id"] == "dep-1" + + +def test_get_available_deployment_raises_when_addressed_dict_is_blocked(): + import litellm + + router = _router_with_two_deployments([True, True]) + with pytest.raises(litellm.ServiceUnavailableError): + router.get_available_deployment(model="dep-0", request_kwargs={}) + + +def _router_with_two_pass_through_deployments(blocked_flags): + import litellm + + model_list = [] + for idx, blocked in enumerate(blocked_flags): + model_list.append( + { + "model_name": "gpt-4o", + "litellm_params": { + "model": f"openai/gpt-4o-{idx}", + "api_key": "sk-fake-for-tests", + "use_in_pass_through": True, + }, + "model_info": {"id": f"pt-{idx}", "blocked": blocked}, + } + ) + return litellm.Router(model_list=model_list) + + +def test_get_available_deployment_for_pass_through_skips_blocked(): + router = _router_with_two_pass_through_deployments([True, False]) + deployment = router.get_available_deployment_for_pass_through( + model="gpt-4o", request_kwargs={} + ) + assert deployment["model_info"]["id"] == "pt-1" + + +def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): + import litellm + + router = _router_with_two_pass_through_deployments([True, True]) + with pytest.raises(litellm.ServiceUnavailableError): + router.get_available_deployment_for_pass_through( + model="pt-0", request_kwargs={} + ) + + +def test_get_deployment_credentials_returns_none_for_blocked_deployment(): + router = _router_with_two_deployments([True, False]) + assert router.get_deployment_credentials(model_id="dep-0") is None + assert router.get_deployment_credentials(model_id="dep-1") is not None + + +def test_get_deployment_credentials_with_provider_returns_none_for_blocked_deployment(): + router = _router_with_two_deployments([True, False]) + assert router.get_deployment_credentials_with_provider(model_id="dep-0") is None + assert router.get_deployment_credentials_with_provider(model_id="dep-1") is not None + + +def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): + """ + Exercises Router._is_deployment_blocked so router_code_coverage.py (AST call graph) + marks the helper as covered by router-named tests. + """ + import types + + import litellm + + router = _router_with_two_deployments([True, False]) + blocked_dep = router.get_deployment("dep-0") + unblocked_dep = router.get_deployment("dep-1") + assert blocked_dep is not None and unblocked_dep is not None + assert litellm.Router._is_deployment_blocked(blocked_dep) is True + assert litellm.Router._is_deployment_blocked(unblocked_dep) is False + + # No model_info on deployment object → treated as not blocked + assert litellm.Router._is_deployment_blocked(object()) is False + missing_blocked = types.SimpleNamespace() + assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False + assert litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) + ) is True diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 50dd5e78f37..41dfdb21d1b 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -654,12 +654,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { Input Tokens - {Math.max( - 0, - (userSpendData.metadata?.total_prompt_tokens || 0) - - (userSpendData.metadata?.total_cache_read_input_tokens || 0) - - (userSpendData.metadata?.total_cache_creation_input_tokens || 0) - ).toLocaleString()} + {(userSpendData.metadata?.total_prompt_tokens || 0).toLocaleString()} diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx index 8ae90c1cbbc..554fe86bf66 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx @@ -52,6 +52,8 @@ const AdvancedSettings: React.FC = ({ form.setFieldsValue({ input_cost_per_token: undefined, output_cost_per_token: undefined, + cache_read_input_token_cost: undefined, + cache_creation_input_token_cost: undefined, input_cost_per_second: undefined, }); } @@ -211,6 +213,24 @@ const AdvancedSettings: React.FC = ({ > + + + + + + ) : ( , ac if (formValues.output_cost_per_token !== undefined && formValues.output_cost_per_token !== null && formValues.output_cost_per_token !== "") { formValues.output_cost_per_token = Number(formValues.output_cost_per_token) / 1000000; } + + // Cache Read Cost: if blank, default to Input Cost (already token-unit converted above) + if ( + formValues.cache_read_input_token_cost !== undefined && + formValues.cache_read_input_token_cost !== null && + formValues.cache_read_input_token_cost !== "" + ) { + formValues.cache_read_input_token_cost = + Number(formValues.cache_read_input_token_cost) / 1000000; + } else if ( + formValues.input_cost_per_token !== undefined && + formValues.input_cost_per_token !== null && + formValues.input_cost_per_token !== "" + ) { + formValues.cache_read_input_token_cost = Number(formValues.input_cost_per_token); + } else { + delete formValues.cache_read_input_token_cost; + } + + // Cache Write Cost: explicit value if provided, else leave unset so the + // backend keeps the model-level default (per-second pricing, model_prices + // entries, etc.). Sending 0 here would overwrite that default. + // The backend falls back to input_cost_per_token when this key is absent. + if ( + formValues.cache_creation_input_token_cost !== undefined && + formValues.cache_creation_input_token_cost !== null && + formValues.cache_creation_input_token_cost !== "" + ) { + formValues.cache_creation_input_token_cost = + Number(formValues.cache_creation_input_token_cost) / 1000000; + } else { + delete formValues.cache_creation_input_token_cost; + } // Keep input_cost_per_second as is, no conversion needed // Iterate through the key-value pairs in formValues @@ -119,7 +152,13 @@ export const prepareModelAddRequest = async (formValues: Record, ac } // Handle the pricing fields - else if (key === "input_cost_per_token" || key === "output_cost_per_token" || key === "input_cost_per_second") { + else if ( + key === "input_cost_per_token" || + key === "output_cost_per_token" || + key === "input_cost_per_second" || + key === "cache_read_input_token_cost" || + key === "cache_creation_input_token_cost" + ) { if (value !== undefined && value !== null && value !== "") { litellmParamsObj[key] = Number(value); } diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 95a43862de5..5ed4c0468b8 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -263,6 +263,34 @@ export default function ModelInfoView({ updatedLitellmParams.output_cost_per_token = Number(values.output_cost) / 1_000_000; } + // Cache Read Cost: explicit value if provided, else fall back to input cost (when input cost touched). + if (form.isFieldTouched("cache_read_cost") || form.isFieldTouched("input_cost")) { + if ( + values.cache_read_cost !== undefined && + values.cache_read_cost !== null && + values.cache_read_cost !== "" + ) { + updatedLitellmParams.cache_read_input_token_cost = Number(values.cache_read_cost) / 1_000_000; + } else if (updatedLitellmParams.input_cost_per_token !== undefined) { + updatedLitellmParams.cache_read_input_token_cost = updatedLitellmParams.input_cost_per_token; + } + } + + // Cache Write Cost: explicit value if provided, else clear the override + // so the backend falls back to the model-level default. Sending 0 here + // would persist a zero rate even when the user intended to unset it. + if (form.isFieldTouched("cache_write_cost")) { + if ( + values.cache_write_cost !== undefined && + values.cache_write_cost !== null && + values.cache_write_cost !== "" + ) { + updatedLitellmParams.cache_creation_input_token_cost = Number(values.cache_write_cost) / 1_000_000; + } else { + delete updatedLitellmParams.cache_creation_input_token_cost; + } + } + if (values.litellm_credential_name) { updatedLitellmParams.litellm_credential_name = values.litellm_credential_name; } else { @@ -638,6 +666,22 @@ export default function ModelInfoView({ output_cost: localModelData.litellm_params?.output_cost_per_token ? localModelData.litellm_params.output_cost_per_token * 1_000_000 : localModelData.model_info?.output_cost_per_token * 1_000_000 || null, + cache_read_cost: + localModelData.litellm_params?.cache_read_input_token_cost !== undefined && + localModelData.litellm_params?.cache_read_input_token_cost !== null + ? localModelData.litellm_params.cache_read_input_token_cost * 1_000_000 + : localModelData.model_info?.cache_read_input_token_cost !== undefined && + localModelData.model_info?.cache_read_input_token_cost !== null + ? localModelData.model_info.cache_read_input_token_cost * 1_000_000 + : null, + cache_write_cost: + localModelData.litellm_params?.cache_creation_input_token_cost !== undefined && + localModelData.litellm_params?.cache_creation_input_token_cost !== null + ? localModelData.litellm_params.cache_creation_input_token_cost * 1_000_000 + : localModelData.model_info?.cache_creation_input_token_cost !== undefined && + localModelData.model_info?.cache_creation_input_token_cost !== null + ? localModelData.model_info.cache_creation_input_token_cost * 1_000_000 + : null, cache_control: localModelData.litellm_params?.cache_control_injection_points ? true : false, cache_control_injection_points: localModelData.litellm_params?.cache_control_injection_points || [], model_access_group: Array.isArray(localModelData.model_info?.access_groups) @@ -725,6 +769,52 @@ export default function ModelInfoView({ )} +
+ Cache Read Cost (per 1M tokens) + {isEditing ? ( + + + + ) : ( +
+ {localModelData?.litellm_params?.cache_read_input_token_cost !== undefined && + localModelData?.litellm_params?.cache_read_input_token_cost !== null + ? (localModelData.litellm_params.cache_read_input_token_cost * 1_000_000).toFixed(4) + : localModelData?.model_info?.cache_read_input_token_cost !== undefined && + localModelData?.model_info?.cache_read_input_token_cost !== null + ? (localModelData.model_info.cache_read_input_token_cost * 1_000_000).toFixed(4) + : "Not Set"} +
+ )} +
+ +
+ Cache Write Cost (per 1M tokens) + {isEditing ? ( + + + + ) : ( +
+ {localModelData?.litellm_params?.cache_creation_input_token_cost !== undefined && + localModelData?.litellm_params?.cache_creation_input_token_cost !== null + ? (localModelData.litellm_params.cache_creation_input_token_cost * 1_000_000).toFixed(4) + : localModelData?.model_info?.cache_creation_input_token_cost !== undefined && + localModelData?.model_info?.cache_creation_input_token_cost !== null + ? (localModelData.model_info.cache_creation_input_token_cost * 1_000_000).toFixed(4) + : "Not Set"} +
+ )} +
+
API Base {isEditing ? ( From 73e32a31bfe595525f9b5a532f5db3478228eb47 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 19 May 2026 04:58:14 +0530 Subject: [PATCH 08/11] feat(prometheus): add user_email and user_alias to user budget metrics (#28155) * feat(prometheus): add user_email and user_alias to user budget metrics User budget Prometheus gauges now expose human-readable labels alongside user_id, matching team and API key budget metrics for Grafana filtering. Co-authored-by: Cursor * fix(prometheus): gate user budget email/alias labels behind opt-in flag Address greptile review: adding labels to existing metrics is a breaking cardinality change. Gate behind prometheus_user_budget_label_include_email_alias=True (default: False) so existing dashboards and recording rules are unaffected. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- litellm/__init__.py | 1 + litellm/integrations/prometheus.py | 6 ++ litellm/types/integrations/prometheus.py | 30 ++++-- tests/otel_tests/test_prometheus.py | 12 +-- .../test_prometheus_user_team_metrics.py | 98 +++++++++++++++++++ 5 files changed, 129 insertions(+), 18 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index b9da0524095..c868ae55b4f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -416,6 +416,7 @@ custom_prometheus_metadata_labels: List[str] = [] custom_prometheus_tags: List[str] = [] prometheus_metrics_config: Optional[List] = None prometheus_emit_stream_label: bool = False +prometheus_user_budget_label_include_email_alias: bool = False prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000 prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0 prometheus_end_user_metrics_cleanup_interval_seconds: Optional[float] = 60.0 diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 30af0dcb8ed..2c63455565c 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3540,6 +3540,10 @@ class PrometheusLogger(CustomLogger): user_object.budget_reset_at = user_info.budget_reset_at if user_object.max_budget is None and user_info.max_budget is not None: user_object.max_budget = user_info.max_budget + if user_info.user_email is not None: + user_object.user_email = user_info.user_email + if user_info.user_alias is not None: + user_object.user_alias = user_info.user_alias return user_object @@ -3556,6 +3560,8 @@ class PrometheusLogger(CustomLogger): """ enum_values = UserAPIKeyLabelValues( user=user.user_id, + user_email=user.user_email or "", + user_alias=user.user_alias or "", ) _labels = prometheus_label_factory( diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 43a287f29bc..7b5c5ab2969 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -160,6 +160,7 @@ class UserAPIKeyLabelNames(Enum): END_USER = "end_user" USER = "user" USER_EMAIL = "user_email" + USER_ALIAS = "user_alias" API_KEY_HASH = "hashed_api_key" API_KEY_ALIAS = "api_key_alias" TEAM = "team" @@ -533,17 +534,9 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, ] - litellm_user_max_budget_metric = [ - UserAPIKeyLabelNames.USER.value, - ] + litellm_user_max_budget_metric = litellm_remaining_user_budget_metric - litellm_user_budget_remaining_hours_metric = [ - UserAPIKeyLabelNames.USER.value, - ] - - litellm_user_budget_remaining_hours_metric = [ - UserAPIKeyLabelNames.USER.value, - ] + litellm_user_budget_remaining_hours_metric = litellm_remaining_user_budget_metric litellm_remaining_api_key_requests_for_model = [ UserAPIKeyLabelNames.API_KEY_HASH.value, @@ -730,6 +723,22 @@ class PrometheusMetricLabels: ): custom_labels.append(UserAPIKeyLabelNames.STREAM.value) + _user_budget_metrics = { + "litellm_remaining_user_budget_metric", + "litellm_user_max_budget_metric", + "litellm_user_budget_remaining_hours_metric", + } + if ( + label_name in _user_budget_metrics + and litellm.prometheus_user_budget_label_include_email_alias is True + ): + for label in [ + UserAPIKeyLabelNames.USER_EMAIL.value, + UserAPIKeyLabelNames.USER_ALIAS.value, + ]: + if label not in default_labels and label not in custom_labels: + custom_labels.append(label) + if label_name in PrometheusMetricLabels._org_label_metrics: for label in [ UserAPIKeyLabelNames.ORG_ID.value, @@ -759,6 +768,7 @@ class UserAPIKeyLabelValues: end_user: Optional[str] = None user: Optional[str] = None user_email: Optional[str] = None + user_alias: Optional[str] = None hashed_api_key: Optional[str] = None api_key_alias: Optional[str] = None team: Optional[str] = None diff --git a/tests/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py index c9490af07cb..90c71037609 100644 --- a/tests/otel_tests/test_prometheus.py +++ b/tests/otel_tests/test_prometheus.py @@ -610,22 +610,18 @@ def extract_user_budget_metrics(metrics_text: str, user_id: str) -> Dict[str, fl # Escape user_id for regex pattern matching escaped_user_id = re.escape(user_id) - # Get remaining budget - remaining_pattern = ( - f'litellm_remaining_user_budget_metric{{user="{escaped_user_id}"}} ([0-9.]+)' - ) + # Get remaining budget (user_email and user_alias may also be present as labels) + remaining_pattern = rf'litellm_remaining_user_budget_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)' remaining_match = re.search(remaining_pattern, metrics_text) metrics["remaining"] = float(remaining_match.group(1)) if remaining_match else None # Get total budget - total_pattern = ( - f'litellm_user_max_budget_metric{{user="{escaped_user_id}"}} ([0-9.]+)' - ) + total_pattern = rf'litellm_user_max_budget_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)' total_match = re.search(total_pattern, metrics_text) metrics["total"] = float(total_match.group(1)) if total_match else None # Get remaining hours - hours_pattern = f'litellm_user_budget_remaining_hours_metric{{user="{escaped_user_id}"}} ([0-9.]+)' + hours_pattern = rf'litellm_user_budget_remaining_hours_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)' hours_match = re.search(hours_pattern, metrics_text) metrics["remaining_hours"] = float(hours_match.group(1)) if hours_match else None diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index 19ae819c85a..12f30ab6024 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -460,6 +460,104 @@ async def test_assemble_user_object_does_not_override_metadata_max_budget( ), "max_budget from metadata must not be replaced by the DB value" +async def test_assemble_user_object_populates_user_email_and_alias_from_db( + prometheus_logger, +): + db_user = MagicMock() + db_user.max_budget = None + db_user.budget_reset_at = None + db_user.user_email = "alice@example.com" + db_user.user_alias = "Alice" + + with patch("litellm.proxy.auth.auth_checks.get_user_object") as mock_get_user: + mock_get_user.return_value = db_user + user_object = await prometheus_logger._assemble_user_object( + user_id="user-abc-123", + spend=10.0, + max_budget=None, + response_cost=0.5, + ) + + assert user_object.user_email == "alice@example.com" + assert user_object.user_alias == "Alice" + + +def test_set_user_budget_metrics_default_no_email_alias_labels( + prometheus_logger, +): + """By default (flag off), only user label is emitted.""" + import litellm + from litellm.proxy._types import LiteLLM_UserTable + + litellm.prometheus_user_budget_label_include_email_alias = False + + user = LiteLLM_UserTable( + user_id="user-abc-123", + user_email="alice@example.com", + user_alias="Alice", + spend=25.0, + max_budget=100.0, + budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc), + ) + + prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() + prometheus_logger.litellm_user_max_budget_metric = MagicMock() + prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() + + prometheus_logger._set_user_budget_metrics(user) + + prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with( + user="user-abc-123", + ) + + +def test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in( + prometheus_logger, +): + """When prometheus_user_budget_label_include_email_alias=True, email+alias labels appear.""" + import litellm + from litellm.proxy._types import LiteLLM_UserTable + + litellm.prometheus_user_budget_label_include_email_alias = True + + user = LiteLLM_UserTable( + user_id="user-abc-123", + user_email="alice@example.com", + user_alias="Alice", + spend=25.0, + max_budget=100.0, + budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc), + ) + + prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() + prometheus_logger.litellm_user_max_budget_metric = MagicMock() + prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() + + try: + prometheus_logger._set_user_budget_metrics(user) + + prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with( + user="user-abc-123", + user_email="alice@example.com", + user_alias="Alice", + ) + prometheus_logger.litellm_remaining_user_budget_metric.labels().set.assert_called_once_with( + 75.0 + ) + prometheus_logger.litellm_user_max_budget_metric.labels.assert_called_once_with( + user="user-abc-123", + user_email="alice@example.com", + user_alias="Alice", + ) + prometheus_logger.litellm_user_budget_remaining_hours_metric.labels.assert_called_once_with( + user="user-abc-123", + user_email="alice@example.com", + user_alias="Alice", + ) + finally: + litellm.prometheus_user_budget_label_include_email_alias = False + + async def test_set_user_budget_metrics_after_api_request_no_inf_when_metadata_budget_none( prometheus_logger, ): From a7f3dbcbe37caba5cfedf6f502fceacb8f658a05 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 18 May 2026 16:39:02 -0700 Subject: [PATCH 09/11] test(callbacks): harden flaky proxy callback-leak detector (#28195) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test(callbacks): TEMP diagnostic probe for callback-leak flake Hardened leak detector (sample N, flag sustained monotonic per-type growth, normalize instance addresses) + a temporary always-fail probe on test_check_num_callbacks_on_lowest_latency that dumps the per-type series and raw reprs via the JUnit failure message, to settle real-leak vs bounded-pollution on CCI. Diagnostic block is clearly marked and will be reverted before the PR. * test(callbacks): harden proxy callback-leak detector, drop diagnostic CCI diagnostic confirmed the 85->95 jump is a bounded one-time registration from the test's own switch to latency-based-routing (+LowestLatencyLoggingHandler, +SlackAlerting), flat at 95 for 2.5 min under load — not a leak. Final detector: settle past the deliberate config/update, sample N times, flag only sustained monotonic per-type growth, normalize instance addresses, name the leaking type on failure. Removes the temporary always-fail probe. * test(callbacks): address review - drop redundant settle, close terminal-burst blind spot - test_check_num_callbacks: remove leftover sleep(30) before sleep(SETTLE_SECONDS) (60s -> 30s dead wait). - Add _terminal_suspects + _detect_leaks_confirmed: when monotonic net growth is confined to the final interval (escapes the >=2-interval guard), take one confirmation sample. A real ongoing leak keeps climbing and is flagged; a one-time terminal registration plateaus and is ignored. --- tests/test_callbacks_on_proxy.py | 261 +++++++++++++++++++++---------- 1 file changed, 182 insertions(+), 79 deletions(-) diff --git a/tests/test_callbacks_on_proxy.py b/tests/test_callbacks_on_proxy.py index 0b55d820532..17c0db9260f 100644 --- a/tests/test_callbacks_on_proxy.py +++ b/tests/test_callbacks_on_proxy.py @@ -9,12 +9,155 @@ import pytest import asyncio import aiohttp import os +import re import dotenv +from collections import Counter from dotenv import load_dotenv import pytest load_dotenv() +# A *leak* is sustained, monotonic growth of one callback TYPE across the whole +# sampling window. A one-time bump that then plateaus is benign pollution from +# other tests sharing this proxy (this suite runs `pytest -n 4` against a single +# proxy container, so other workers legitimately add team/key-scoped callbacks +# while this test sleeps). We therefore sample N times and only flag a type +# whose normalized count never decreases, grows in >=2 distinct intervals, and +# nets >= LEAK_MIN_NET_GROWTH overall. +NUM_SAMPLES = 4 +SAMPLE_INTERVAL_SECONDS = 20 +LEAK_MIN_NET_GROWTH = 5 +LEAK_MIN_GROWING_INTERVALS = 2 +# A routing-strategy switch / alerting config is a *known, bounded, one-time* +# registration (CCI diagnostic 2026-05-16: total 85->95 on the first interval +# after switching to latency-based-routing, then flat at 95 for 2.5 min under +# load). We absorb that step by settling before the baseline sample, so only +# growth *after* the deliberate perturbation can count as a leak. +SETTLE_SECONDS = 30 + +# Strip instance-identity noise so N leaked instances of one class collapse to +# one rising counter instead of N opaque, unrelated-looking strings. +_ADDR_RE = re.compile(r" at 0x[0-9a-fA-F]+") +_OBJ_RE = re.compile(r"<([\w.]+) object") + + +def _normalize_callback(cb_str: str) -> str: + """Reduce a callback's str() to a stable type key (drops 0x… addresses).""" + s = _ADDR_RE.sub("", cb_str) + m = _OBJ_RE.search(s) + if m: + return m.group(1).split(".")[-1] + # bound methods: ">" -> "Cls.m" + bm = re.search(r"bound method ([\w.]+)", s) + if bm: + return bm.group(1) + return s.strip() + + +def _summarize(all_litellm_callbacks) -> Counter: + return Counter(_normalize_callback(str(c)) for c in all_litellm_callbacks) + + +def _detect_leaks(samples): + """ + samples: list[Counter] taken in time order. + + Returns {callback_type: [counts across samples]} for types that grew + monotonically (never decreased), in >=LEAK_MIN_GROWING_INTERVALS intervals, + and netted >=LEAK_MIN_NET_GROWTH overall — i.e. a real leak, not a one-shot + step from a parallel test. + """ + leaks = {} + all_types = set().union(*[set(s) for s in samples]) if samples else set() + for t in all_types: + series = [s.get(t, 0) for s in samples] + deltas = [b - a for a, b in zip(series, series[1:])] + net = series[-1] - series[0] + non_decreasing = all(d >= 0 for d in deltas) + growing_intervals = sum(1 for d in deltas if d > 0) + if ( + non_decreasing + and net >= LEAK_MIN_NET_GROWTH + and growing_intervals >= LEAK_MIN_GROWING_INTERVALS + ): + leaks[t] = series + return leaks + + +def _terminal_suspects(samples): + """ + Types whose net growth clears the threshold monotonically but is confined + to the *final* interval — `growing_intervals == 1` with that one growing + interval being the last. `_detect_leaks`' `>= 2` guard silently passes + these, so a real leak that accumulates entirely in the last sampled window + is indistinguishable from a one-time terminal step *without one more + sample*. Returns the set of such types so the caller can re-confirm. + """ + suspects = set() + all_types = set().union(*[set(s) for s in samples]) if samples else set() + for t in all_types: + series = [s.get(t, 0) for s in samples] + deltas = [b - a for a, b in zip(series, series[1:])] + if not deltas: + continue + net = series[-1] - series[0] + non_decreasing = all(d >= 0 for d in deltas) + growing = [i for i, d in enumerate(deltas) if d > 0] + if ( + non_decreasing + and net >= LEAK_MIN_NET_GROWTH + and growing == [len(deltas) - 1] + ): + suspects.add(t) + return suspects + + +async def _detect_leaks_confirmed(session, samples): + """ + `_detect_leaks`, plus a single confirmation sample when growth is confined + to the final interval (see `_terminal_suspects`). A genuine ongoing leak + keeps climbing -> now grows in >= 2 intervals -> flagged; a one-time + terminal registration plateaus -> still 1 growing interval -> ignored. + Returns `(leaks, samples)` (samples may have one extra entry appended). + """ + leaks = _detect_leaks(samples) + if not leaks and _terminal_suspects(samples): + await asyncio.sleep(SAMPLE_INTERVAL_SECONDS) + _, _, all_cb = await get_active_callbacks(session=session) + samples = samples + [_summarize(all_cb)] + leaks = _detect_leaks(samples) + return leaks, samples + + +def _format_report(samples, leaks) -> str: + lines = ["Callback count per type across samples (time order):"] + all_types = sorted(set().union(*[set(s) for s in samples])) + for t in all_types: + series = [s.get(t, 0) for s in samples] + marker = " <-- LEAK" if t in leaks else "" + lines.append(f" {t}: {series}{marker}") + totals = [sum(s.values()) for s in samples] + lines.append(f"TOTAL callbacks per sample: {totals}") + if leaks: + lines.append( + "Leaking callback types (sustained monotonic growth): " + + ", ".join(sorted(leaks)) + ) + return "\n".join(lines) + + +async def _sample_callbacks(session, num_samples, interval): + """Take `num_samples` callback snapshots `interval`s apart.""" + samples = [] + alerts = [] + for i in range(num_samples): + if i > 0: + await asyncio.sleep(interval) + num_cb, num_alert, all_cb = await get_active_callbacks(session=session) + samples.append(_summarize(all_cb)) + alerts.append(num_alert) + return samples, alerts + async def config_update(session, routing_strategy=None): url = "http://0.0.0.0:4000/config/update" @@ -97,105 +240,65 @@ async def get_current_routing_strategy(session): @pytest.mark.asyncio @pytest.mark.order1 +@pytest.mark.flaky(reruns=2, reruns_delay=5) async def test_check_num_callbacks(): """ - Test 1: num callbacks should NOT increase over time - -> check current callbacks - -> sleep for 30 seconds - -> check current callbacks - -> sleep for 30 seconds - -> check current callbacks + PROD invariant: no callback TYPE should grow without bound over time. + + This suite runs `pytest -n 4` against one shared proxy, so the raw count is + noisy — other workers legitimately add team/key-scoped callbacks that then + plateau. We settle first, then sample several times, and only fail on + *sustained, monotonic* per-type growth (a genuine leak), naming the type. """ - from litellm._uuid import uuid - async with aiohttp.ClientSession() as session: - await asyncio.sleep(30) - num_callbacks_1, _, all_litellm_callbacks_1 = await get_active_callbacks( - session=session - ) - assert num_callbacks_1 > 0 - await asyncio.sleep(30) + # Absorb proxy warmup / in-flight parallel registration before baseline. + await asyncio.sleep(SETTLE_SECONDS) - num_callbacks_2, _, all_litellm_callbacks_2 = await get_active_callbacks( - session=session + samples, _ = await _sample_callbacks( + session, NUM_SAMPLES, SAMPLE_INTERVAL_SECONDS ) - print("all_litellm_callbacks_1", all_litellm_callbacks_1) + assert sum(samples[0].values()) > 0, "expected some callbacks registered" - print( - "diff in callbacks=", - set(all_litellm_callbacks_1) - set(all_litellm_callbacks_2), - ) - - assert abs(num_callbacks_1 - num_callbacks_2) <= 4 - - await asyncio.sleep(30) - - num_callbacks_3, _, all_litellm_callbacks_3 = await get_active_callbacks( - session=session - ) - - print( - "diff in callbacks = all_litellm_callbacks3 - all_litellm_callbacks2 ", - set(all_litellm_callbacks_3) - set(all_litellm_callbacks_2), - ) - - assert abs(num_callbacks_3 - num_callbacks_2) <= 4 + leaks, samples = await _detect_leaks_confirmed(session, samples) + report = _format_report(samples, leaks) + print(report) + assert not leaks, f"Callback leak detected.\n{report}" @pytest.mark.asyncio @pytest.mark.order2 +@pytest.mark.flaky(reruns=2, reruns_delay=5) async def test_check_num_callbacks_on_lowest_latency(): """ - Test 1: num callbacks should NOT increase over time - -> Update to lowest latency - -> check current callbacks - -> sleep for 30s - -> check current callbacks - -> sleep for 30s - -> check current callbacks - -> update back to original routing-strategy + Same PROD invariant as test_check_num_callbacks, but after switching the + router to latency-based-routing. That switch is a *known, bounded* one-time + registration (it adds the latency strategy handler + Slack alerting); we + settle past it before baselining so only post-switch growth counts as a + leak. Also asserts the alerting count is stable. """ - from litellm._uuid import uuid - async with aiohttp.ClientSession() as session: await asyncio.sleep(30) original_routing_strategy = await get_current_routing_strategy(session=session) await config_update(session=session, routing_strategy="latency-based-routing") - await asyncio.sleep(30) + try: + # Absorb the deliberate one-time config/update registration step. + await asyncio.sleep(SETTLE_SECONDS) - num_callbacks_1, num_alerts_1, all_litellm_callbacks_1 = ( - await get_active_callbacks(session=session) - ) + samples, alerts = await _sample_callbacks( + session, NUM_SAMPLES, SAMPLE_INTERVAL_SECONDS + ) - await asyncio.sleep(30) - - num_callbacks_2, num_alerts_2, all_litellm_callbacks_2 = ( - await get_active_callbacks(session=session) - ) - - print( - "diff in callbacks all_litellm_callbacks_2 - all_litellm_callbacks_1 =", - set(all_litellm_callbacks_2) - set(all_litellm_callbacks_1), - ) - - assert abs(num_callbacks_1 - num_callbacks_2) <= 4 - - await asyncio.sleep(30) - - num_callbacks_3, num_alerts_3, all_litellm_callbacks_3 = ( - await get_active_callbacks(session=session) - ) - - print( - "diff in callbacks all_litellm_callbacks_3 - all_litellm_callbacks_2 =", - set(all_litellm_callbacks_3) - set(all_litellm_callbacks_2), - ) - - assert abs(num_callbacks_2 - num_callbacks_3) <= 4 - - assert num_alerts_1 == num_alerts_2 == num_alerts_3 - - await config_update(session=session, routing_strategy=original_routing_strategy) + leaks, samples = await _detect_leaks_confirmed(session, samples) + report = _format_report(samples, leaks) + print(report) + assert not leaks, f"Callback leak detected.\n{report}" + assert ( + len(set(alerts)) == 1 + ), f"alerting count changed across samples: {alerts}" + finally: + await config_update( + session=session, routing_strategy=original_routing_strategy + ) From 761ab1920977e1def32beea58f19d7429940be08 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 18 May 2026 18:00:18 -0700 Subject: [PATCH 10/11] fix(bedrock): sanitize batch metadata to prevent Pydantic ValidationError (#28202) * fix(bedrock): sanitize batch metadata to prevent Pydantic ValidationError Proxy guardrail hooks (Model Armor, OpenAI Moderations) and internal processing inject non-string values (dicts, floats) into the request metadata. When the Bedrock batch handler passes this metadata directly to LiteLLMBatch (which inherits OpenAI's Batch Pydantic model with metadata: Dict[str, str]), Pydantic raises a ValidationError. This causes the router retry loop to re-submit the same Bedrock job multiple times before ultimately failing. Add _get_openai_compatible_batch_metadata() that serializes non-string values to JSON strings via safe_dumps, skips None values and internal logging keys, ensuring the response object always validates. * test(bedrock): add tests for batch metadata sanitization Covers _get_openai_compatible_batch_metadata: string passthrough, dict/float serialization, None/internal key exclusion, and LiteLLMBatch compatibility. --------- Co-authored-by: Noah Nistler <60981020+noahnistler@users.noreply.github.com> --- .../llms/bedrock/batches/transformation.py | 26 +++- .../test_batch_metadata_sanitization.py | 119 ++++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 0602b1c2f62..620bc91732d 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -5,6 +5,7 @@ from typing import Any, Dict, List, Literal, Optional, Union, cast from httpx import Headers, Response +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -263,9 +264,32 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): cancelling_at=None, cancelled_at=None, request_counts=None, - metadata=original_request.get("metadata", {}), + metadata=self._get_openai_compatible_batch_metadata( + original_request.get("metadata", {}) + ), ) + @staticmethod + def _get_openai_compatible_batch_metadata(metadata: Any) -> Dict[str, str]: + """ + OpenAI Batch metadata only accepts string values. + """ + if not isinstance(metadata, dict): + return {} + + sanitized_metadata: Dict[str, str] = {} + for key, value in metadata.items(): + if key == "standard_logging_guardrail_information" or value is None: + continue + + str_key = str(key) + if isinstance(value, str): + sanitized_metadata[str_key] = value + else: + sanitized_metadata[str_key] = safe_dumps(value) + + return sanitized_metadata + def transform_retrieve_batch_request( self, batch_id: str, diff --git a/tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py b/tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py new file mode 100644 index 00000000000..8de47331614 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py @@ -0,0 +1,119 @@ +""" +Test that BedrockBatchesConfig._get_openai_compatible_batch_metadata +sanitizes non-string metadata values injected by proxy guardrail hooks. + +The OpenAI Batch Pydantic model requires metadata: Dict[str, str]. +Proxy hooks (Model Armor, OpenAI Moderations, queue time tracking) inject +dicts, floats, and other non-string values that cause a ValidationError +when constructing LiteLLMBatch. This test suite verifies the sanitization +layer prevents that. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + +class TestGetOpenaiCompatibleBatchMetadata: + """Tests for _get_openai_compatible_batch_metadata.""" + + def test_string_values_pass_through_unchanged(self): + metadata = {"user_key": "user_value", "run_id": "abc123"} + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert result == {"user_key": "user_value", "run_id": "abc123"} + + def test_dict_values_serialized_to_json_string(self): + metadata = { + "_model_armor_response": { + "sanitizationResult": {"filterMatchState": "MATCH_FOUND"} + } + } + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert "_model_armor_response" in result + assert isinstance(result["_model_armor_response"], str) + assert "MATCH_FOUND" in result["_model_armor_response"] + + def test_float_values_serialized_to_string(self): + metadata = {"queue_time_seconds": 0.5} + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert result == {"queue_time_seconds": "0.5"} + + def test_none_values_excluded(self): + metadata = {"key": "value", "empty": None} + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert "empty" not in result + assert result == {"key": "value"} + + def test_standard_logging_guardrail_information_excluded(self): + metadata = { + "standard_logging_guardrail_information": {"some": "logging_data"}, + "user_key": "keep_me", + } + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert "standard_logging_guardrail_information" not in result + assert result == {"user_key": "keep_me"} + + def test_non_dict_input_returns_empty_dict(self): + assert BedrockBatchesConfig._get_openai_compatible_batch_metadata(None) == {} + assert BedrockBatchesConfig._get_openai_compatible_batch_metadata("string") == {} + assert BedrockBatchesConfig._get_openai_compatible_batch_metadata(123) == {} + + def test_empty_dict_returns_empty_dict(self): + assert BedrockBatchesConfig._get_openai_compatible_batch_metadata({}) == {} + + def test_mixed_metadata_from_guardrails(self): + """Simulate real metadata contaminated by proxy guardrails.""" + metadata = { + "_model_armor_response": {"sanitizationResult": {"key": "val"}}, + "_model_armor_status": "success", + "_openai_moderation_response": {"id": "mod-123", "flagged": False}, + "queue_time_seconds": 1.23, + "headers": {"Authorization": "Bearer sk-xxx"}, + "standard_logging_guardrail_information": {"internal": True}, + "user_metadata_key": "user_value", + "none_field": None, + } + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + + # All values must be strings + for key, value in result.items(): + assert isinstance(value, str), f"metadata[{key!r}] is {type(value)}, not str" + + # Excluded keys + assert "standard_logging_guardrail_information" not in result + assert "none_field" not in result + + # Preserved keys + assert result["_model_armor_status"] == "success" + assert result["user_metadata_key"] == "user_value" + + def test_result_compatible_with_litellm_batch(self): + """Verify sanitized metadata can construct a LiteLLMBatch without error.""" + import time + + from litellm.types.utils import LiteLLMBatch + + metadata = { + "_model_armor_response": {"blocked": True}, + "queue_time_seconds": 0.05, + "user_key": "value", + } + sanitized = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + + # This would raise ValidationError before the fix + batch = LiteLLMBatch( + id="arn:aws:bedrock:us-east-1:123:model-invocation-job/test", + object="batch", + endpoint="/v1/chat/completions", + input_file_id="file-123", + completion_window="24h", + status="validating", + created_at=int(time.time()), + metadata=sanitized, + ) + assert batch.metadata == sanitized From 761c280a6e2401ae64c74e091542453655beb714 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 18 May 2026 18:14:13 -0700 Subject: [PATCH 11/11] fix(deepseek): use native /anthropic/v1/messages endpoint and sanitize tools (#28200) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(deepseek): route messages api through anthropic config Add a DeepSeek-specific Anthropic Messages config so deepseek/... models use the native messages endpoint and preserve thinking blocks. Strip Anthropic custom tool type markers that DeepSeek rejects while keeping hosted tool types intact. * fix(deepseek): normalize anthropic messages api base Handle OpenAI-style DeepSeek api_base values ending in /v1 or /v1/messages by stripping those suffixes before adding the /anthropic messages path. * chore(deepseek): format messages transformation * chore(deepseek): add test package markers * fix(deepseek): tighten anthropic url path check and fall back to DEEPSEEK_API_BASE Author: mateo-berri <277851410+mateo-berri@users.noreply.github.com> * fix(tests): normalize smart quotes in realtime guardrail refusal check gpt-realtime nondeterministically returns refusals with Unicode curly apostrophes (e.g. 'I’m sorry, but I can’t assist with that.'), but the safe_markers tuple in test_text_message_blocked_by_guardrail_no_ai_response only contains straight ASCII apostrophes. The substring match then fails even though the response is a clear refusal, flipping CI red. Normalize the AI text to ASCII quotes before the marker check so both straight and curly variants count as safe outcomes. * fix(deepseek): drop redundant anthropic v1/messages endswith check * fix(deepseek): strip /beta suffix in anthropic messages URL normalization Co-authored-by: Yassin Kortam --------- Co-authored-by: Felipe Rodrigues Gare Carnielli Co-authored-by: Cursor Agent Co-authored-by: Yassin Kortam --- .../llms/deepseek/messages/transformation.py | 133 ++++++++++++ litellm/utils.py | 6 + .../test_realtime_guardrails_openai.py | 9 +- tests/test_litellm/llms/deepseek/__init__.py | 0 .../llms/deepseek/messages/__init__.py | 0 ...pseek_anthropic_messages_transformation.py | 189 ++++++++++++++++++ 6 files changed, 336 insertions(+), 1 deletion(-) create mode 100644 litellm/llms/deepseek/messages/transformation.py create mode 100644 tests/test_litellm/llms/deepseek/__init__.py create mode 100644 tests/test_litellm/llms/deepseek/messages/__init__.py create mode 100644 tests/test_litellm/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py diff --git a/litellm/llms/deepseek/messages/transformation.py b/litellm/llms/deepseek/messages/transformation.py new file mode 100644 index 00000000000..ad60478960e --- /dev/null +++ b/litellm/llms/deepseek/messages/transformation.py @@ -0,0 +1,133 @@ +""" +DeepSeek Anthropic-compatible messages transformation config. +""" + +from typing import Any, Dict, List, Optional, Tuple + +import litellm +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams + + +class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): + """ + DeepSeek exposes an Anthropic-compatible Messages API at + https://api.deepseek.com/anthropic. + + It accepts the native Anthropic Messages conversation shape, including + thinking blocks in assistant history, but rejects Anthropic's explicit + custom-tool discriminator (`{"type": "custom"}`). + """ + + @property + def custom_llm_provider(self) -> Optional[str]: + return "deepseek" + + @staticmethod + def get_api_key(api_key: Optional[str] = None) -> Optional[str]: + return api_key or get_secret_str("DEEPSEEK_API_KEY") or litellm.api_key + + @staticmethod + def get_api_base(api_base: Optional[str] = None) -> str: + return ( + api_base + or get_secret_str("DEEPSEEK_ANTHROPIC_API_BASE") + or get_secret_str("DEEPSEEK_API_BASE") + or "https://api.deepseek.com/anthropic" + ) + + def validate_anthropic_messages_environment( + self, + headers: dict, + model: str, + messages: List[Any], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> Tuple[dict, Optional[str]]: + dynamic_api_key = self.get_api_key(api_key=api_key) + + if ( + "x-api-key" not in headers + and "authorization" not in headers + and dynamic_api_key is not None + ): + headers["x-api-key"] = dynamic_api_key + + if "anthropic-version" not in headers: + headers["anthropic-version"] = "2023-06-01" + if "content-type" not in headers: + headers["content-type"] = "application/json" + + headers = self._update_headers_with_anthropic_beta( + headers=headers, + optional_params=optional_params, + custom_llm_provider=self.custom_llm_provider or "deepseek", + ) + + return headers, api_base + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + base_url = self.get_api_base(api_base=api_base).rstrip("/") + + if base_url.endswith("/v1/messages") and "/anthropic/" in base_url: + return base_url + if base_url.endswith("/v1/messages"): + base_url = base_url[: -len("/v1/messages")] + if base_url.endswith("/v1"): + base_url = base_url[: -len("/v1")] + if base_url.endswith("/beta"): + base_url = base_url[: -len("/beta")] + + if not base_url.endswith("/anthropic") and "/anthropic/" not in base_url: + base_url = f"{base_url}/anthropic" + + return f"{base_url}/v1/messages" + + @staticmethod + def _sanitize_tools_for_deepseek(tools: Any) -> Any: + if not isinstance(tools, list): + return tools + + sanitized_tools = [] + for tool in tools: + if isinstance(tool, dict) and tool.get("type") == "custom": + sanitized_tool = dict(tool) + sanitized_tool.pop("type", None) + sanitized_tools.append(sanitized_tool) + else: + sanitized_tools.append(tool) + return sanitized_tools + + def transform_anthropic_messages_request( + self, + model: str, + messages: List[Dict], + anthropic_messages_optional_request_params: Dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Dict: + anthropic_messages_request = super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + if "tools" in anthropic_messages_request: + anthropic_messages_request["tools"] = self._sanitize_tools_for_deepseek( + anthropic_messages_request["tools"] + ) + return anthropic_messages_request diff --git a/litellm/utils.py b/litellm/utils.py index 54cea313b0c..001c89fee4c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8539,6 +8539,12 @@ class ProviderConfigManager: ) return MinimaxMessagesConfig() + elif litellm.LlmProviders.DEEPSEEK == provider: + from litellm.llms.deepseek.messages.transformation import ( + DeepSeekAnthropicMessagesConfig, + ) + + return DeepSeekAnthropicMessagesConfig() return None @staticmethod diff --git a/tests/llm_translation/realtime/test_realtime_guardrails_openai.py b/tests/llm_translation/realtime/test_realtime_guardrails_openai.py index ec9d73e2d60..50cedba2ac0 100644 --- a/tests/llm_translation/realtime/test_realtime_guardrails_openai.py +++ b/tests/llm_translation/realtime/test_realtime_guardrails_openai.py @@ -232,8 +232,15 @@ async def test_text_message_blocked_by_guardrail_no_ai_response(): assert ( BLOCKED_PHRASE not in real_ai_text ), f"Blocked phrase leaked into AI response: {real_ai_text!r}" + normalized_ai_text = ( + real_ai_text.lower() + .replace("\u2019", "'") + .replace("\u2018", "'") + .replace("\u201c", '"') + .replace("\u201d", '"') + ) assert any( - marker in real_ai_text.lower() for marker in safe_markers + marker in normalized_ai_text for marker in safe_markers ), f"AI responded with non-guardrail content even though message was blocked: {real_ai_text!r}" finally: diff --git a/tests/test_litellm/llms/deepseek/__init__.py b/tests/test_litellm/llms/deepseek/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/deepseek/messages/__init__.py b/tests/test_litellm/llms/deepseek/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py b/tests/test_litellm/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py new file mode 100644 index 00000000000..7c5f0483ded --- /dev/null +++ b/tests/test_litellm/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py @@ -0,0 +1,189 @@ +import litellm +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) +from litellm.llms.deepseek.messages.transformation import ( + DeepSeekAnthropicMessagesConfig, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager + + +def test_deepseek_provider_uses_anthropic_messages_config(): + config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="deepseek-v4-pro", + provider=litellm.LlmProviders.DEEPSEEK, + ) + + assert isinstance(config, DeepSeekAnthropicMessagesConfig) + assert config.custom_llm_provider == "deepseek" + + +def test_deepseek_anthropic_messages_config_defaults(): + config = DeepSeekAnthropicMessagesConfig() + + assert config.custom_llm_provider == "deepseek" + assert config.get_api_base() == "https://api.deepseek.com/anthropic" + + +def test_anthropic_provider_keeps_default_config_for_deepseek_named_model(): + config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="deepseek-v4-pro", + provider=litellm.LlmProviders.ANTHROPIC, + ) + + assert isinstance(config, AnthropicMessagesConfig) + assert not isinstance(config, DeepSeekAnthropicMessagesConfig) + + +def test_deepseek_anthropic_messages_url_defaults_to_anthropic_endpoint(): + config = DeepSeekAnthropicMessagesConfig() + + assert ( + config.get_complete_url( + api_base=None, + api_key=None, + model="deepseek-v4-pro", + optional_params={}, + litellm_params={}, + ) + == "https://api.deepseek.com/anthropic/v1/messages" + ) + assert ( + config.get_complete_url( + api_base="https://api.deepseek.com/anthropic/v1", + api_key=None, + model="deepseek-v4-pro", + optional_params={}, + litellm_params={}, + ) + == "https://api.deepseek.com/anthropic/v1/messages" + ) + assert ( + config.get_complete_url( + api_base="https://api.deepseek.com/anthropic", + api_key=None, + model="deepseek-v4-pro", + optional_params={}, + litellm_params={}, + ) + == "https://api.deepseek.com/anthropic/v1/messages" + ) + assert ( + config.get_complete_url( + api_base="https://api.deepseek.com", + api_key=None, + model="deepseek-v4-pro", + optional_params={}, + litellm_params={}, + ) + == "https://api.deepseek.com/anthropic/v1/messages" + ) + assert ( + config.get_complete_url( + api_base="https://api.deepseek.com/v1", + api_key=None, + model="deepseek-v4-pro", + optional_params={}, + litellm_params={}, + ) + == "https://api.deepseek.com/anthropic/v1/messages" + ) + assert ( + config.get_complete_url( + api_base="https://api.deepseek.com/v1/messages", + api_key=None, + model="deepseek-v4-pro", + optional_params={}, + litellm_params={}, + ) + == "https://api.deepseek.com/anthropic/v1/messages" + ) + + +def test_deepseek_anthropic_messages_headers_use_deepseek_key(): + config = DeepSeekAnthropicMessagesConfig() + + headers, api_base = config.validate_anthropic_messages_environment( + headers={}, + model="deepseek-v4-pro", + messages=[], + optional_params={}, + litellm_params={}, + api_key="sk-deepseek", + api_base="https://example.test/anthropic", + ) + + assert api_base == "https://example.test/anthropic" + assert headers["x-api-key"] == "sk-deepseek" + assert headers["anthropic-version"] == "2023-06-01" + assert headers["content-type"] == "application/json" + + +def test_deepseek_anthropic_messages_preserves_thinking_and_sanitizes_custom_tools(): + config = DeepSeekAnthropicMessagesConfig() + messages = [ + { + "role": "user", + "content": "Use the tool.", + }, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "I should call the tool.", + "signature": "sig", + }, + { + "type": "tool_use", + "id": "toolu_123", + "name": "get_weather", + "input": {"city": "Sao Paulo"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_123", + "content": "Sunny", + } + ], + }, + ] + + request = config.transform_anthropic_messages_request( + model="deepseek-v4-pro", + messages=messages, + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "tools": [ + { + "type": "custom", + "name": "get_weather", + "description": "Get weather", + "input_schema": {"type": "object"}, + }, + { + "type": "web_search_20260209", + "name": "web_search", + "max_uses": 1, + }, + ], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert request["messages"] == messages + assert request["thinking"] == {"type": "enabled", "budget_tokens": 1024} + assert request["tools"][0] == { + "name": "get_weather", + "description": "Get weather", + "input_schema": {"type": "object"}, + } + assert request["tools"][1]["type"] == "web_search_20260209"