From e16d79d635123db7d4c956b57536ebc43cd00e27 Mon Sep 17 00:00:00 2001 From: Jodok Batlogg Date: Tue, 27 Oct 2026 03:24:09 +0100 Subject: [PATCH 01/12] feat(decisions): support LM Studio decision models (#48448) * feat(decisions): support LM Studio decision models Co-Authored-By: Codex * fix(ci): resolve OSV dependency findings Upgrade the extension MSAL dependency and document expiring exceptions for unused Windows filelock APIs while Python 3.10 remains supported. Co-Authored-By: Codex * Clean up ignored vulnerabilities in osv-scanner.tomlchore: remove unrelated dependency and scanner changes Removed ignored vulnerabilities related to Windows issues and Python version compatibility. * chore: remove unrelated dependency and scanner changes * chore: remove unrelated dependency and scanner changes --------- Co-authored-by: Codex Co-authored-by: ishaan-j-berri --- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 2 + .../lm_studio/decisions/transformation.py | 16 ++ litellm/utils.py | 2 + tests/unit/decisions/test_main.py | 140 +++++++++++++++++- .../decisions_endpoints/test_endpoints.py | 50 +++++++ 6 files changed, 209 insertions(+), 4 deletions(-) create mode 100644 litellm/llms/lm_studio/decisions/transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 117ba9604e9..c2c13a12eea 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1729,6 +1729,9 @@ if TYPE_CHECKING: from .llms.hosted_vllm.decisions.transformation import ( HostedVLLMDecisionsConfig as HostedVLLMDecisionsConfig, ) + from .llms.lm_studio.decisions.transformation import ( + LMStudioDecisionsConfig as LMStudioDecisionsConfig, + ) from .llms.laya.decisions.transformation import ( LayaDecisionsConfig as LayaDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index a56004d1274..6aecfd2ff8c 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -170,6 +170,7 @@ LLM_CONFIG_NAMES: Final = ( "StrandsDeciderDecisionsConfig", "DatabricksDecisionsConfig", "HostedVLLMDecisionsConfig", + "LMStudioDecisionsConfig", "LayaDecisionsConfig", "BespokeDecisionsConfig", "OpenAIDecisionsConfig", @@ -746,6 +747,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.hosted_vllm.decisions.transformation", "HostedVLLMDecisionsConfig", ), + "LMStudioDecisionsConfig": (".llms.lm_studio.decisions.transformation", "LMStudioDecisionsConfig"), "LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"), "BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"), "OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"), diff --git a/litellm/llms/lm_studio/decisions/transformation.py b/litellm/llms/lm_studio/decisions/transformation.py new file mode 100644 index 00000000000..84e5ff14443 --- /dev/null +++ b/litellm/llms/lm_studio/decisions/transformation.py @@ -0,0 +1,16 @@ +from typing import Final + +from litellm.constants import PLACEHOLDER_API_KEY +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + + +class LMStudioDecisionsConfig(BaseDecisionsConfig): + api_key_env = ("LM_STUDIO_API_KEY",) + api_base_env = ("LM_STUDIO_API_BASE",) + api_key_required = False + supports_images = True + inlines_remote_images = True + + def resolve_api_key(self, api_key: str | None) -> str | None: + resolved: Final = super().resolve_api_key(api_key) + return None if resolved == PLACEHOLDER_API_KEY else resolved diff --git a/litellm/utils.py b/litellm/utils.py index 1aed15ba83e..fdfad7fc665 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9610,6 +9610,8 @@ class ProviderConfigManager: return litellm.DatabricksDecisionsConfig() if provider == LlmProviders.HOSTED_VLLM: return litellm.HostedVLLMDecisionsConfig() + if provider == LlmProviders.LM_STUDIO: + return litellm.LMStudioDecisionsConfig() if provider == LlmProviders.LAYA: return litellm.LayaDecisionsConfig() if provider == LlmProviders.BESPOKE: diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 5b32008b279..d13d212b4de 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -38,6 +38,7 @@ from litellm.types.decisions import ( from litellm.types.router import ModelInfo from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +from tests.unit.conftest import AsyncOnlyImageFetch _QUESTIONS: Final[Mapping[str, object]] = MappingProxyType( { @@ -682,7 +683,9 @@ _PPLX_USAGE: Final = {"model": "pplx-decider-v1-27b", "usage": {"input_tokens": ("reply", "detail"), ( pytest.param(_PPLX_USAGE, "DecisionsResponse answers: Field required", id="usage"), - pytest.param({"error": "Rate limit reached"}, "DecisionsResponse answers: Field required", id="rate_limit_text"), + pytest.param( + {"error": "Rate limit reached"}, "DecisionsResponse answers: Field required", id="rate_limit_text" + ), pytest.param({"error": "Request timed out"}, "DecisionsResponse answers: Field required", id="timeout_text"), pytest.param( {**_PPLX_USAGE, "answers": {"rate_limit": None}}, @@ -694,7 +697,9 @@ _PPLX_USAGE: Final = {"model": "pplx-decider-v1-27b", "usage": {"input_tokens": "DecisionsResponse answers.is_defect", id="answer_type_timeout_text", ), - pytest.param(["Rate limit reached"], "DecisionsResponse Input should be a valid dictionary", id="not_an_object"), + pytest.param( + ["Rate limit reached"], "DecisionsResponse Input should be a valid dictionary", id="not_an_object" + ), ), ) async def test_a_malformed_upstream_reply_is_a_server_error_whatever_its_text( @@ -1525,7 +1530,12 @@ async def test_databricks_router_deployment_sends_its_own_key_to_its_own_base( _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset( - {LlmProviders.OPENAI.value, LlmProviders.OPENROUTER.value, LlmProviders.HOSTED_VLLM.value} + { + LlmProviders.OPENAI.value, + LlmProviders.OPENROUTER.value, + LlmProviders.HOSTED_VLLM.value, + LlmProviders.LM_STUDIO.value, + } ) _DEDICATED_DECIDER_PROVIDERS: Final[tuple[str, ...]] = tuple( sorted( @@ -3791,7 +3801,12 @@ async def test_a_returned_model_registered_without_prices_is_billed_at_the_reque monkeypatch.setitem( litellm.model_cost, "typesafe/jev-test-alias", - {"litellm_provider": "typesafe", "mode": "evaluation", "input_cost_per_token": 0.001, "output_cost_per_token": 0}, + { + "litellm_provider": "typesafe", + "mode": "evaluation", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0, + }, ) monkeypatch.setitem(litellm.model_cost, "typesafe/jev-test-unpriced", unpriced_row) respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json={**_RESPONSE, "model": "jev-test-unpriced"}) @@ -3893,3 +3908,120 @@ async def test_native_decisions_distinguishes_missing_usage_from_reported_zero( expected = litellm.completion_cost(completion_response=response) assert response._hidden_params["response_cost"] == pytest.approx(expected) assert (expected == 0) == (usage["input_tokens"] == usage["output_tokens"] == 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", (True, False), ids=("asystemone", "systemone")) +@pytest.mark.parametrize("api_base", ("http://lmstudio.local:1234", "http://lmstudio.local:1234/v1")) +async def test_lm_studio_systemone_preserves_questions_and_omits_placeholder_auth( + use_async: bool, api_base: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("LM_STUDIO_API_KEY", raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-openai-key") + route: Final = respx_mock.post("http://lmstudio.local:1234/v1/systemone").respond(json=_RESPONSE) + kwargs: Final = { + "model": "lm_studio/local-decision", + "state": "review", + "questions": _QUESTIONS, + "api_base": api_base, + } + + response: Final = await litellm.asystemone(**kwargs) if use_async else litellm.systemone(**kwargs) + + assert json.loads(route.calls.last.request.content) == { + "model": "local-decision", + "state": "review", + "questions": dict(_QUESTIONS), + } + assert "authorization" not in route.calls.last.request.headers + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert response.answers["sentiment"].choice == "positive" + assert response.answers["severity"].score == 1 + assert response.usage == DecisionsUsage(input_tokens=_INPUT_TOKENS, output_tokens=_OUTPUT_TOKENS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", (True, False), ids=("adecisions", "decisions")) +async def test_lm_studio_decisions_translates_openai_questions_and_answers( + use_async: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("LM_STUDIO_API_KEY", raising=False) + route: Final = respx_mock.post("http://lmstudio.local:1234/v1/systemone").respond(json=_RESPONSE) + kwargs: Final = { + "model": "lm_studio/local-decision", + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + "api_base": "http://lmstudio.local:1234/v1", + } + + response: Final = await litellm.adecisions(**kwargs) if use_async else litellm.decisions(**kwargs) + + assert json.loads(route.calls.last.request.content) == { + "model": "local-decision", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert response.answers[0].type == "predicate" + assert response.answers[0].name == "is_defect" + assert response.answers[0].probability == 0.9 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("explicit_key", (None, "deployment-key")) +async def test_lm_studio_decisions_resolves_environment_connection_and_explicit_key_override( + explicit_key: str | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setenv("LM_STUDIO_API_BASE", "http://lmstudio.local:1234/v1") + monkeypatch.setenv("LM_STUDIO_API_KEY", "environment-key") + route: Final = respx_mock.post("http://lmstudio.local:1234/v1/systemone").respond(json=_RESPONSE) + + response: Final = await litellm.asystemone( + model="lm_studio/local-decision", state="review", questions=_QUESTIONS, api_key=explicit_key + ) + + assert route.calls.last.request.headers["authorization"] == f"Bearer {explicit_key or 'environment-key'}" + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + + +@pytest.mark.asyncio +async def test_lm_studio_decisions_requires_a_server_before_http( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("LM_STUDIO_API_BASE", raising=False) + with pytest.raises(litellm.BadRequestError, match="api_base is required"): + await litellm.asystemone(model="lm_studio/local-decision", state="review", questions=_QUESTIONS) + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("remote_image", (False, True), ids=("inline", "remote")) +async def test_lm_studio_decisions_sends_inline_images_to_systemone( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + remote_image: bool, + async_only_image_fetch: AsyncOnlyImageFetch, +) -> None: + monkeypatch.delenv("LM_STUDIO_API_KEY", raising=False) + image: Final = "https://images.example/input.png" if remote_image else "data:image/png;base64,AA==" + expected_image: Final = async_only_image_fetch.data_url if remote_image else image + route: Final = respx_mock.post("http://lmstudio.local:1234/v1/systemone").respond(json=_RESPONSE) + + response: Final = await litellm.adecisions( + model="lm_studio/local-decision", + input=[ + { + "role": "user", + "content": [{"type": "input_text", "text": "review"}, {"type": "input_image", "image_url": image}], + } + ], + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + api_base="http://lmstudio.local:1234/v1", + ) + + assert json.loads(route.calls.last.request.content) == { + "model": "local-decision", + "state": "review", + "images": [expected_image], + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert response.answers[0].probability == 0.9 diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index 4384507e841..a0a79f9b7cf 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -1784,3 +1784,53 @@ def test_a_whole_number_request_size_limit_is_named_as_configured_in_the_413( assert response.status_code == 413, response.text assert response.json() == {"error": "Request size is too large. Max size is 1 MB"} assert not upstream.called + + +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/v1/decisions")) +@pytest.mark.parametrize("api_key", (None, "deployment-key")) +def test_proxy_routes_lm_studio_decisions_with_optional_authentication( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + endpoint: str, + api_key: str | None, +) -> None: + monkeypatch.delenv("LM_STUDIO_API_KEY", raising=False) + deployment_router: Final = litellm.Router( + model_list=[ + { + "model_name": "local-decision", + "litellm_params": { + "model": "lm_studio/served-decision", + "api_base": "http://lmstudio.local:1234/v1", + "api_key": api_key, + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", deployment_router) + upstream: Final = respx_mock.post("http://lmstudio.local:1234/v1/systemone").respond(json=_RESPONSE) + fields: Final = ( + {"state": "review", "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}} + if endpoint.endswith("systemone") + else { + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + } + ) + + response: Final = client.post(endpoint, json={"model": "local-decision", **fields}) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == { + "model": "served-decision", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert upstream.calls.last.request.headers.get("authorization") == (f"Bearer {api_key}" if api_key else None) + expected: Final = ( + _RESPONSE["answers"] + if endpoint.endswith("systemone") + else [{"type": "predicate", "name": "is_defect", "probability": 0.9}] + ) + assert response.json()["answers"] == expected From a0d0ce1aa5923a01c37cd9fdb57afc70c0dca521 Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Mon, 26 Oct 2026 19:31:31 -0700 Subject: [PATCH 02/12] fix(token-counter): safely count inline auto image tiles (#48422) (#48526) --- litellm/litellm_core_utils/token_counter.py | 61 +++++++++++++---- .../litellm_core_utils/test_token_counter.py | 67 +++++++++++++++++++ 2 files changed, 115 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 90e79d4d853..ee5f70fb394 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -261,6 +261,10 @@ def get_image_dimensions( _header, encoded = data.split(",", 1) img_data = base64.b64decode(encoded) + return _image_dimensions_from_bytes(img_data) + + +def _image_dimensions_from_bytes(img_data: bytes) -> tuple[int, int]: img_type: Final = get_image_type(img_data) if img_type == "png": @@ -307,8 +311,37 @@ def get_image_dimensions( return DEFAULT_IMAGE_WIDTH, DEFAULT_IMAGE_HEIGHT +def _auto_image_dimensions(data: str) -> tuple[int, int] | None: + """Read bounded inline image data without fetching caller-supplied URLs.""" + if not data.startswith("data:image/"): + return None + max_bytes: Final = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) + max_encoded_length: Final = 4 * ((max_bytes + 2) // 3) + # Bound allocations before splitting or decoding the data URL, including its header. + if max_bytes <= 0 or len(data) > max_encoded_length + 128: + return None + header, separator, encoded = data.partition(",") + if not separator or len(header) > 128 or not header.endswith(";base64") or len(encoded) > max_encoded_length: + return None + try: + image_data: Final = base64.b64decode(encoded, validate=True) + if len(image_data) > max_bytes or get_image_type(image_data) not in ( + "png", + "gif", + "jpeg", + "webp", + ): + return None + width, height = _image_dimensions_from_bytes(image_data) + except (ValueError, TypeError, struct.error): + return None + if width <= 0 or height <= 0: + return None + return width, height + + def calculate_img_tokens( - data, + data: str, mode: Literal["low", "high", "auto"] = "auto", base_tokens: int = 85, # openai default - https://openai.com/pricing use_default_image_token_count: bool = False, @@ -318,7 +351,8 @@ def calculate_img_tokens( Args: data (str): The URL or base64 encoded string of the image. - mode (Literal["low", "high", "auto"]): The mode to use for calculating the number of tokens. + mode (Literal["low", "high", "auto"]): Auto uses high-detail tiles for readable inline images; + URLs and unreadable images keep the base estimate without network I/O. base_tokens (int): The base number of tokens for an image. use_default_image_token_count (bool): When True, will NOT make a GET request to the image URL and instead return the default image dimensions. @@ -328,18 +362,19 @@ def calculate_img_tokens( if use_default_image_token_count: verbose_logger.debug("Using default image token count: %s", DEFAULT_IMAGE_TOKEN_COUNT) return DEFAULT_IMAGE_TOKEN_COUNT - if mode == "low" or mode == "auto": + if mode == "low": return base_tokens - elif mode == "high": - # Run the async function using the helper - width, height = get_image_dimensions( - data=data, - ) - resized_width, resized_height = resize_image_high_res(width=width, height=height) - tiles_needed_high_res = calculate_tiles_needed(resized_width=resized_width, resized_height=resized_height) - tile_tokens: Final = (base_tokens * 2) * tiles_needed_high_res - total_tokens: Final = base_tokens + tile_tokens - return total_tokens + if mode not in ("auto", "high"): + return None + dimensions: Final = _auto_image_dimensions(data) if mode == "auto" else get_image_dimensions(data=data) + if dimensions is None: + return base_tokens + width, height = dimensions + resized_width, resized_height = resize_image_high_res(width=width, height=height) + tiles_needed_high_res = calculate_tiles_needed(resized_width=resized_width, resized_height=resized_height) + tile_tokens: Final = (base_tokens * 2) * tiles_needed_high_res + total_tokens: Final = base_tokens + tile_tokens + return total_tokens TokenCounterFunction = Callable[[str], int] diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 88af656ad2e..c452e12f2aa 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -3,6 +3,7 @@ import asyncio import base64 import importlib +import io import json import os import subprocess @@ -19,6 +20,7 @@ from unittest.mock import MagicMock, sentinel import anyio.to_thread import pytest import tiktoken +from PIL import Image from tokenizers import Regex, Tokenizer, models, pre_tokenizers from unittest.mock import AsyncMock, patch @@ -2015,6 +2017,71 @@ def _png_data_url(width: int, height: int) -> str: return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() +@pytest.mark.parametrize( + ("width", "height", "expected"), [(256, 256, 262), (512, 512, 262), (768, 768, 772), (1024, 1024, 772)] +) +@pytest.mark.parametrize("detail", [None, "auto", "high"]) +def test_inline_auto_images_count_high_detail_tiles(width: int, height: int, expected: int, detail: str | None) -> None: + image_url = {"url": _png_data_url(width, height)} + if detail is not None: + image_url["detail"] = detail + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": image_url}]}] + + assert token_counter_new(model="gpt-4o", messages=messages) == expected + + +@pytest.mark.parametrize("image_format", ["PNG", "JPEG", "GIF", "WEBP"]) +def test_inline_auto_images_read_supported_formats(image_format: str) -> None: + output = io.BytesIO() + Image.new("RGB", (1024, 1024)).save(output, format=image_format) + url = f"data:image/{image_format.lower()};base64," + base64.b64encode(output.getvalue()).decode() + + assert calculate_img_tokens(url) == 765 + + +@pytest.mark.parametrize( + "data", + [ + "https://example.com/image.png", + "http://127.0.0.1/image.png", + "http://169.254.169.254/latest/meta-data/", + "file:///etc/passwd", + "data:application/pdf;base64,JVBERi0xLjQK", + "data:image/png;base64,%%%", + "data:image/png;base64,iVBORw0KGgo=", + "data:image/jpeg;base64,/9j/4AAQ", + "data:image/png,not-base64", + "", + _png_data_url(0, 1024), + _png_data_url(1024, 0), + ], +) +def test_auto_image_fallback_does_not_fetch_urls_or_raise(data: str) -> None: + # The unit suite disables sockets; a fetch must also not be attempted and swallowed. + with patch("httpx.Client.get", side_effect=AssertionError("auto counting must not fetch URLs")) as get: + assert calculate_img_tokens(data) == 85 + get.assert_not_called() + + +def test_auto_image_decode_respects_size_limit(monkeypatch: pytest.MonkeyPatch) -> None: + counter = importlib.import_module("litellm.litellm_core_utils.token_counter") + monkeypatch.setattr(counter, "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 24 / (1024 * 1024)) + + with patch("base64.b64decode", side_effect=AssertionError("oversized input must not be decoded")) as decode: + assert calculate_img_tokens(_png_data_url(1024, 1024)) == 85 + decode.assert_not_called() + + +def test_low_and_default_image_counts_skip_inline_decoding() -> None: + with patch("base64.b64decode", side_effect=AssertionError("fixed counts must not decode")) as decode: + assert calculate_img_tokens(_png_data_url(1024, 1024), mode="low") == 85 + assert ( + calculate_img_tokens(_png_data_url(1024, 1024), use_default_image_token_count=True) + == litellm.constants.DEFAULT_IMAGE_TOKEN_COUNT + ) + decode.assert_not_called() + + @pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)]) def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None: assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound() From cd1ccb8d6885b318748922b29d832884af0adf52 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 26 Oct 2026 19:37:14 -0700 Subject: [PATCH 03/12] fix(chatgpt): forward client credentials for image generation (#48522) * fix(chatgpt): forward client credentials for image generation * fix(chatgpt): validate image headers at the typed boundary * refactor(chatgpt): use read-only image header types * fix(chatgpt): preserve image header credentials and workspace --- litellm/images/main.py | 74 ++++++++++++++++++++++++++++++++-- tests/unit/images/test_main.py | 55 +++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 4 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 29bac3ffe05..9c88e737a40 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars import importlib -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from functools import partial from typing import TYPE_CHECKING, Final, Literal, cast, overload @@ -85,6 +85,47 @@ def _get_ImageEditRequestUtils() -> type["ImageEditRequestUtils"]: return _ImageEditRequestUtils_cache +def _chatgpt_image_params(kwargs: Mapping[str, object]) -> GenericLiteLLMParams: + from litellm.responses.utils import ResponsesAPIRequestUtils + + headers: Final = ResponsesAPIRequestUtils.merge_client_forwarded_headers( + TypeAdapter[dict[str, object] | None](dict[str, object] | None).validate_python(kwargs.get("extra_headers")), + TypeAdapter[dict[str, str] | None](dict[str, str] | None).validate_python(kwargs.get("headers")), + ) + return GenericLiteLLMParams.model_validate({**kwargs, "extra_headers": headers}) + + +def _chatgpt_image_headers(model: str, params: GenericLiteLLMParams, headers: Mapping[str, str]) -> Mapping[str, str]: + from litellm.llms.chatgpt.common_utils import get_chatgpt_request_credential + from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig + + credential: Final = get_chatgpt_request_credential(params) + credential_headers: Final = ( + { + "Authorization": f"Bearer {credential.access_token}", + **({"ChatGPT-Account-Id": credential.account_id} if credential.account_id else {}), + } + if credential is not None + else {} + ) + explicit_account_id: Final = next( + (value for name, value in headers.items() if name.lower() == "chatgpt-account-id"), None + ) + auth_headers: Final = ChatGPTResponsesAPIConfig().validate_environment( + headers=credential_headers, model=model, litellm_params=params + ) + return { + **{k: v for k, v in auth_headers.items() if k.lower() not in ("content-type", "accept")}, + **{ + k: v + for k, v in headers.items() + if k.lower() not in ("authorization", "chatgpt-account-id", "cookie", "x-litellm-api-key") + }, + **({"ChatGPT-Account-Id": explicit_account_id} if credential is None and explicit_account_id else {}), + "Accept": "application/json", + } + + ##### Image Generation ####################### @client async def aimage_generation(*args, **kwargs) -> ImageResponse: @@ -121,6 +162,12 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: model=model, custom_llm_provider=provider_hints.get("custom_llm_provider"), api_base=provider_hints.get("api_base"), + litellm_params=( + _chatgpt_image_params(TypeAdapter(dict[str, object]).validate_python(kwargs)) + if provider_hints.get("custom_llm_provider") == "chatgpt" + or (isinstance(model, str) and model.startswith("chatgpt/")) + else None + ), ) # Await normally @@ -240,12 +287,18 @@ def image_generation( if extra_headers is not None: headers.update(extra_headers) model_response: ImageResponse = litellm.utils.ImageResponse() + chatgpt_params: Final = ( + _chatgpt_image_params(TypeAdapter(dict[str, object]).validate_python({**kwargs, "api_key": api_key})) + if custom_llm_provider == "chatgpt" or (model is not None and model.startswith("chatgpt/")) + else None + ) dynamic_api_key: str | None = None if model is not None or custom_llm_provider is not None: model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, + litellm_params=chatgpt_params, ) else: model = "dall-e-2" @@ -268,7 +321,10 @@ def image_generation( "size", "style", ] - non_default_params: Final = filter_out_litellm_params(kwargs, excluding=openai_params) + non_default_params: Final = filter_out_litellm_params( + kwargs, + excluding=[*openai_params, "extra_headers", "headers"] if chatgpt_params is not None else openai_params, + ) image_generation_config: BaseImageGenerationConfig | None = None if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values(): @@ -490,7 +546,7 @@ def image_generation( or custom_llm_provider == LlmProviders.LITELLM_PROXY.value or custom_llm_provider in litellm.openai_compatible_providers ): - if extra_headers is not None: + if extra_headers is not None and chatgpt_params is None: optional_params["extra_headers"] = extra_headers # Forward OpenAI organization if present (set by proxy pre-call utils) organization: Final[str | None] = kwargs.get("organization", None) @@ -506,7 +562,17 @@ def image_generation( organization=organization, aimg_generation=aimg_generation, client=client, - headers=headers, + headers=( + dict( + _chatgpt_image_headers( + model, + chatgpt_params.model_copy(update={"api_key": dynamic_api_key or api_key}), + TypeAdapter(dict[str, str]).validate_python(headers), + ) + ) + if chatgpt_params is not None + else headers + ), litellm_params=litellm_params_dict, ) elif custom_llm_provider == "bedrock": diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index ff43e892eab..38af99ff505 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -1,6 +1,8 @@ import asyncio +import base64 import json from datetime import datetime +from pathlib import Path from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -405,3 +407,56 @@ async def test_azure_image_generation_request_body(): call_args = mock_post.call_args request_json = call_args.kwargs.get("json", {}) assert request_json == expected_body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("transport", ["headers", "extra_headers", "scoped", "api_key"]) +@pytest.mark.parametrize("saved_key", [None, "fake-saved-gateway-key"]) +async def test_chatgpt_images_forward_client_login_and_workspace( + asynchronous: bool, + transport: str, + saved_key: str | None, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + claims: Final = {"https://api.openai.com/auth": {"chatgpt_account_id": "client-account"}} + payload: Final = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=") + token: Final = f"test.{payload}.test" + headers: Final = {"aUtHoRiZaTiOn": f"Bearer {token}", "ChatGPT-Account-Id": "selected-workspace"} + credentials: Final = ( + {"provider_specific_header": {"custom_llm_provider": "chatgpt", "extra_headers": headers}} + if transport == "scoped" + else {"api_key": token, "extra_headers": {"cHaTgPt-AcCoUnT-Id": "selected-workspace"}} + if transport == "api_key" + else {transport: {**headers, "cookie": "fake-cookie", "x-litellm-api-key": "fake-gateway-key"}} + ) + route: Final = respx_mock.post("https://chatgpt.com/backend-api/codex/images/generations").mock( + return_value=httpx.Response(200, json={"created": 1700000000, "data": [{"b64_json": "aW1n"}]}) + ) + params: Final = { + "model": "chatgpt/gpt-image-2", + "prompt": "A blue circle", + "api_key": saved_key, + **( + { + "client": AsyncOpenAI( + api_key="fake", base_url="https://chatgpt.com/backend-api/codex", http_client=httpx.AsyncClient() + ) + } + if asynchronous + else {} + ), + **credentials, + } + response: Final = await litellm.aimage_generation(**params) if asynchronous else litellm.image_generation(**params) + assert response.data[0].b64_json == "aW1n" + request: Final = route.calls.last.request + assert request.headers["authorization"] == f"Bearer {token}" + assert request.headers["chatgpt-account-id"] == "selected-workspace" + assert request.headers.get_list("content-type") == ["application/json"] + assert "cookie" not in request.headers + assert "x-litellm-api-key" not in request.headers + assert json.loads(request.content) == {"model": "gpt-image-2", "prompt": "A blue circle"} From cc09b52616de5fa744095cc73e01ba764c951586 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 22:42:38 -0400 Subject: [PATCH 04/12] fix(proxy): include Claude web search fees in streamed usage cost (#48528) --- litellm/proxy/common_request_processing.py | 1 + .../proxy/test_common_request_processing.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 61cf4d696a1..8e681c615ec 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4687,6 +4687,7 @@ class ProxyBaseLLMRequestProcessing: total_tokens=total_tokens, completion_tokens_details=usage.get("completion_tokens_details"), prompt_tokens_details=usage.get("prompt_tokens_details"), + server_tool_use=usage.get("server_tool_use"), ) @staticmethod diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 82dd3479321..561577094df 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -10175,6 +10175,32 @@ class TestInjectCostIntoUsageDict: class TestProcessChunkWithCostInjection: + @pytest.mark.parametrize("model", ["claude-opus-5-5", "vertex_ai/claude-opus-5-5"]) + @pytest.mark.parametrize("search_requests", [None, "null", 0, 1, 2]) + def test_claude_search_stream_cost_matches_complete_usage( + self, monkeypatch: pytest.MonkeyPatch, model: str, search_requests: int | Literal["null"] | None + ) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + usage: Final = litellm.Usage( + prompt_tokens=10760, + completion_tokens=399, + total_tokens=11159, + server_tool_use={"web_search_requests": search_requests} if isinstance(search_requests, int) else None, + ) + wire_usage: Final = usage.model_dump() | ({"server_tool_use": None} if search_requests == "null" else {}) + event: Final = {"object": "chat.completion.chunk", "choices": [], "usage": wire_usage} + chunk: Final = f"data: {json.dumps(event)}\n\n".encode() + + result: Final = ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection(chunk, model) + + assert isinstance(result, bytes) + payload: Final = json.loads(result.decode().removeprefix("data: ")) + expected_cost: Final = litellm.completion_cost( + model=model, completion_response=litellm.ModelResponse(usage=usage) + ) + assert payload["usage"]["cost"] == pytest.approx(expected_cost) + assert "cost" not in event["usage"] + def test_complete_usage_frame_chunk_is_injected(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) chunk = ( From f2b59be6ea8a8fec61dace1a20bcde2b5d8af780 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 22:43:23 -0400 Subject: [PATCH 05/12] fix(vertex): keep Messages authentication off the event loop (#48523) --- .../transformation.py | 92 +++++++++++++++---- ...erimental_pass_through_messages_handler.py | 22 ++++- ...artner_models_anthropic_messages_config.py | 80 ++++++++++++++++ 3 files changed, 172 insertions(+), 22 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index cfd5b5d3980..eefeda190ca 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -26,34 +26,92 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert def should_strip_billing_metadata(self) -> bool: return True + @staticmethod + def _vertex_auth_parameters( + litellm_params: dict[str, object], # mutable-ok: delegates to the existing Vertex parameter readers + ) -> tuple[str | None, str | None, str | None]: + return ( + VertexBase.safe_get_vertex_ai_project(litellm_params), + VertexBase.safe_get_vertex_ai_location(litellm_params), + VertexBase.safe_get_vertex_ai_credentials(litellm_params), + ) + def validate_anthropic_messages_environment( self, - headers: dict, + headers: dict[str, str], model: str, messages: Sequence[object], - optional_params: dict, - litellm_params: dict, + optional_params: dict[str, object], + litellm_params: dict[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> tuple[dict, str | None]: + ) -> tuple[dict[str, str], str | None]: """ OPTIONAL Validate the environment for the request """ - # Work on a local copy — router shallow-copies litellm_params so the caller's - # headers dict may be the shared deployment extra_headers object. - headers = dict(headers) - vertex_ai_project: Final = VertexBase.safe_get_vertex_ai_project(litellm_params) - vertex_ai_location: Final = VertexBase.safe_get_vertex_ai_location(litellm_params) - - vertex_credentials: Final = VertexBase.safe_get_vertex_ai_credentials(litellm_params) + vertex_ai_project, vertex_ai_location, vertex_credentials = self._vertex_auth_parameters(litellm_params) access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_ai_project, custom_llm_provider="vertex_ai", ) - headers["Authorization"] = f"Bearer {access_token}" + return self._finalize_vertex_messages_environment( + headers, + model, + messages, + optional_params, + api_base, + access_token, + project_id, + vertex_ai_project, + vertex_ai_location, + ) + + async def avalidate_anthropic_messages_environment( + self, + headers: dict[str, str], # mutable-ok: mirrors the provider validation contract + model: str, + messages: Sequence[object], + optional_params: dict[str, object], # mutable-ok: mirrors the provider validation contract + litellm_params: dict[str, object], # mutable-ok: mirrors the provider validation contract + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict[str, str], str | None]: # mutable-ok: the caller consumes request headers + vertex_ai_project, vertex_ai_location, vertex_credentials = self._vertex_auth_parameters(litellm_params) + access_token, project_id = await self._ensure_access_token_async( + credentials=vertex_credentials, + project_id=vertex_ai_project, + custom_llm_provider="vertex_ai", + ) + return self._finalize_vertex_messages_environment( + headers, + model, + messages, + optional_params, + api_base, + access_token, + project_id, + vertex_ai_project, + vertex_ai_location, + ) + + def _finalize_vertex_messages_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[object], + optional_params: dict, # mutable-ok: preserves the existing provider optional-parameter contract + api_base: str | None, + access_token: str, + project_id: str, + vertex_ai_project: str | None, + vertex_ai_location: str | None, + ) -> tuple[dict[str, str], str | None]: # mutable-ok: the caller consumes request headers + # Router shallow-copies litellm_params, so headers may be shared between requests. + request_headers: Final = dict(headers) + request_headers["Authorization"] = f"Bearer {access_token}" api_base = self.get_complete_vertex_url( custom_api_base=api_base, @@ -65,14 +123,14 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert model=model, ) - headers["content-type"] = "application/json" + request_headers["content-type"] = "application/json" # Add beta headers for Vertex AI tools: Final = optional_params.get("tools", []) beta_values: Final[set[str]] = set() # Get existing beta headers if any - existing_beta: Final = headers.get("anthropic-beta") + existing_beta: Final = request_headers.get("anthropic-beta") if existing_beta: beta_values.update(b.strip() for b in existing_beta.split(",")) @@ -117,9 +175,9 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value) if beta_values: - headers["anthropic-beta"] = ",".join(beta_values) + request_headers["anthropic-beta"] = ",".join(beta_values) - return headers, api_base + return request_headers, api_base def get_complete_url( self, diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index f4920505722..ff8b75af4d4 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -6,6 +6,7 @@ import uuid from typing import Any, Dict, Final, List import httpx +import responses import respx import pytest from fastapi.testclient import TestClient @@ -1709,16 +1710,20 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet @pytest.mark.asyncio @pytest.mark.parametrize("client_headers", _CLIENT_BETA_HEADERS) async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_vertex( - local_beta_headers_config, client_headers -): + local_beta_headers_config: None, client_headers: dict[str, str] +) -> None: """Vertex rawPredict takes the beta as the `anthropic-beta` header and 400s on `safeguards` without it, so the beta rides along with the field.""" from litellm.llms.anthropic.pass_through.messages import handler - from litellm.llms.vertex_ai.vertex_llm_base import VertexBase safeguards, safeguard_results = _claude_code_auto_mode_request() captured: dict[str, object] = {} - with patch.object(VertexBase, "_ensure_access_token", return_value=("test-token", "test-project")): + with responses.RequestsMock(assert_all_requests_are_fired=False) as oauth: + oauth.add( + responses.POST, + "https://oauth2.googleapis.com/token", + json={"access_token": "test-token", "expires_in": 3600}, + ) response = await handler.anthropic_messages( max_tokens=16, messages=[{"role": "user", "content": "hi"}], @@ -1726,7 +1731,14 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet custom_llm_provider="vertex_ai", vertex_project="test-project", vertex_location="global", - vertex_credentials="{}", + vertex_credentials=json.dumps( + { + "type": "authorized_user", + "client_id": "synthetic-client", + "client_secret": "synthetic-secret", + "refresh_token": "synthetic-refresh", + } + ), client=_upstream_answering_with(safeguard_results, captured), safeguards=safeguards, extra_headers=client_headers, diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index ae9b9684f2d..136c444c9f1 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -2,18 +2,98 @@ import copy import json import os from collections.abc import Mapping +from datetime import datetime, timedelta +from queue import SimpleQueue +from threading import get_ident from typing import Final from unittest.mock import MagicMock, patch import pytest +import requests +import responses from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS from litellm.types.router import GenericLiteLLMParams +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("infer_project", (False, True), ids=("explicit-project", "credential-project")) +@pytest.mark.asyncio +async def test_messages_auth_load_and_refresh_leave_the_event_loop(stream: bool, infer_project: bool) -> None: + config: Final = VertexAIPartnerModelsAnthropicMessagesConfig() + loop_thread: Final = get_ident() + token_threads: Final[SimpleQueue[int]] = SimpleQueue() + now: Final = datetime(2025, 1, 1) + model: Final = "claude-sonnet-5" + headers: Final = {"Authorization": "Bearer incoming", "anthropic-beta": "retained-beta", "x-request-id": "req-1"} + optional_params: Final = { + "stream": stream, + "tools": [{"type": ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value, "name": "web_search"}], + } + litellm_params: Final = { + "vertex_ai_project": None if infer_project else "auth-project", + "vertex_ai_location": "global", + "vertex_credentials": json.dumps( + { + "type": "authorized_user", + "client_id": "synthetic-client", + "client_secret": "synthetic-secret", + "refresh_token": "synthetic-refresh", + "quota_project_id": "auth-project", + } + ), + } + + async def validate() -> tuple[Mapping[str, object], str | None]: + return await config.avalidate_anthropic_messages_environment( + headers=headers, model=model, messages=[], optional_params=optional_params, litellm_params=litellm_params + ) + + with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth: + + def token_response(_request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + token_threads.put(get_ident()) + return 200, {}, json.dumps({"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600}) + + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + cold_headers, cold_url = await validate() + assert token_threads.get_nowait() != loop_thread + assert cold_headers["Authorization"] == "Bearer token-1" + assert await validate() == (cold_headers, cold_url) + assert len(oauth.calls) == 1 + + clock.return_value = now + timedelta(hours=2) + refreshed_headers, refreshed_url = await validate() + assert token_threads.get_nowait() != loop_thread + assert refreshed_headers["Authorization"] == "Bearer token-2" + assert refreshed_url == cold_url + assert len(oauth.calls) == 2 + assert config.validate_anthropic_messages_environment( + headers=headers, model=model, messages=[], optional_params=optional_params, litellm_params=litellm_params + ) == (refreshed_headers, refreshed_url) + assert len(oauth.calls) == 2 + + endpoint: Final = "streamRawPredict" if stream else "rawPredict" + assert cold_url == ( + "https://aiplatform.googleapis.com/v1/projects/auth-project/locations/global/" + f"publishers/anthropic/models/{model}:{endpoint}" + ) + assert refreshed_headers == { + "Authorization": "Bearer token-2", + "anthropic-beta": cold_headers["anthropic-beta"], + "x-request-id": "req-1", + "content-type": "application/json", + } + beta_header: Final = cold_headers["anthropic-beta"] + assert isinstance(beta_header, str) + assert set(beta_header.split(",")) == {"retained-beta", ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value} + assert headers == {"Authorization": "Bearer incoming", "anthropic-beta": "retained-beta", "x-request-id": "req-1"} + + def test_validate_environment_uses_vertex_ai_location(): config = VertexAIPartnerModelsAnthropicMessagesConfig() headers = {} From ebe0871a28bbd585b73988c78edebf675985b9de Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Mon, 26 Oct 2026 19:49:32 -0700 Subject: [PATCH 06/12] fix(decisions): preserve omitted probabilities (#48449) (#48524) --- .../llms/base_llm/decisions/transformation.py | 3 +- .../llms/openai/decisions/transformation.py | 10 ++-- .../test_base_decisions_transformation.py | 46 ++++++++++++++++ .../test_openai_decisions_transformation.py | 52 +++++++++++++++++++ 4 files changed, 104 insertions(+), 7 deletions(-) diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index 14c6a2d2a66..b3df4654a6b 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -351,9 +351,10 @@ def _ir_score_answer(question: DecisionsIRScoreQuestion, answer: ScoreAnswer) -> DecisionsIRScoreProbability( value=index, label=answer.legend.get(str(index), level.label), - probability=answer.probabilities.get(str(index), 0.0), + probability=answer.probabilities[str(index)], ) for index, level in enumerate(question.levels) + if str(index) in answer.probabilities ), extra=MappingProxyType(answer.model_extra or {}), ) diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 25d403a0dff..4518fe6a77e 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -270,10 +270,9 @@ def _openai_choice_answer(question: DecisionsIRChoiceQuestion, answer: Decisions name=question.name, choice=answer.choice, probabilities=tuple( - OpenAIChoiceProbability( - value=option.value, probability=probabilities.get(systemone_choice_key(option.value), 0.0) - ) + OpenAIChoiceProbability(value=option.value, probability=probabilities[systemone_choice_key(option.value)]) for option in question.choices + if systemone_choice_key(option.value) in probabilities ), confidence=answer.confidence, ) @@ -285,10 +284,9 @@ def _openai_score_answer(question: DecisionsIRScoreQuestion, answer: DecisionsIR name=question.name, score=answer.score, probabilities=tuple( - OpenAIScoreProbability( - value=index, label=decisions_text(level.label), probability=probabilities.get(index, 0.0) - ) + OpenAIScoreProbability(value=index, label=decisions_text(level.label), probability=probabilities[index]) for index, level in enumerate(question.levels) + if index in probabilities ), confidence=answer.confidence, ) diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index 3aebf56ee73..798cbe5fe00 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -355,3 +355,49 @@ def test_invalid_provider_response_is_an_api_error(payload: object, provider: st with pytest.raises(BaseLLMException, match="unexpected response") as error: config.transform_decisions_response("drex-v2.0", provider, httpx.Response(200, json=payload), request) assert error.value.status_code == 500 + + +@pytest.mark.parametrize( + ("probabilities", "expected"), + ( + ({}, ()), + ({"1": 0.8}, ((1, "medium", 0.8),)), + ({"0": 0.0, "1": 0.8}, ((0, "low", 0.0), (1, "medium", 0.8))), + ({"2": 0.0}, ((2, "high", 0.0),)), + ), + ids=("empty", "missing", "explicit_zero", "only_zero"), +) +def test_score_probabilities_preserve_missing_entries_and_explicit_zeros( + probabilities: Mapping[str, float], expected: tuple[tuple[int, str, float], ...] +) -> None: + ir: Final = systemone_request_to_ir( + _SYSTEMONE_BODY.validate_python( + { + "state": "Classify the fixture.", + "questions": {"score": {"type": "score", "criteria": ["low", "medium", "high"]}}, + } + ) + ) + payload: Final = { + "model": "fixture", + "answers": { + "score": { + "type": "score", + "score": 1, + "confidence": 0.8, + "legend": {"0": "low", "1": "medium", "2": "high"}, + "probabilities": probabilities, + } + }, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + + response: Final = parse_systemone_response(payload, ir) + answer: Final = response.answers[0] + + assert isinstance(answer, DecisionsIRScoreAnswer) + assert tuple((item.value, item.label, item.probability) for item in answer.probabilities) == expected + assert ( + ir_to_systemone_response(response, ir).model_dump(mode="json")["answers"]["score"]["probabilities"] + == probabilities + ) diff --git a/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py b/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py index 2a07462fe81..eddfe578256 100644 --- a/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py +++ b/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py @@ -299,3 +299,55 @@ def test_an_openai_response_becomes_systemone_answers_with_the_callers_score_lab } assert response.usage is not None assert (response.usage.cached_tokens, response.usage.cache_write_tokens) == (256, 64) + + +@pytest.mark.parametrize( + ("choice_probabilities", "score_probabilities"), + ( + ((), ()), + ( + ({"value": True, "probability": 0.8},), + ({"value": 1, "label": "major", "probability": 0.8},), + ), + ( + ( + {"value": True, "probability": 0.0}, + {"value": "escalate", "probability": 0.8}, + ), + ( + {"value": 0, "label": "minor", "probability": 0.0}, + {"value": 1, "label": "major", "probability": 0.8}, + ), + ), + ( + ({"value": True, "probability": 0.0},), + ({"value": 1, "label": "major", "probability": 0.0},), + ), + ), + ids=("empty", "missing", "explicit_zero", "only_zero"), +) +def test_openai_probabilities_preserve_missing_entries_and_explicit_zeros( + choice_probabilities: tuple[Mapping[str, object], ...], + score_probabilities: tuple[Mapping[str, object], ...], +) -> None: + ir: Final = _openai_ir(_OPENAI_REQUEST) + payload: Final = { + **_OPENAI_RESPONSE, + "answers": [ + _PREDICATE_ANSWER, + {**_CHOICE_ANSWER, "probabilities": choice_probabilities}, + { + "type": "score", + "name": "severity", + "score": 1, + "confidence": 0.8, + "probabilities": score_probabilities, + }, + _REFUSAL_ANSWER, + ], + } + + response: Final = ir_to_openai_response(OpenAIDecisionsConfig().parse_response(payload, ir), ir, "requested") + + assert response.answers[1].model_dump(mode="json")["probabilities"] == list(choice_probabilities) + assert response.answers[2].model_dump(mode="json")["probabilities"] == list(score_probabilities) From 53a4cee0d82b5358821383acc4937dfc2f43fd1e Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 23:00:03 -0400 Subject: [PATCH 07/12] fix(ui): honor dark and local logos across Connect and Chat (#48532) --- .../proxy_setting_endpoints.py | 30 +++++--- .../test_proxy_setting_endpoints.py | 77 +++++++++++++------ .../ui-theme/UIThemeSettings.test.tsx | 33 +++++++- .../(dashboard)/ui-theme/UIThemeSettings.tsx | 9 ++- .../src/app/chat/layout.integration.test.tsx | 58 +++++++++++++- .../chat/ChatShell.integration.test.tsx | 16 +++- .../src/components/chat/ChatShell.tsx | 25 ++++++ .../src/components/chat/design.md | 2 + .../src/components/navbar.test.tsx | 32 +++++--- .../src/components/navbar.tsx | 4 +- .../src/contexts/ThemeContext.tsx | 38 ++++++--- ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 +++ 12 files changed, 267 insertions(+), 67 deletions(-) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index e04ddbc75cc..1bb88d5ff69 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -234,6 +234,9 @@ class DefaultTeamSettingsResponse(SettingsResponse): class UIThemeSettingsResponse(SettingsResponse): """Response model for UI theme settings""" + logo_url_via_proxy: bool = False + logo_url_dark_via_proxy: bool = False + _TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_PERMISSIONS)) @@ -1377,7 +1380,7 @@ async def update_sso_settings( tags=["UI Theme Settings"], response_model=UIThemeSettingsResponse, ) -async def get_ui_theme_settings(): +async def get_ui_theme_settings() -> UIThemeSettingsResponse: """ Get UI theme configuration from the litellm_settings. Returns current logo settings for UI customization. @@ -1390,18 +1393,23 @@ async def get_ui_theme_settings(): # Load existing config config: Final = await proxy_config.get_config() - result: Final = await _get_settings_with_schema( - settings_key="ui_theme_config", - settings_class=UIThemeConfig, - config=config, + result: Final = _SettingsWithSchema.model_validate( + await _get_settings_with_schema( + settings_key="ui_theme_config", + settings_class=UIThemeConfig, + config=config, + ) ) - - stored_values: Final = result.get("values", {}) - result["values"] = { - **stored_values, - **{field: _resolve_ui_theme_field(stored_values, field) for field in _UI_THEME_FIELD_ENV_VARS}, + values: Final = { + **result.values, + **{field: _resolve_ui_theme_field(result.values, field) for field in _UI_THEME_FIELD_ENV_VARS}, } - return result + return UIThemeSettingsResponse( + values=values, + field_schema=result.field_schema, + logo_url_via_proxy=values["logo_url"] is None and bool(os.getenv("UI_LOGO_PATH", "").strip()), + logo_url_dark_via_proxy=values["logo_url_dark"] is None and bool(os.getenv("UI_LOGO_PATH_DARK", "").strip()), + ) def _validate_public_image_url(value: str | None, field_name: str) -> None: diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 57b68e6db05..4360655c8ca 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1359,35 +1359,45 @@ class TestProxySettingEndpoints: assert values["logo_url"] == "https://cdn.example.com/logo.png" assert values["favicon_url"] == "https://cdn.example.com/favicon.ico" + @pytest.mark.parametrize("env_logo", [ + "https://env.example.com/logo.png", "/mnt/secret/internal/logo.png", + ]) def test_get_ui_theme_settings_stored_value_wins_over_env( - self, mock_auth, monkeypatch - ): + self, monkeypatch: pytest.MonkeyPatch, env_logo: str, + ) -> None: """A stored ui_theme_config field outranks the env var for that field. The env fallback only fills fields the stored config leaves blank, so the UI-driven flow is unchanged while an unstored field still resolves. """ from litellm.proxy.proxy_server import proxy_config + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import UIThemeSettingsResponse - stored_config = { + stored_config: Final[dict[str, object]] = { "litellm_settings": { - "ui_theme_config": {"logo_url": "https://db.example.com/logo.png"} + "ui_theme_config": { + "logo_url": "https://db.example.com/logo.png", + "logo_url_dark": "https://db.example.com/logo-dark.png", + } } } - async def mock_get_config(): + async def mock_get_config() -> dict[str, object]: return stored_config monkeypatch.setattr(proxy_config, "get_config", mock_get_config) - monkeypatch.setenv("UI_LOGO_PATH", "https://env.example.com/logo.png") + monkeypatch.setenv("UI_LOGO_PATH", env_logo) + monkeypatch.setenv("UI_LOGO_PATH_DARK", env_logo) monkeypatch.setenv("LITELLM_FAVICON_URL", "https://env.example.com/favicon.ico") - response = client.get("/get/ui_theme_settings") + response: Final = client.get("/get/ui_theme_settings") assert response.status_code == 200 - values = response.json()["values"] - assert values["logo_url"] == "https://db.example.com/logo.png" - assert values["favicon_url"] == "https://env.example.com/favicon.ico" + result: Final = UIThemeSettingsResponse.model_validate(response.json()) + assert result.values["logo_url"] == "https://db.example.com/logo.png" + assert result.values["logo_url_dark"] == "https://db.example.com/logo-dark.png" + assert result.values["favicon_url"] == "https://env.example.com/favicon.ico" + assert (result.logo_url_via_proxy, result.logo_url_dark_via_proxy) == (False, False) def test_get_ui_theme_settings_reports_unset_when_absent_everywhere( self, mock_proxy_config, monkeypatch @@ -1403,22 +1413,45 @@ class TestProxySettingEndpoints: assert values["logo_url"] is None assert values["favicon_url"] is None - def test_get_ui_theme_settings_does_not_disclose_local_path_env_value( - self, mock_proxy_config, monkeypatch - ): - """This endpoint is public, so an env-configured local filesystem branding - path must never be surfaced to anonymous callers; only public http(s) URLs. - """ - monkeypatch.setenv("UI_LOGO_PATH", "/mnt/secret/internal/logo.png") + @pytest.mark.usefixtures("mock_proxy_config") + @pytest.mark.parametrize("light,public_light,light_via_proxy", [ + (None, None, False), + (" ", None, False), + ("https://cdn.example.com/light.png", "https://cdn.example.com/light.png", False), + ("/mnt/secret/internal/light.png", None, True), + ]) + @pytest.mark.parametrize("dark,public_dark,dark_via_proxy", [ + (None, None, False), + (" ", None, False), + ("https://cdn.example.com/dark.png", "https://cdn.example.com/dark.png", False), + ("/mnt/secret/internal/dark.png", None, True), + ]) + def test_get_ui_theme_settings_separates_private_logo_rendering_from_settings( + self, monkeypatch: pytest.MonkeyPatch, + light: str | None, public_light: str | None, light_via_proxy: bool, + dark: str | None, public_dark: str | None, dark_via_proxy: bool, + ) -> None: + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import UIThemeSettingsResponse + + for key, value in (("UI_LOGO_PATH", light), ("UI_LOGO_PATH_DARK", dark)): + if value is None: + monkeypatch.delenv(key, raising=False) + else: + monkeypatch.setenv(key, value) monkeypatch.setenv("LITELLM_FAVICON_URL", "file:///etc/favicon.ico") - response = client.get("/get/ui_theme_settings") + response: Final = client.get("/get/ui_theme_settings") assert response.status_code == 200 - values = response.json()["values"] - # the local path / file scheme is withheld rather than disclosed - assert values["logo_url"] is None - assert values["favicon_url"] is None + result: Final = UIThemeSettingsResponse.model_validate(response.json()) + assert (result.values["logo_url"], result.values["logo_url_dark"]) == (public_light, public_dark) + assert (result.logo_url_via_proxy, result.logo_url_dark_via_proxy) == (light_via_proxy, dark_via_proxy) + assert result.values["favicon_url"] is None + assert "/mnt/secret/internal/" not in response.text + assert "file:///etc/favicon.ico" not in response.text + for hint in ("logo_url_via_proxy", "logo_url_dark_via_proxy"): + assert hint not in result.values + assert hint not in json.dumps(result.field_schema) def test_get_ui_settings(self, mock_auth, monkeypatch): """Test retrieving UI settings with allowlist sanitization""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.test.tsx index 4b58590ffd1..86ef33c0d8b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.test.tsx @@ -10,7 +10,8 @@ const setLogoUrl = vi.fn(); const setLogoUrlDark = vi.fn(); const setFaviconUrl = vi.fn(); -vi.mock("@/contexts/ThemeContext", () => ({ +vi.mock("@/contexts/ThemeContext", async (importOriginal) => ({ + ...(await importOriginal()), useTheme: () => ({ logoUrl: null, setLogoUrl, @@ -22,7 +23,7 @@ vi.mock("@/contexts/ThemeContext", () => ({ })); vi.mock("@/components/networking", () => ({ - getProxyBaseUrl: () => "", + getProxyBaseUrl: () => "/gateway", getGlobalLitellmHeaderName: () => "Authorization", })); @@ -111,6 +112,34 @@ describe("UIThemeSettings", () => { expect(bodyOf(patchCalls()[0]).logo_url_dark).toBe("https://a.test/logo-dark.png"); }); + it.each(["Save Changes", "Reset to Default"])( + "should keep local logo routes out of settings and clear preview overrides with %s", + async (action) => { + const user = userEvent.setup(); + fetchMock.mockImplementation(async () => + Response.json({ + values: { logo_url: null, logo_url_dark: null }, + logo_url_via_proxy: true, + logo_url_dark_via_proxy: true, + }), + ); + + render(); + + await waitFor(() => expect(setLogoUrlDark).toHaveBeenCalledWith("/gateway/get_image?theme=dark")); + expect(setLogoUrl).toHaveBeenCalledWith("/gateway/get_image"); + expect(screen.getByPlaceholderText(LOGO_PLACEHOLDER)).toHaveValue(""); + expect(screen.getByPlaceholderText(LOGO_DARK_PLACEHOLDER)).toHaveValue(""); + + await user.click(screen.getByRole("button", { name: action })); + await waitFor(() => expect(patchCalls()).toHaveLength(1)); + expect(bodyOf(patchCalls()[0])).toEqual({ logo_url: null, logo_url_dark: null, favicon_url: null }); + + expect(setLogoUrl).toHaveBeenLastCalledWith(null); + expect(setLogoUrlDark).toHaveBeenLastCalledWith(null); + }, + ); + it("should surface a backend failure when saving fails", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx index b6bc0f28dd5..c0b3bd26001 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx @@ -4,7 +4,7 @@ import { Card, CardContent } from "@/components/ui/card"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; -import { useTheme } from "@/contexts/ThemeContext"; +import { resolveThemeImages, useTheme } from "@/contexts/ThemeContext"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import { toast } from "@/lib/toast"; @@ -40,12 +40,13 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc }); if (response.ok) { const data = await response.json(); + const images = resolveThemeImages(data, proxyBaseUrl); setLogoUrlInput(data.values?.logo_url || ""); setLogoUrlDarkInput(data.values?.logo_url_dark || ""); setFaviconUrlInput(data.values?.favicon_url || ""); - setLogoUrl(data.values?.logo_url || null); - setLogoUrlDark(data.values?.logo_url_dark || null); - setFaviconUrl(data.values?.favicon_url || null); + setLogoUrl(images.logoUrl); + setLogoUrlDark(images.logoUrlDark); + setFaviconUrl(images.faviconUrl); } } catch (error) { console.error("Error fetching theme settings:", error); diff --git a/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx index a7189ea6cd7..2efe698ec69 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx @@ -1,7 +1,8 @@ import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, screen, testQueryClient } from "@/../tests/test-utils"; +import { renderWithProviders, screen, testQueryClient, within } from "@/../tests/test-utils"; import { setServerRootPath } from "@/lib/serverRootPath"; +import { getProxyBaseUrl } from "@/components/networking"; import ChatLayout from "./layout"; const { router, route } = vi.hoisted(() => ({ @@ -33,9 +34,13 @@ vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ const plugin = { name: "observability", display_name: "Observability", url: "https://plugin.example.com" }; const originalLocation = Object.getOwnPropertyDescriptor(window, "location")!; const assign = vi.fn(); +const LIGHT_LOGO = "https://cdn.test/light.svg"; +const DARK_LOGO = "https://cdn.test/dark.svg"; +let themeResponse: unknown = { values: {} }; describe("Chat layout product navigation", () => { beforeEach(() => { + themeResponse = { values: {} }; route.pathname = "/ui/chat"; localStorage.clear(); localStorage.setItem("litellm_plugin_mode", plugin.name); @@ -55,7 +60,7 @@ describe("Chat layout product navigation", () => { expect(new Headers(init?.headers).get("Authorization")).toBe("Bearer chat-layout-test-token"); return Response.json([plugin]); case "/get/ui_theme_settings": - return Response.json({ values: {} }); + return Response.json(themeResponse); case "/sso/get/ui_settings": return Response.json({}); default: @@ -75,6 +80,55 @@ describe("Chat layout product navigation", () => { vi.clearAllMocks(); }); + it.each([ + ["public", [LIGHT_LOGO, DARK_LOGO], [false, false], [LIGHT_LOGO, DARK_LOGO]], + ["local", [null, null], [true, true], ["/get_image", "/get_image?theme=dark"]], + ["public light and local dark", [LIGHT_LOGO, null], [false, true], [LIGHT_LOGO, "/get_image?theme=dark"]], + ["local light and public dark", [null, DARK_LOGO], [true, false], ["/get_image", DARK_LOGO]], + ["public light only", [LIGHT_LOGO, null], [false, false], [LIGHT_LOGO, LIGHT_LOGO]], + ["public dark only", [null, DARK_LOGO], [false, false], [DARK_LOGO, DARK_LOGO]], + ["local light only", [null, null], [true, false], ["/get_image", "/get_image"]], + ["local dark only", [null, null], [false, true], ["/get_image?theme=dark", "/get_image?theme=dark"]], + ] as const)( + "should render %s organization logos from the theme response", + async ( + _label, + [logo_url, logo_url_dark], + [logo_url_via_proxy, logo_url_dark_via_proxy], + [lightSource, darkSource], + ) => { + themeResponse = { values: { logo_url, logo_url_dark }, logo_url_via_proxy, logo_url_dark_via_proxy }; + const user = userEvent.setup(); + renderWithProviders( + +

Chat content

+ , + ); + + await user.click(screen.getByRole("button", { name: "Open chat navigation" })); + const logo = await screen.findByRole("img", { name: "Organization logo" }); + const [light, dark] = within(logo).getAllByAltText(""); + const source = (path: string) => (path.startsWith("/") ? `${getProxyBaseUrl()}${path}` : path); + expect(light).toHaveAttribute("src", source(lightSource)); + expect(dark).toHaveAttribute("src", source(darkSource)); + expect(light).toHaveClass("dark:hidden"); + expect(dark).toHaveClass("hidden", "dark:block"); + }, + ); + + it("should omit organization branding when no logo is configured", async () => { + const user = userEvent.setup(); + renderWithProviders( + +

Chat content

+ , + ); + + await user.click(screen.getByRole("button", { name: "Open chat navigation" })); + expect(await screen.findByRole("dialog", { name: "LiteLLM Chat" })).toBeInTheDocument(); + expect(screen.queryByRole("img", { name: "Organization logo" })).not.toBeInTheDocument(); + }); + it.each([ { pathname: "/ui/chat", root: "/", gateway: "/ui/" }, { pathname: "/gateway/ui/chat/models/", root: "/gateway", gateway: "/gateway/ui/" }, diff --git a/ui/litellm-dashboard/src/components/chat/ChatShell.integration.test.tsx b/ui/litellm-dashboard/src/components/chat/ChatShell.integration.test.tsx index c1af91845ee..3c7b9ef1b1a 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatShell.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatShell.integration.test.tsx @@ -1,9 +1,13 @@ -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders as render } from "@/../tests/test-utils"; +import { renderWithProviders } from "@/../tests/test-utils"; +import type { ReactElement } from "react"; +import { ThemeProvider } from "@/contexts/ThemeContext"; import ChatShell from "./ChatShell"; +const render = (ui: ReactElement) => renderWithProviders({ui}); + const { mockPush, mockUsePathname, mockUseChatShell, mockUseMediaQuery } = vi.hoisted(() => ({ mockPush: vi.fn(), mockUsePathname: vi.fn(() => "/ui/chat"), @@ -34,7 +38,15 @@ vi.mock("usehooks-ts", () => ({ useMediaQuery: mockUseMediaQuery })); vi.mock("./ConversationList", () => ({ default: () =>
})); describe("ChatShell", () => { + beforeEach(() => { + vi.stubGlobal( + "fetch", + vi.fn(async () => Response.json({ values: {} })), + ); + }); + afterEach(() => { + vi.unstubAllGlobals(); mockPush.mockClear(); mockUsePathname.mockReturnValue("/ui/chat"); mockUseMediaQuery.mockReturnValue(true); diff --git a/ui/litellm-dashboard/src/components/chat/ChatShell.tsx b/ui/litellm-dashboard/src/components/chat/ChatShell.tsx index c089a709b64..280e2d10b8c 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatShell.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatShell.tsx @@ -1,6 +1,7 @@ "use client"; import React, { useState } from "react"; +import Image from "next/image"; import { usePathname, useRouter } from "next/navigation"; import { Plus, @@ -34,6 +35,7 @@ import ThemeToggle from "@/components/ThemeToggle/ThemeToggle"; import { useLogout } from "@/app/(dashboard)/hooks/useLogout"; import { uiHref } from "@/utils/uiHref"; import { useChatShell } from "@/contexts/ChatShellContext"; +import { useTheme } from "@/contexts/ThemeContext"; import ViewSwitcher from "@/components/Navbar/ViewSwitcher"; import ConversationList from "./ConversationList"; @@ -94,6 +96,9 @@ const ChatShell: React.FC<{ children: React.ReactNode }> = ({ children }) => { const [searchOpen, setSearchOpen] = useState(false); if (isDesktop && sidebarOpen) setSidebarOpen(false); const collapsed = isDesktop && sidebarCollapsed; + const { logoUrl, logoUrlDark } = useTheme(); + const lightLogo = logoUrl || logoUrlDark; + const darkLogo = logoUrlDark || logoUrl; const routes = getChatRoutes(); const isChatsRoute = pathname === routes.chats; const accountName = userEmail?.split("@")[0] || userId || "Account"; @@ -216,6 +221,26 @@ const ChatShell: React.FC<{ children: React.ReactNode }> = ({ children }) => { hideSearchButton />
+ {lightLogo && darkLogo && ( +
+ + +
+ )}