diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 709df5e64df..e368a1a2275 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2320,6 +2320,9 @@ async def _run_centralized_common_checks( None if isinstance(global_spend_result, BaseException) else global_spend_result ) + if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: + user_api_key_auth_obj.org_id = team_object.organization_id + # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and # master_key tokens get admin from the token; the DB row for the diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index fffa0bf86d2..af9aa901022 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -38,11 +38,44 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( update_batch_in_database, ) from litellm.proxy.utils import handle_exception_on_proxy, is_known_model +from litellm.repositories.table_repositories import ManagedFileRepository from litellm.types.llms.openai import LiteLLMBatchCreateRequest router = APIRouter() +async def _resolve_managed_input_file_storage_url(input_file_id: str) -> "str | None": + """Resolve a managed (unified) input_file_id to its backend storage_url. + + Provider batch handlers (e.g. Vertex AI, which parses a `publishers/` + segment out of the file URI) need a real storage location; the opaque + unified token crashes them. Returns None only when there is no database or + the row has no storage_url yet, so callers fall back to the original id + (which the managed-files deployment hook can still map). Fails closed + rather than dispatch a token that cannot be resolved: 404 when no + managed-file row exists, 503 when the lookup itself errors so the caller + can retry. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return None + try: + db_file = await ManagedFileRepository(prisma_client).table.find_first(where={"unified_file_id": input_file_id}) + except Exception as e: + verbose_proxy_logger.warning("create_batch: managed file lookup failed for %s: %s", input_file_id, e) + raise HTTPException( + status_code=503, + detail={"error": "Could not resolve managed file; please retry"}, + ) + if db_file is None: + raise HTTPException( + status_code=404, + detail={"error": f"Managed file not found: {input_file_id}"}, + ) + return db_file.storage_url or None + + @router.post( "/{provider}/v1/batches", dependencies=[Depends(user_api_key_auth)], @@ -224,6 +257,11 @@ async def create_batch( ) model = target_model_names[0] _create_batch_data["model"] = model + + resolved_storage_url = await _resolve_managed_input_file_storage_url(input_file_id) + if resolved_storage_url is not None: + _create_batch_data["input_file_id"] = resolved_storage_url + if llm_router is None: raise HTTPException( status_code=500, diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index ca3bb0ee361..7b166185865 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -412,6 +412,19 @@ class HeadroomGuardrail(CustomGuardrail): ) if key in body } + tokens_before = stats.get("tokens_before") + tokens_after = stats.get("tokens_after") + if ( + "tokens_saved" not in stats + and isinstance(tokens_before, (int, float)) + and not isinstance(tokens_before, bool) + and isinstance(tokens_after, (int, float)) + and not isinstance(tokens_after, bool) + ): + # Spend tracking (extract_compression_saved_tokens) reads only + # tokens_saved, which the live compression service omits; derive it + # so savings are counted, but let a service-sent value win. + stats["tokens_saved"] = tokens_before - tokens_after return filtered, True, stats async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str: diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 0359d974d19..c0c91ae103f 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3796,6 +3796,167 @@ async def test_centralized_common_checks_user_http_exception_isolates_to_user_on setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_org_id,expected_org_id", + [ + (None, "org-from-team", "org-from-team"), + ("org-pinned-on-key", "org-from-team", "org-pinned-on-key"), + (None, None, None), + ], +) +async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, team_org_id, expected_org_id): + """LIT-4688 regression: a key minted without an organization_id but attached + to an org-linked team must leave auth with org_id set from the team, so the + spend writer (which reads user_api_key_dict.org_id, no team fallback) + credits the org and the org budget cap can actually trip. A key with an + explicitly pinned org_id must win over the team's org.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u", team_id="t1", org_id=key_org_id) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + org_id_seen_by_common_checks = [] + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + side_effect=lambda **kw: org_id_seen_by_common_checks.append(kw["valid_token"].org_id), + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.org_id == expected_org_id + assert org_id_seen_by_common_checks == [expected_org_id] + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_cli_session_token_org_backfilled_from_team(monkeypatch): + """LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted + with a real team_id but no org_id, and their auth path decrypts the blob + without the combined_view team join, so their spend never reached the org. + The centralized-checks backfill must complete the credential from the team + the same way the SQL view does for DB keys.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-lit4688") + + cli_user = LiteLLM_UserTable(user_id="cli-user", user_role="internal_user", teams=["t-cli"], models=[]) + blob = ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=cli_user, team_id="t-cli", team_alias="cli-team") + token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(blob) + assert token is not None + assert token.is_session_token is True + assert token.team_id == "t-cli" + assert token.org_id is None + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + org_linked_team = LiteLLM_TeamTableCachedObj(team_id="t-cli", organization_id="org-infoops") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=org_linked_team, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + assert token.org_id == "org-infoops" + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_org_backfill_survives_team_fetch_failure(): + """When the team DB fetch fails, the token-derived fallback team carries no + organization_id, so the backfill must leave org_id as None rather than + crash or mis-attribute.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u", team_id="t1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=Exception("DB down"), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.org_id is None + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_master_key_auth_substitutes_alias_for_api_key(): """ diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 26c654cd154..fe09241ee5c 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -76,9 +76,7 @@ CREDS: Dict[str, Dict[str, str]] = { } # A real model-encoded file id: decodes to "azure/gpt-4o", strips to "file-original123". -AZURE_FILE_ID = encode_file_id_with_model( - "file-original123", "azure/gpt-4o", id_type="file" -) +AZURE_FILE_ID = encode_file_id_with_model("file-original123", "azure/gpt-4o", id_type="file") def make_batch( @@ -166,9 +164,7 @@ def harness(): router = MagicMock(spec=Router) router.acreate_batch = AsyncMock(return_value=make_batch()) - router.get_deployment_credentials_with_provider = MagicMock( - side_effect=_creds_lookup - ) + router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) read_body = AsyncMock(side_effect=lambda request: body_holder["body"]) pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock())) @@ -186,11 +182,7 @@ def harness(): pre_call, ) ) - stack.enter_context( - patch.object( - ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers - ) - ) + stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers)) stack.enter_context( patch.object( endpoints, @@ -200,14 +192,13 @@ def harness(): ) stack.enter_context(patch.object(endpoints, "is_known_model", is_known_model)) stack.enter_context(patch.object(litellm, "acreate_batch", litellm_acreate)) - stack.enter_context( - patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False) - ) + stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False)) stack.enter_context(patch.object(proxy_server, "llm_router", router)) stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging)) stack.enter_context(patch.object(proxy_server, "general_settings", {})) stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock())) stack.enter_context(patch.object(proxy_server, "version", "test-version")) + stack.enter_context(patch.object(proxy_server, "prisma_client", None)) h = Harness( body=body_holder, @@ -283,9 +274,7 @@ async def test_create__model_encoded_file_id(harness): } # 4. OUTPUT SHAPE - ids re-encoded with the model; input_file_id restored. - assert resp.id == encode_file_id_with_model( - "batch-provider-id", "azure/gpt-4o", id_type="batch" - ) + assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch") assert resp.input_file_id == AZURE_FILE_ID @@ -307,12 +296,8 @@ async def test_create__model_encoded_file_id__encodes_output_and_error_ids(harne resp = await call_create(harness) - assert resp.output_file_id == encode_file_id_with_model( - "file-out-raw", "azure/gpt-4o" - ) - assert resp.error_file_id == encode_file_id_with_model( - "file-err-raw", "azure/gpt-4o" - ) + assert resp.output_file_id == encode_file_id_with_model("file-out-raw", "azure/gpt-4o") + assert resp.error_file_id == encode_file_id_with_model("file-err-raw", "azure/gpt-4o") @pytest.mark.asyncio @@ -358,9 +343,7 @@ async def test_create__model_from_body(harness): payload = harness.acreate_kwargs() assert payload["custom_llm_provider"] == "vertex_ai" assert payload["input_file_id"] == "file-plain" - assert resp.id == encode_file_id_with_model( - "batch-provider-id", "vertex-model", id_type="batch" - ) + assert resp.id == encode_file_id_with_model("batch-provider-id", "vertex-model", id_type="batch") @pytest.mark.asyncio @@ -495,10 +478,9 @@ async def test_create__unified_file_id_single_model(harness): "completion_window": "24h", }, ) - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz" - ), patch.object( - endpoints, "get_models_from_unified_file_id", return_value=["gpt-4o-mini"] + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gpt-4o-mini"]), ): resp = await call_create(harness) @@ -522,10 +504,9 @@ async def test_create__unified_file_id_not_exactly_one_model_400(harness, models "completion_window": "24h", }, ) - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz" - ), patch.object( - endpoints, "get_models_from_unified_file_id", return_value=models + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=models), ): with pytest.raises(ProxyException) as exc: await call_create(harness) @@ -535,6 +516,182 @@ async def test_create__unified_file_id_not_exactly_one_model_400(harness, models harness.litellm_acreate.assert_not_called() +@pytest.mark.asyncio +async def test_create__unified_file_id_resolves_real_storage_url(harness): + """A base64 unified_file_id is a LiteLLM-internal token, not a real + provider-side file reference (e.g. Vertex AI's batch transformation parses + a `publishers/` segment out of the file URI and crashes on the opaque + base64 string). The real backend location (`storage_url`) must be looked + up from LiteLLM_ManagedFileTable and substituted before dispatch. + + Regression lock on the lookup key: LiteLLM_ManagedFileTable.unified_file_id + stores the raw base64 file id (see schema.prisma and the enterprise + managed-files hook, which queries with the raw id), NOT the decoded + litellm_proxy:... string. Querying with the decoded string never matches + and silently falls back.""" + set_body( + harness, + { + "input_file_id": "litellm_proxy_unified_id", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + fake_db_file = MagicMock( + storage_url="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.0/abc", + ) + find_first = AsyncMock(return_value=fake_db_file) + fake_repo_instance = MagicMock() + fake_repo_instance.table.find_first = find_first + fake_repo_cls = MagicMock(return_value=fake_repo_instance) + + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]), + patch.object(proxy_server, "prisma_client", MagicMock()), + patch.object(endpoints, "ManagedFileRepository", fake_repo_cls), + ): + resp = await call_create(harness) + + assert harness.router_kwargs()["input_file_id"] == fake_db_file.storage_url + find_first.assert_awaited_once_with(where={"unified_file_id": "litellm_proxy_unified_id"}) + assert resp.input_file_id == "litellm_proxy_unified_id" + assert resp._hidden_params["unified_file_id"] == "unified-xyz" + + +@pytest.mark.asyncio +async def test_create__unified_file_id_db_error_fails_closed_503(harness): + """A lookup error leaves the token unresolved, so it fails closed with a + retryable 503 rather than dispatching the opaque id into the provider crash + it cannot parse. Nothing is dispatched.""" + set_body( + harness, + { + "input_file_id": "litellm_proxy_unified_id", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + find_first = AsyncMock(side_effect=Exception("db unavailable")) + fake_repo_instance = MagicMock() + fake_repo_instance.table.find_first = find_first + fake_repo_cls = MagicMock(return_value=fake_repo_instance) + + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]), + patch.object(proxy_server, "prisma_client", MagicMock()), + patch.object(endpoints, "ManagedFileRepository", fake_repo_cls), + ): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "503" + harness.router_acreate.assert_not_called() + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__multi_model_unified_file_with_loadbalancing_keeps_router_branch(harness): + """Regression guard: a multi-model managed file dispatched with an explicit + router model under load balancing must keep taking the load-balanced router + branch, exactly as on the base revision, where the managed-files deployment + hook remaps the unified id per model. Routing it into the unified branch + instead would trip that branch's "exactly one model" 400 and break a path + that works today, so the unified-file resolution must not steal the + load-balanced branch.""" + set_body( + harness, + { + "input_file_id": "litellm_proxy_unified_id", + "model": "vertex-model", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + harness.is_known_model.return_value = True + + with ( + patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True), + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["model-a", "model-b"]), + ): + await call_create(harness) + + assert harness.router_acreate.call_count == 1 + assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id" + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__unified_file_id_missing_row_fails_closed_404(harness): + """With a database present, a managed unified id that has no row cannot be + resolved to a real storage location, so it fails closed with a 404 rather + than dispatching the opaque token, which would hit the Vertex + publishers-segment IndexError this PR exists to prevent.""" + set_body( + harness, + { + "input_file_id": "litellm_proxy_unified_id", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + find_first = AsyncMock(return_value=None) + fake_repo_instance = MagicMock() + fake_repo_instance.table.find_first = find_first + fake_repo_cls = MagicMock(return_value=fake_repo_instance) + + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]), + patch.object(proxy_server, "prisma_client", MagicMock()), + patch.object(endpoints, "ManagedFileRepository", fake_repo_cls), + ): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "404" + harness.router_acreate.assert_not_called() + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__unified_file_id_legacy_row_without_storage_url_dispatches_raw( + harness, +): + """A managed file whose row predates the storage_url column still dispatches + the original id (the managed-files deployment hook maps it); the row exists, + so this is not the missing-row fail-closed case.""" + set_body( + harness, + { + "input_file_id": "litellm_proxy_unified_id", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + fake_db_file = MagicMock(storage_url=None) + find_first = AsyncMock(return_value=fake_db_file) + fake_repo_instance = MagicMock() + fake_repo_instance.table.find_first = find_first + fake_repo_cls = MagicMock(return_value=fake_repo_instance) + + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]), + patch.object(proxy_server, "prisma_client", MagicMock()), + patch.object(endpoints, "ManagedFileRepository", fake_repo_cls), + ): + await call_create(harness) + + assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id" + + @pytest.mark.asyncio async def test_create__model_encoded_beats_unified(harness): """Precedence row: a file id that is BOTH model-encoded and (pretend) unified @@ -547,10 +704,9 @@ async def test_create__model_encoded_beats_unified(harness): "completion_window": "24h", }, ) - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz" - ), patch.object( - endpoints, "get_models_from_unified_file_id", return_value=["something-else"] + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["something-else"]), ): await call_create(harness) @@ -579,9 +735,7 @@ async def test_create__loadbalancing_routes_to_router(harness): with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True): await call_create(harness) - harness.is_known_model.assert_called_once_with( - model="lb-model", llm_router=harness.router - ) + harness.is_known_model.assert_called_once_with(model="lb-model", llm_router=harness.router) assert harness.router_acreate.call_count == 1 harness.litellm_acreate.assert_not_called() harness.creds_resolver.assert_not_called() @@ -630,9 +784,7 @@ async def test_create__team_expiry_injected(harness): }, ) - await call_create( - harness, user=_user_with_expiry({"anchor": "created_at", "seconds": 3600}) - ) + await call_create(harness, user=_user_with_expiry({"anchor": "created_at", "seconds": 3600})) assert harness.acreate_kwargs()["output_expires_after"] == { "anchor": "created_at", @@ -738,12 +890,7 @@ async def test_create__exception_calls_failure_hook(harness): await call_create(harness) harness.logging.post_call_failure_hook.assert_called_once() - assert ( - harness.logging.post_call_failure_hook.call_args.kwargs[ - "original_exception" - ].args[0] - == "provider boom" - ) + assert harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom" # =========================================================================== # @@ -770,9 +917,7 @@ async def test_create__exception_calls_failure_hook(harness): # A real model-encoded BATCH id: decodes to "azure/gpt-4o", strips to # "batch_orig123". Distinct from AZURE_FILE_ID so retrieve tests can't pass by # accidentally reusing the create fixture's value. -AZURE_BATCH_ID = encode_file_id_with_model( - "batch_orig123", "azure/gpt-4o", id_type="batch" -) +AZURE_BATCH_ID = encode_file_id_with_model("batch_orig123", "azure/gpt-4o", id_type="batch") # A realistic decoded unified batch id (what _is_base64_encoded_unified_file_id # returns). model_id / llm_batch_id are parsed out of this by the real helpers. @@ -823,9 +968,7 @@ def retrieve_harness(): router = MagicMock(spec=Router) router.aretrieve_batch = AsyncMock(return_value=make_batch()) - router.get_deployment_credentials_with_provider = MagicMock( - side_effect=_creds_lookup - ) + router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock())) get_headers = MagicMock(return_value={}) @@ -846,11 +989,7 @@ def retrieve_harness(): pre_call, ) ) - stack.enter_context( - patch.object( - ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers - ) - ) + stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers)) stack.enter_context( patch.object( endpoints, @@ -865,24 +1004,12 @@ def retrieve_harness(): provider_from_query, ) ) - stack.enter_context( - patch.object(endpoints, "get_batch_from_database", get_batch_from_db) - ) - stack.enter_context( - patch.object(endpoints, "update_batch_in_database", update_batch_in_db) - ) - stack.enter_context( - patch.object(endpoints, "resolve_input_file_id_to_unified", resolve_input) - ) - stack.enter_context( - patch.object( - endpoints, "resolve_output_file_ids_to_unified", resolve_output - ) - ) + stack.enter_context(patch.object(endpoints, "get_batch_from_database", get_batch_from_db)) + stack.enter_context(patch.object(endpoints, "update_batch_in_database", update_batch_in_db)) + stack.enter_context(patch.object(endpoints, "resolve_input_file_id_to_unified", resolve_input)) + stack.enter_context(patch.object(endpoints, "resolve_output_file_ids_to_unified", resolve_output)) stack.enter_context(patch.object(litellm, "aretrieve_batch", litellm_aretrieve)) - stack.enter_context( - patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False) - ) + stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False)) stack.enter_context(patch.object(proxy_server, "llm_router", router)) stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging)) stack.enter_context(patch.object(proxy_server, "general_settings", {})) @@ -956,9 +1083,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness): } # 4. OUTPUT SHAPE - ids re-encoded with the model for the round-trip. - assert resp.id == encode_file_id_with_model( - "batch-provider-id", "azure/gpt-4o", id_type="batch" - ) + assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch") # write-back to the managed-object table happened, tagged as a retrieve. assert retrieve_harness.update_batch_in_db.call_count == 1 @@ -989,12 +1114,8 @@ async def test_retrieve__model_encoded_id__encodes_output_and_error_ids( resp = await call_retrieve(retrieve_harness, AZURE_BATCH_ID) - assert resp.output_file_id == encode_file_id_with_model( - "file-out-raw", "azure/gpt-4o" - ) - assert resp.error_file_id == encode_file_id_with_model( - "file-err-raw", "azure/gpt-4o" - ) + assert resp.output_file_id == encode_file_id_with_model("file-out-raw", "azure/gpt-4o") + assert resp.error_file_id == encode_file_id_with_model("file-err-raw", "azure/gpt-4o") @pytest.mark.asyncio @@ -1018,9 +1139,7 @@ async def test_retrieve__model_encoded_beats_loadbalancing(retrieve_harness): @pytest.mark.asyncio async def test_retrieve__unified_batch_id_routes_to_router(retrieve_harness): - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID - ): + with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): resp = await call_retrieve(retrieve_harness, "batch-unified-blob") # DISPATCH - router fired, direct litellm did not. @@ -1125,9 +1244,7 @@ async def test_retrieve__fallback_provider_precedence_path_over_header( @pytest.mark.asyncio -@pytest.mark.parametrize( - "status", ["completed", "complete", "failed", "cancelled", "expired"] -) +@pytest.mark.parametrize("status", ["completed", "complete", "failed", "cancelled", "expired"]) async def test_retrieve__db_terminal_state_short_circuits(retrieve_harness, status): # "complete" is the DB-normalized alias of "completed"; it is not a valid # constructor literal but reaches the endpoint via a stored row, so set it @@ -1151,9 +1268,7 @@ async def test_retrieve__db_terminal_unified_resolves_file_ids(retrieve_harness) db_response = make_batch(id="batch-from-db", status="completed") retrieve_harness.get_batch_from_db.return_value = (MagicMock(), db_response) - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID - ): + with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): await call_retrieve(retrieve_harness, "batch-unified-blob") # Terminal short-circuit still resolves raw provider file ids to unified. @@ -1186,9 +1301,7 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn async def test_retrieve__uses_aretrieve_batch_route_type(retrieve_harness): await call_retrieve(retrieve_harness, "batch-raw-xyz") - assert ( - retrieve_harness.pre_call.call_args.kwargs["route_type"] == "aretrieve_batch" - ) + assert retrieve_harness.pre_call.call_args.kwargs["route_type"] == "aretrieve_batch" @pytest.mark.asyncio @@ -1200,9 +1313,7 @@ async def test_retrieve__exception_calls_failure_hook(retrieve_harness): retrieve_harness.logging.post_call_failure_hook.assert_called_once() assert ( - retrieve_harness.logging.post_call_failure_hook.call_args.kwargs[ - "original_exception" - ].args[0] + retrieve_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom" ) @@ -1275,9 +1386,7 @@ def list_harness(): router = MagicMock(spec=Router) router.alist_batches = AsyncMock(return_value=FakeListPage([])) - router.get_deployment_credentials_with_provider = MagicMock( - side_effect=_creds_lookup - ) + router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) read_body = AsyncMock(side_effect=lambda request: body_holder["body"]) pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock())) @@ -1295,11 +1404,7 @@ def list_harness(): pre_call, ) ) - stack.enter_context( - patch.object( - ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers - ) - ) + stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers)) stack.enter_context( patch.object( endpoints, @@ -1432,21 +1537,15 @@ async def test_list__managed_files_beats_model_param(list_harness): ) @pytest.mark.asyncio async def test_list__model_from_body_routes_and_encodes(list_harness): - list_harness.litellm_alist.return_value = FakeListPage( - [make_batch(id="batch-1"), make_batch(id="batch-2")] - ) + list_harness.litellm_alist.return_value = FakeListPage([make_batch(id="batch-1"), make_batch(id="batch-2")]) resp = await call_list(list_harness, body={"model": "azure/gpt-4o"}) assert list_harness.litellm_alist.call_count == 1 list_harness.router_alist.assert_not_called() list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") - assert resp.data[0].id == encode_file_id_with_model( - "batch-1", "azure/gpt-4o", id_type="batch" - ) - assert resp.data[1].id == encode_file_id_with_model( - "batch-2", "azure/gpt-4o", id_type="batch" - ) + assert resp.data[0].id == encode_file_id_with_model("batch-1", "azure/gpt-4o", id_type="batch") + assert resp.data[1].id == encode_file_id_with_model("batch-2", "azure/gpt-4o", id_type="batch") # --------------------------------------------------------------------------- # @@ -1577,12 +1676,7 @@ async def test_list__exception_calls_failure_hook(list_harness): await call_list(list_harness) list_harness.logging.post_call_failure_hook.assert_called_once() - assert ( - list_harness.logging.post_call_failure_hook.call_args.kwargs[ - "original_exception" - ].args[0] - == "provider boom" - ) + assert list_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom" # =========================================================================== # @@ -1645,9 +1739,7 @@ def cancel_harness(): router = MagicMock(spec=Router) router.acancel_batch = AsyncMock(return_value=make_batch()) - router.get_deployment_credentials_with_provider = MagicMock( - side_effect=_creds_lookup - ) + router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock())) # add_litellm_data_to_request is a passthrough that returns the data it got. @@ -1666,11 +1758,7 @@ def cancel_harness(): pre_call, ) ) - stack.enter_context( - patch.object( - ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers - ) - ) + stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers)) stack.enter_context( patch.object( endpoints, @@ -1685,22 +1773,16 @@ def cancel_harness(): provider_from_query, ) ) - stack.enter_context( - patch.object(endpoints, "update_batch_in_database", update_batch_in_db) - ) + stack.enter_context(patch.object(endpoints, "update_batch_in_database", update_batch_in_db)) stack.enter_context(patch.object(litellm, "acancel_batch", litellm_acancel)) - stack.enter_context( - patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False) - ) + stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False)) stack.enter_context(patch.object(proxy_server, "llm_router", router)) stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging)) stack.enter_context(patch.object(proxy_server, "general_settings", {})) stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock())) stack.enter_context(patch.object(proxy_server, "version", "test-version")) stack.enter_context(patch.object(proxy_server, "prisma_client", MagicMock())) - stack.enter_context( - patch.object(proxy_server, "add_litellm_data_to_request", add_data) - ) + stack.enter_context(patch.object(proxy_server, "add_litellm_data_to_request", add_data)) yield CancelHarness( data=data_holder, @@ -1765,9 +1847,7 @@ async def test_cancel__model_encoded_id(cancel_harness): } # OUTPUT SHAPE - response id re-encoded with the DECODED model. - assert resp.id == encode_file_id_with_model( - "batch-provider-id", "azure/gpt-4o", id_type="batch" - ) + assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch") # write-back tagged as a cancel. assert cancel_harness.update_batch_in_db.call_count == 1 @@ -1786,9 +1866,7 @@ async def test_cancel__model_encoded_id_forwards_deployment_model(cancel_harness @pytest.mark.asyncio async def test_cancel__model_encoded_beats_unified(cancel_harness): - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID - ): + with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): await call_cancel(cancel_harness, AZURE_BATCH_ID) assert cancel_harness.litellm_acancel.call_count == 1 @@ -1804,9 +1882,7 @@ async def test_cancel__model_encoded_beats_unified(cancel_harness): @pytest.mark.asyncio async def test_cancel__unified_batch_id_routes_to_router(cancel_harness): - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID - ): + with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): resp = await call_cancel(cancel_harness, "batch-unified-blob") # DISPATCH - router fired, litellm did not, no creds lookup. @@ -1845,8 +1921,9 @@ async def test_cancel__unified_missing_model_id_400(cancel_harness): @pytest.mark.asyncio async def test_cancel__unified_no_router_500(cancel_harness): - with patch.object(proxy_server, "llm_router", None), patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID + with ( + patch.object(proxy_server, "llm_router", None), + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID), ): with pytest.raises(ProxyException) as exc: await call_cancel(cancel_harness, "batch-unified-blob") @@ -1885,9 +1962,7 @@ async def test_cancel__fallback_provider_path_param(cancel_harness): @pytest.mark.asyncio async def test_cancel__fallback_provider_from_data_body(cancel_harness): - await call_cancel( - cancel_harness, "batch-raw-xyz", data_extra={"custom_llm_provider": "bedrock"} - ) + await call_cancel(cancel_harness, "batch-raw-xyz", data_extra={"custom_llm_provider": "bedrock"}) assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "bedrock" @@ -1954,10 +2029,7 @@ async def test_cancel__exception_calls_failure_hook(cancel_harness): cancel_harness.logging.post_call_failure_hook.assert_called_once() assert ( - cancel_harness.logging.post_call_failure_hook.call_args.kwargs[ - "original_exception" - ].args[0] - == "provider boom" + cancel_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom" ) @@ -1979,9 +2051,10 @@ async def test_create__loadbalancing_no_router_500(harness): }, ) harness.is_known_model.return_value = True - with patch.object( - litellm, "enable_loadbalancing_on_batch_endpoints", True - ), patch.object(proxy_server, "llm_router", None): + with ( + patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True), + patch.object(proxy_server, "llm_router", None), + ): with pytest.raises(ProxyException) as exc: await call_create(harness) @@ -2000,12 +2073,10 @@ async def test_create__unified_no_router_500(harness): "completion_window": "24h", }, ) - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz" - ), patch.object( - endpoints, "get_models_from_unified_file_id", return_value=["gpt-4o-mini"] - ), patch.object( - proxy_server, "llm_router", None + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gpt-4o-mini"]), + patch.object(proxy_server, "llm_router", None), ): with pytest.raises(ProxyException) as exc: await call_create(harness) @@ -2015,9 +2086,10 @@ async def test_create__unified_no_router_500(harness): @pytest.mark.asyncio async def test_retrieve__unified_no_router_500(retrieve_harness): - with patch.object( - endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID - ), patch.object(proxy_server, "llm_router", None): + with ( + patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID), + patch.object(proxy_server, "llm_router", None), + ): with pytest.raises(ProxyException) as exc: await call_retrieve(retrieve_harness, "batch-unified-blob") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 776df985d46..4dc527ca45d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -11,6 +11,9 @@ Tests cover: - /v1/compress non-2xx surfaces as httpx.HTTPStatusError (raise_for_status), not a status_code check on the returned response -- both are handled - unreachable_fallback="fail_open" forwards the request uncompressed instead of raising +- tokens_saved is derived from tokens_before/tokens_after when the compression + service omits it, passed through verbatim when present, and skipped (without + breaking compression) when the token counts are not numeric - CCR: headroom_retrieve tool injected when compressed messages contain hashes - CCR: async_should_run_agentic_loop returns True when response has headroom_retrieve tool calls - CCR: async_build_agentic_loop_plan calls retrieve endpoint and builds follow-up messages @@ -32,6 +35,9 @@ from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( has_headroom_retrieve_tool, HEADROOM_RETRIEVE_TOOL_NAME, ) +from litellm.proxy.spend_tracking.compression_savings import ( + extract_compression_saved_tokens, +) from litellm.types.utils import GenericGuardrailAPIInputs FAKE_API_BASE = "https://headroom.example.com" @@ -165,6 +171,113 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages( assert "headroom" in _applied_guardrails(request_data) +def _recorded_guardrail_response(request_data: dict) -> dict: + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(entries) == 1 + return entries[0]["guardrail_response"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_derives_tokens_saved_when_service_omits_it( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + # _make_compress_response omits tokens_saved, matching the live service. + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + request_data: dict = {"model": "gpt-4o"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + stats = _recorded_guardrail_response(request_data) + assert stats["tokens_saved"] == 900 + + # Spend tracking reads the entry under the spend-log metadata key. + entry = request_data["metadata"]["standard_logging_guardrail_information"][0] + assert extract_compression_saved_tokens({"guardrail_information": [entry]}) == 900 + + +@pytest.mark.asyncio +async def test_apply_guardrail_passes_through_service_sent_tokens_saved( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + # Deliberately different from tokens_before - tokens_after (900): the + # service-sent value must win over the derived one. + mock_response.json.return_value["tokens_saved"] = 123 + request_data: dict = {"model": "gpt-4o"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert _recorded_guardrail_response(request_data)["tokens_saved"] == 123 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tokens_before, tokens_after", + [ + ("1000", "100"), + (True, False), + (None, None), + ], +) +async def test_apply_guardrail_skips_derivation_for_non_numeric_token_counts( + guardrail: HeadroomGuardrail, + tokens_before, + tokens_after, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + mock_response.json.return_value["tokens_before"] = tokens_before + mock_response.json.return_value["tokens_after"] = tokens_after + request_data: dict = {"model": "gpt-4o"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert "tokens_saved" not in _recorded_guardrail_response(request_data) + # Compression itself is unaffected by the skipped derivation. + assert result.get("structured_messages") == COMPRESSED_MESSAGES + + @pytest.mark.asyncio async def test_apply_guardrail_injects_retrieve_tool_when_hashes_present( guardrail: HeadroomGuardrail, diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 289012659a1..09d0032e6c9 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2959,11 +2959,6 @@ "count": 1 } }, - "src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/budget_duration_dropdown.tsx": { "local/filename-pascal-case": { "count": 1 @@ -3549,11 +3544,6 @@ "count": 1 } }, - "src/components/routing_groups/RoutingGroupsTable.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, "src/components/routing_groups/index.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.test.tsx deleted file mode 100644 index 58395371bbe..00000000000 --- a/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.test.tsx +++ /dev/null @@ -1,148 +0,0 @@ -import { render, screen, waitFor } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import { describe, expect, it, vi } from "vitest"; -import { TableHeaderSortDropdown } from "./TableHeaderSortDropdown"; - -describe("TableHeaderSortDropdown", () => { - it("should render", () => { - const onSortChange = vi.fn(); - render(); - expect(screen.getByRole("button")).toBeInTheDocument(); - }); - - it("should open dropdown menu when button is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Ascending")).toBeInTheDocument(); - expect(screen.getByText("Descending")).toBeInTheDocument(); - expect(screen.getByText("Reset")).toBeInTheDocument(); - }); - }); - - it("should call onSortChange with asc when ascending option is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Ascending")).toBeInTheDocument(); - }); - - const ascendingOption = screen.getByText("Ascending"); - await user.click(ascendingOption); - - expect(onSortChange).toHaveBeenCalledTimes(1); - expect(onSortChange).toHaveBeenCalledWith("asc"); - }); - - it("should call onSortChange with desc when descending option is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Descending")).toBeInTheDocument(); - }); - - const descendingOption = screen.getByText("Descending"); - await user.click(descendingOption); - - expect(onSortChange).toHaveBeenCalledTimes(1); - expect(onSortChange).toHaveBeenCalledWith("desc"); - }); - - it("should call onSortChange with false when reset option is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Reset")).toBeInTheDocument(); - }); - - const resetOption = screen.getByText("Reset"); - await user.click(resetOption); - - expect(onSortChange).toHaveBeenCalledTimes(1); - expect(onSortChange).toHaveBeenCalledWith(false); - }); - - it("should highlight ascending option when sort state is asc", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - const ascendingOption = screen.getByText("Ascending"); - const menuItem = ascendingOption.closest(".ant-dropdown-menu-item"); - expect(menuItem).toHaveClass("ant-dropdown-menu-item-selected"); - }); - }); - - it("should highlight descending option when sort state is desc", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - const descendingOption = screen.getByText("Descending"); - const menuItem = descendingOption.closest(".ant-dropdown-menu-item"); - expect(menuItem).toHaveClass("ant-dropdown-menu-item-selected"); - }); - }); - - it("should not highlight any option when sort state is false", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Ascending")).toBeInTheDocument(); - }); - - const ascendingOption = screen.getByText("Ascending"); - const menuItem = ascendingOption.closest(".ant-dropdown-menu-item"); - expect(menuItem).not.toHaveClass("ant-dropdown-menu-item-selected"); - }); - - it("should stop event propagation when button is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - const onParentClick = vi.fn(); - - render( -
- -
, - ); - - const button = screen.getByRole("button"); - await user.click(button); - - expect(onParentClick).not.toHaveBeenCalled(); - }); -}); diff --git a/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx deleted file mode 100644 index c83257c5c83..00000000000 --- a/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx +++ /dev/null @@ -1,82 +0,0 @@ -import React from "react"; -import { Button, Dropdown, MenuProps } from "antd"; -import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, XIcon } from "@heroicons/react/outline"; - -export type SortState = "asc" | "desc" | false; - -interface TableHeaderSortDropdownProps { - /** - * Current sort state: "asc", "desc", or false for neutral - */ - sortState: SortState; - /** - * Callback when sort state changes - * @param newState - The new sort state: "asc", "desc", or false - */ - onSortChange: (newState: SortState) => void; - /** - * Optional column ID for identification - */ - columnId?: string; -} - -export const TableHeaderSortDropdown: React.FC = ({ sortState, onSortChange }) => { - const handleMenuClick: MenuProps["onClick"] = ({ key }) => { - if (key === "asc") { - onSortChange("asc"); - } else if (key === "desc") { - onSortChange("desc"); - } else if (key === "reset") { - onSortChange(false); - } - }; - - const menuItems: MenuProps["items"] = [ - { - key: "asc", - label: "Ascending", - icon: , - }, - { - key: "desc", - label: "Descending", - icon: , - }, - { - key: "reset", - label: "Reset", - icon: , - }, - ]; - - // Determine which icon to display based on current sort state - const renderIcon = () => { - if (sortState === "asc") { - return ; - } else if (sortState === "desc") { - return ; - } else { - return ; - } - }; - - return ( - -