diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 5a180c13c53..f98077ea2f0 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -122,6 +122,11 @@ jobs: uv run --no-sync ruff check . cd .. + - name: Run Ruff linting (test tree) + if: steps.changes.outputs.decision != 'skip' + run: | + uv run --no-sync ruff check --config ruff-tests.toml tests + - name: Check strict-rule budget (delta vs base) if: steps.changes.outputs.decision != 'skip' run: | diff --git a/Makefile b/Makefile index 5ae2638fbaa..580d663ba53 100644 --- a/Makefile +++ b/Makefile @@ -160,6 +160,7 @@ lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) # Linting targets lint-ruff: $(LINT_DEP_INSTALL) cd litellm && $(UV_RUN) ruff check . && cd .. + $(UV_RUN) ruff check --config ruff-tests.toml tests # faster linter for developing ... # inspiration from: diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 2eb4232fef9..286f7528896 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1039,16 +1039,29 @@ class LiteLLMUnknownProvider(BadRequestError): class GuardrailRaisedException(Exception): + """ + Raised both when a guardrail judged content and when it could not judge it at all, since a + guardrail that fails closed refuses the request the same way a policy violation does. + + ``blocked_content`` separates the two. Set it only where the guardrail actually reached a + verdict on the payload; leave it alone for an unreachable backend, a timeout, or a response + the integration could not parse. Callers that treat a block as something other than a plain + failure, such as the batch path dropping one record and submitting the rest, must gate on it, + because dropping a record no guardrail ever inspected is a silent loss of enforcement. + """ + def __init__( self, guardrail_name: str | None = None, message: str = "", should_wrap_with_default_message: bool = True, status_code: int = 400, + blocked_content: bool = False, ): default_message: Final = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}" self.guardrail_name = guardrail_name self.status_code = status_code + self.blocked_content = blocked_content self.message = default_message if should_wrap_with_default_message else message super().__init__(self.message) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 0172c789d1e..f2e390625f5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -65,6 +65,41 @@ _guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.Cont ) +def is_guardrail_intervention(e: Exception) -> bool: + """ + Returns True if the exception represents an intentional guardrail block + (this was logged previously as an API failure - guardrail_failed_to_respond). + + Guardrails signal intentional blocks by raising: + - GuardrailRaisedException (generic guardrail API, tool permission) + - BlockedPiiEntityError (Presidio PII detection) + - SensitiveDataRouteException (sensitive-data reroute to on-premise model) + - HTTPException with a block-signalling status (400, 403, 422) + - ModifyResponseException (passthrough mode violation) + + Only the statuses guardrails use in-tree to signal a deliberate rejection + count as an intervention: 400 (content policy), 403 (e.g. akto) and 422 + (e.g. llm_as_a_judge). Other 4xx codes are commonly propagated from an + upstream guardrail provider response (401 bad key, 408 timeout, 429 rate + limit, or a raw upstream status), which are technical failures, not + blocks, so they stay guardrail_failed_to_respond. + """ + if isinstance(e, ModifyResponseException): + return True + if isinstance( + e, + ( + GuardrailRaisedException, + BlockedPiiEntityError, + SensitiveDataRouteException, + ), + ): + return True + if HTTPException is not None and isinstance(e, HTTPException) and e.status_code in _GUARDRAIL_BLOCK_STATUS_CODES: + return True + return False + + def _strict_guardrail_modes_enabled() -> bool: """Whether guardrail-mode validation raises (default) or logs a warning. @@ -429,11 +464,13 @@ class CustomGuardrail(CustomLogger): f"Sensitive data detected by {self.guardrail_name} (routing skipped: request has no session_id)" ), guardrail_name=self.guardrail_name, + blocked_content=True, ) else: raise GuardrailRaisedException( message=f"Sensitive data detected by {self.guardrail_name}", guardrail_name=self.guardrail_name, + blocked_content=True, ) @staticmethod @@ -1068,42 +1105,8 @@ class CustomGuardrail(CustomLogger): @staticmethod def _is_guardrail_intervention(e: Exception) -> bool: - """ - Returns True if the exception represents an intentional guardrail block - (this was logged previously as an API failure - guardrail_failed_to_respond). - - Guardrails signal intentional blocks by raising: - - GuardrailRaisedException (generic guardrail API, tool permission) - - BlockedPiiEntityError (Presidio PII detection) - - SensitiveDataRouteException (sensitive-data reroute to on-premise model) - - HTTPException with a block-signalling status (400, 403, 422) - - ModifyResponseException (passthrough mode violation) - - Only the statuses guardrails use in-tree to signal a deliberate rejection - count as an intervention: 400 (content policy), 403 (e.g. akto) and 422 - (e.g. llm_as_a_judge). Other 4xx codes are commonly propagated from an - upstream guardrail provider response (401 bad key, 408 timeout, 429 rate - limit, or a raw upstream status), which are technical failures, not - blocks, so they stay guardrail_failed_to_respond. - """ - if isinstance(e, ModifyResponseException): - return True - if isinstance( - e, - ( - GuardrailRaisedException, - BlockedPiiEntityError, - SensitiveDataRouteException, - ), - ): - return True - if ( - HTTPException is not None - and isinstance(e, HTTPException) - and e.status_code in _GUARDRAIL_BLOCK_STATUS_CODES - ): - return True - return False + """Retained spelling for existing callers; prefer ``is_guardrail_intervention``.""" + return is_guardrail_intervention(e) def _process_error( self, diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index b45513c5ea3..ee506a69ef9 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -4,10 +4,10 @@ CLI Token Utilities SDK-level utilities for reading the credential minted by `lite login`. Non-secret metadata lives in ~/.litellm/token.json. The secret material (the -bearer key, plus a JWT when one is issued) lives in the OS keychain when the -machine has one, and in that same 0600 file otherwise. This module hides the -split from callers, and migrates a legacy plaintext file into the keychain the -first time it reads one. +bearer key, the refresh token that renews it, and a JWT when one is issued) +lives in the OS keychain when the machine has one, and in that same 0600 file +otherwise. This module hides the split from callers, and migrates a plaintext +file into the keychain the first time it reads one. This module has no dependencies on proxy code and can be safely imported at the SDK level. """ @@ -111,13 +111,20 @@ class CliTokenSecret(BaseModel): secret minted for one server is never handed to another, even if the metadata file is edited underneath us. `timestamp` is the sign-in this secret came from, which is what decides it against a secret still on disk. + + Every field a thief could sign in with belongs here, which is why the + refresh token is one of them: it buys a fresh key from the proxy on demand, + so leaving it on disk would leave the login readable there. `key` is + optional because the file can hold a refresh token without one, and moving + that into the keychain must not invent a key to go with it. """ model_config = ConfigDict(frozen=True) base_url: str - key: str + key: str | None = None jwt_token: str = "" + refresh_token: str | None = None timestamp: float = 0.0 @@ -155,7 +162,7 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN staged: Final = _stage_token_file(_without_secret(stamped)) if isinstance(staged, CredentialNotSaved): return staged - outcome: Final = SecretStored() if stamped.key is None else vault.write(_encode_secret(stamped, stamped.key)) + outcome: Final = vault.write(_encode_secret(stamped)) if _holds_a_secret(stamped) else SecretStored() if isinstance(outcome, SecretStored): return outcome if _commit_token_file(staged) else CredentialNotRecorded() discard_staged_json(staged) @@ -390,8 +397,9 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - secret is usually left on disk by a keychain that would not take it, which makes the file the fresher of the two. It is the older one when a login the keychain did take could not replace the file afterwards, and serving that one would put a superseded credential back in use. Equal - stamps are one login sitting in both stores, left by a migration whose scrub was refused, so - that branch retries the migration rather than trading one credential for another. + stamps are one login sitting in both stores, left by a migration whose scrub was refused or by + an upgrade that took the key into the keychain and left the refresh token behind, so that branch + rejoins the halves and retries the migration rather than trading one credential for another. A scrub the file refuses leaves that superseded secret where it lies, which is the state the login already named when it could not replace the file, and which `lite logout` reports rather @@ -400,21 +408,46 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - superseded one back out. """ secret: Final = _decode_secret(blob, record.base_url) - if secret is None or (record.key is not None and secret.timestamp <= record.timestamp): - return _migrate_file_secret(record, vault) + if secret is None or (_holds_a_secret(record) and secret.timestamp <= record.timestamp): + return _migrate_file_secret(_rejoined(record, secret), vault, replacing=secret) _scrub_file_secret(record) return record.model_copy( update=MappingProxyType( { "key": secret.key, "jwt_token": secret.jwt_token, + "refresh_token": secret.refresh_token, "timestamp": max(secret.timestamp, record.timestamp), } ) ) -def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord | None: +def _rejoined(record: CliTokenRecord, secret: CliTokenSecret | None) -> CliTokenRecord: + """Put one sign-in's secret material back together when each store holds part of it. + + Upgrading from the release that kept only the key in the keychain leaves the refresh token + behind in the file, so a single login sits across both stores. Filling in whatever the file is + missing before the migration writes its entry is what stops that write from replacing a live key + with nothing. Only a matching stamp is one login. Two stamps are two logins, and pairing one's + key with the other's refresh token would build a credential neither store ever held. + """ + if secret is None or secret.timestamp != record.timestamp: + return record + return record.model_copy( + update=MappingProxyType( + { + "key": record.key if record.key is not None else secret.key, + "jwt_token": record.jwt_token or secret.jwt_token, + "refresh_token": record.refresh_token if record.refresh_token is not None else secret.refresh_token, + } + ) + ) + + +def _migrate_file_secret( + record: CliTokenRecord, vault: SecretVault, *, replacing: CliTokenSecret | None = None +) -> CliTokenRecord | None: """Move a file-held secret into the vault, but only once the file's copy can be taken away. The scrubbed file is staged first so a directory that will not accept it stops the migration @@ -426,23 +459,28 @@ def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliToken asked to take the new entry back, so the migration finishes on a directory that would only ever have refused it. Rolling back is the last resort, and a rollback the keychain also refuses leaves the secret in both stores until the next read, which retries this same migration. + + Only an entry this migration put there is taken back. `replacing` names one that was already in + the keychain, whose material the new entry carries forward, so erasing it would take away the + half the file never had, and a machine that refuses the scrub is exactly the one with nowhere + else to keep it. The next read finds the same two halves and tries the move again. """ - if record.key is None: + if not _holds_a_secret(record): return None staged: Final = _stage_scrubbed_file(record) if staged is None: return record - if not isinstance(vault.write(_encode_secret(record, record.key)), SecretStored): + if not isinstance(vault.write(_encode_secret(record)), SecretStored): discard_staged_json(staged) return record - if not _commit_token_file(staged) and not _overwrite_file_secret(record): + if not _commit_token_file(staged) and not _overwrite_file_secret(record) and replacing is None: vault.erase() return record def _scrub_file_secret(record: CliTokenRecord) -> bool: """Leave no secret material in the token file once the vault holds it""" - if record.key is None and not record.jwt_token: + if not _holds_a_secret(record): return True staged: Final = _stage_scrubbed_file(record) if staged is not None and _commit_token_file(staged): @@ -487,13 +525,22 @@ def _commit_token_file(staged: str) -> bool: return True +def _holds_a_secret(record: CliTokenRecord) -> bool: + """Whether the record carries anything that would sign someone in as this user""" + return record.key is not None or bool(record.jwt_token) or record.refresh_token is not None + + def _without_secret(record: CliTokenRecord) -> CliTokenRecord: - return record.model_copy(update=MappingProxyType({"key": None, "jwt_token": ""})) + return record.model_copy(update=MappingProxyType({"key": None, "jwt_token": "", "refresh_token": None})) -def _encode_secret(record: CliTokenRecord, key: str) -> str: +def _encode_secret(record: CliTokenRecord) -> str: return CliTokenSecret( - base_url=record.base_url, key=key, jwt_token=record.jwt_token, timestamp=record.timestamp + base_url=record.base_url, + key=record.key, + jwt_token=record.jwt_token, + refresh_token=record.refresh_token, + timestamp=record.timestamp, ).model_dump_json() diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6f352b73290..05fc6e07176 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -832,6 +832,10 @@ class LiteLLMRoutes(enum.Enum): # Team guardrail submissions - endpoint scopes results to caller's teams (non-admin) "/guardrails/submissions", "/guardrails/submissions/{guardrail_id}", + # Auto-router dry runs - both gate like the /model/new write they rehearse: + # proxy admin, or team admin naming their own team via team_id + "/auto_router/test_routing", + "/auto_router/validate_complexity_router_config", ] # routes that manage their own allowed/disallowed logic ## Org Admin Routes ## diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 3a39bec5ee6..1fff68677cc 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -331,7 +331,7 @@ sequenceDiagram CLI->>Proxy: Poll /sso/cli/poll/login_id with poll_secret header Proxy->>CLI: Return {"status": "ready", "key": "jwt"} - CLI->>CLI: Save key to the OS keychain (metadata to ~/.litellm/token.json) + CLI->>CLI: Save the secret to the OS keychain (metadata to ~/.litellm/token.json) ``` ### Authentication Commands @@ -365,7 +365,7 @@ The CLI provides these authentication commands: ### Token Storage -The key itself goes into the OS keychain (macOS Keychain, Windows Credential Manager, or the Linux Secret Service) under service `litellm-cli`, account `credential`. Only the non-secret session metadata is written to `~/.litellm/token.json`, in a `0700` directory with `0600` file permissions: +The key itself, together with the refresh token that renews a `--pkce` credential, goes into the OS keychain (macOS Keychain, Windows Credential Manager, or the Linux Secret Service) under service `litellm-cli`, account `credential`. Only the non-secret session metadata is written to `~/.litellm/token.json`, in a `0700` directory with `0600` file permissions: ```json { @@ -378,7 +378,7 @@ The key itself goes into the OS keychain (macOS Keychain, Windows Credential Man } ``` -Keychain storage needs the `keyring` package, which ships with `pip install 'litellm[cli]'`. Headless boxes and CI runners usually have no keychain either. In all of those cases the key stays in the same `0600` file alongside the metadata, exactly as it did before, and `lite login` names which one applies: the package is missing, the machine has no keychain, or you set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. +Keychain storage needs the `keyring` package, which ships with `pip install 'litellm[cli]'`. Headless boxes and CI runners usually have no keychain either. In all of those cases the key and the refresh token stay in the same `0600` file alongside the metadata, exactly as they did before, and `lite login` names which one applies: the package is missing, the machine has no keychain, or you set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. That includes a refresh token left behind by the release that moved only the key. `lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job. diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index e1c0653ebd3..c0f72af7576 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py @@ -356,6 +356,7 @@ class DeepKeepGuardrail(CustomGuardrail): guardrail_name=GUARDRAIL_NAME, message=error_message, should_wrap_with_default_message=False, + blocked_content=True, ) return self._build_return_inputs( diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 16768a4b08f..e3cf645ceaf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -464,6 +464,7 @@ class GenericGuardrailAPI(CustomGuardrail): guardrail_name=GUARDRAIL_NAME, message=error_message, should_wrap_with_default_message=False, + blocked_content=True, ) return self._build_guardrail_return_inputs( diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 4a20adf0e82..6644a3d3902 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -54,6 +54,7 @@ class OvalixGuardrailBlockedException(GuardrailRaisedException): guardrail_name=guardrail_name, message=message, should_wrap_with_default_message=should_wrap_with_default_message, + blocked_content=True, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index 4775a8b3caa..c25f704567e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -169,6 +169,7 @@ class PromptGuardGuardrail(CustomGuardrail): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=(f"Blocked by PromptGuard: {threat_type} (confidence={confidence}, event_id={event_id})"), + blocked_content=True, ) if decision == "redact": diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index cd9da8a58b7..3865ba4ed0e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -211,6 +211,7 @@ class SingulrGuardrail(CustomGuardrail): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=f"Blocked by Singulr: {result.blocking_due_to or 'unknown'}", + blocked_content=True, ) return inputs diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py index eee66f93b7a..7cca1ae2d63 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py @@ -536,12 +536,14 @@ class StraikerGuardrail(CustomGuardrail): request_data: dict, input_type: Literal["request", "response"], message: str, + blocked_content: bool = False, ) -> NoReturn: if input_type == "request": raise GuardrailRaisedException( guardrail_name=self.guardrail_name or GUARDRAIL_NAME, message=message, should_wrap_with_default_message=False, + blocked_content=blocked_content, ) raise ModifyResponseException( message=message, @@ -623,6 +625,7 @@ class StraikerGuardrail(CustomGuardrail): request_data=request_data, input_type=input_type, message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, + blocked_content=True, ) if parsed.action == "GUARDRAIL_INTERVENED": is_streamed_response: Final = input_type == "response" and _is_streamed_request(request_data) @@ -631,6 +634,7 @@ class StraikerGuardrail(CustomGuardrail): request_data=request_data, input_type=input_type, message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, + blocked_content=True, ) return self._intervened_inputs(inputs, parsed) return inputs diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 0514d2ab6f7..3c5625bc272 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -527,7 +527,9 @@ class ToolPermissionGuardrail(CustomGuardrail): if not is_allowed and message is not None: verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": - raise GuardrailRaisedException(guardrail_name=self.guardrail_name, message=message) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, message=message, blocked_content=True + ) return tuple( ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py index ee1aade8ea6..6b8148645aa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py @@ -205,6 +205,7 @@ class VigilGuardGuardrail(CustomGuardrail): guardrail_name=self.guardrail_name, message=self._build_block_reason(analysis), should_wrap_with_default_message=False, + blocked_content=True, ) if decision == "SANITIZED": @@ -245,6 +246,7 @@ class VigilGuardGuardrail(CustomGuardrail): guardrail_name=self.guardrail_name, message=self._build_block_reason(analysis), should_wrap_with_default_message=False, + blocked_content=True, ) if decision == "SANITIZED": diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 0112ad1f6ed..d47bd7fa311 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -2,6 +2,7 @@ AUTO ROUTER MANAGEMENT ENDPOINTS POST /auto_router/test_routing - Route one prompt through an unsaved complexity-router config +POST /auto_router/validate_complexity_router_config - Dry-run the complexity-router write gate without saving """ from collections.abc import Mapping, Sequence @@ -43,6 +44,8 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterCacheStats, AutoRouterRoutingTestRequest, AutoRouterRoutingTestResponse, + ComplexityRouterConfigValidationRequest, + ComplexityRouterConfigValidationResponse, RequestComplexityRouterConfig, ShadowEvalDirection, ShadowEvalJobKeyResponse, @@ -130,12 +133,13 @@ async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) - return await prisma_client.db.query_raw(query, *args) -async def _authorize_routing_test(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> None: +async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> None: """Allow exactly the callers who could create this router. - Routing a prompt can spend money (an `llm` classifier config calls its classifier, a - semantic config embeds the prompt), so this is gated like a write rather than a read: - a proxy admin, or a team admin naming their own team, matching /model/new. + Both dry runs are gated like the write they rehearse rather than as reads: a proxy + admin, or a team admin naming their own team, matching /model/new. Routing a test + prompt can also spend money (an `llm` classifier config calls its classifier, a + semantic config embeds the prompt), so a read-level gate would be too loose anyway. """ from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, @@ -149,7 +153,7 @@ async def _authorize_routing_test(user_api_key_dict: UserAPIKeyAuth, team_id: st raise HTTPException( status_code=403, detail={ # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape - "error": f"User does not have permission to test an auto router. Your role={user_api_key_dict.user_role}. Test as a PROXY_ADMIN, or as a team admin by specifying a team_id." + "error": f"User does not have permission to dry-run an auto router. Your role={user_api_key_dict.user_role}. Call as a PROXY_ADMIN, or as a team admin by specifying a team_id." }, ) @@ -238,6 +242,35 @@ async def _authorize_models_this_test_can_call( ) from e +@router.post( + "/auto_router/validate_complexity_router_config", + tags=["model management"], # mutable-ok: fastapi's decorator signature types tags as a list + dependencies=[Depends(user_api_key_auth)], # mutable-ok: fastapi's decorator signature types dependencies as a list + response_model=ComplexityRouterConfigValidationResponse, + status_code=status.HTTP_200_OK, +) +async def validate_complexity_router_config( + data: ComplexityRouterConfigValidationRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> ComplexityRouterConfigValidationResponse: + """ + Validate a complexity-router config without saving it. + + Runs the same check every write path runs (the router's own pydantic model), so a form can + show the backend's exact verdict while the operator is still editing rather than after a + rejected save. Gated exactly like the save it rehearses: a proxy admin, or a team admin + naming their own team. Nothing is created, routed, or billed. + """ + await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) + + from litellm.router_utils.auto_router_model_naming import ( + validate_complexity_router_config_write, + ) + + error: Final = validate_complexity_router_config_write(data.complexity_router_config) + return ComplexityRouterConfigValidationResponse(valid=error is None, error=error) + + @router.post( "/auto_router/test_routing", tags=["model management"], # mutable-ok: fastapi's decorator signature types tags as a list @@ -270,9 +303,17 @@ async def preview_auto_router_routing( } ``` """ - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + user_model, + ) + from litellm.proxy.utils import get_available_models_for_user - await _authorize_routing_test(user_api_key_dict=user_api_key_dict, team_id=data.team_id) + await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) if llm_router is None: raise HTTPException( @@ -327,9 +368,19 @@ async def preview_auto_router_routing( }, ) + available_models: Final = await get_available_models_for_user( + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + general_settings=general_settings, + user_model=user_model, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + team_id=data.team_id, + user_api_key_cache=user_api_key_cache, + ) return AutoRouterRoutingTestResponse( routed_model=hook_response.model, - routed_model_configured=hook_response.model in frozenset(llm_router.get_model_names()), + routed_model_configured=hook_response.model in frozenset(available_models), routing_decision=hook_response.routing_decision, ) diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py new file mode 100644 index 00000000000..5c886ca0e9b --- /dev/null +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -0,0 +1,548 @@ +""" +Run the configured pre-call guardrails over every record of a batch input file. + +Runs after ``batch_file_validation.check_batch_file_upload``, so every line here is already known +to parse as a JSON object carrying ``custom_id``, ``method``, ``url`` and ``body``. +""" + +from __future__ import annotations + +import asyncio +import copy +import json +import re +import tempfile +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, BinaryIO, Final, NoReturn, TypeAlias +from urllib.parse import urlsplit + +from fastapi import HTTPException +from typing_extensions import assert_never + +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import is_guardrail_intervention +from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport +from litellm.types.utils import CallTypes, CallTypesLiteral + +if TYPE_CHECKING: + from litellm.proxy.utils import ProxyLogging + +EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) + +_SCAN_WINDOW: Final = 32 + +# Past this the rewrite rolls to disk, keeping the router's per-deployment deepcopy of the handle +# as cheap as it is for the spooled upload this replaces. +_REWRITE_SPOOL_BYTES: Final = 1024 * 1024 + +# custom_id is caller-supplied and reaches a log line, so it is stripped of control characters +# and capped rather than rendered as given. +_CONTROL_CHARACTERS: Final = re.compile(r"[\x00-\x1f\x7f]") +_CUSTOM_ID_LOG_LIMIT: Final = 128 +_SUMMARY_LIMIT: Final = 50 + +_SCAN_METADATA_KEY: Final = "litellm_metadata" +_SCAN_METADATA_BAGS: Final = (_SCAN_METADATA_KEY, "metadata") + +# Set by pre_call_hook when a guardrail rerouted the request to a different model. +_ROUTE_APPLIED_KEY: Final = "sensitive_data_routing_applied" + +# Dropped before dispatch and restored afterwards rather than diffed. Guardrail dispatch writes +# its bookkeeping into `metadata`, and a record's own metadata is not scanned content on the +# online path either. `guardrails` is dropped because guardrail selection reads it ahead of the +# proxy-injected list, so leaving it would let a record's own body opt out of the chain its key +# and team selected; online that key can only add to the list, never replace it. +_INJECTED_KEYS: Final = frozenset({_SCAN_METADATA_KEY, "metadata", "guardrails"}) + +# Only what guardrail dispatch reads. The parent OTel span is deliberately left out: parenting one +# guardrail span per record would put tens of thousands of spans on a single upload's trace. +_SCAN_METADATA_KEYS: Final = frozenset( + { + "guardrails", + "_guardrail_pipelines", + "_pipeline_managed_guardrails", + "user_api_key_metadata", + "user_api_key_team_metadata", + "tags", + "headers", + } +) + +_SCANNABLE_CALL_TYPES: Final = frozenset( + { + CallTypes.acompletion, + CallTypes.atext_completion, + CallTypes.aembedding, + CallTypes.aresponses, + CallTypes.anthropic_messages, + } +) + +# Mirrors the record classifier in litellm/llms/bedrock/files/transformation.py, so a record +# litellm already accepts without a url keeps working. +_BODY_SHAPE_CALL_TYPES: Final = ( + ("messages", CallTypes.acompletion), + ("prompt", CallTypes.atext_completion), + ("input", CallTypes.aembedding), +) + + +@dataclass(frozen=True, slots=True) +class UnparseableRecord: + line_number: int + + +@dataclass(frozen=True, slots=True) +class UnscannableRecord: + line_number: int + custom_id: str | None + url: str | None + + +@dataclass(frozen=True, slots=True) +class UnroutableRecord: + line_number: int + custom_id: str | None + guardrail: str | None + + +BatchScanFailure: TypeAlias = UnparseableRecord | UnscannableRecord | UnroutableRecord + + +@dataclass(frozen=True, slots=True) +class _Redaction: + """A rewritten record on its way to the scan spool, held only for the window it was scanned in.""" + + line_number: int + custom_id: str | None + text: str + + +@dataclass(frozen=True, slots=True) +class RecordRedacted: + line_number: int + custom_id: str | None + offset: int + length: int + """Where the re-serialized record sits in the scan spool, so a large file's rewrites stay off the heap.""" + + +@dataclass(frozen=True, slots=True) +class RecordDropped: + line_number: int + custom_id: str | None + guardrail: str | None = None + + +_RecordChange: TypeAlias = RecordRedacted | RecordDropped +_ScanOutcome: TypeAlias = BatchScanFailure | _Redaction | RecordDropped + + +@dataclass(frozen=True, slots=True) +class BatchScanResult: + """What the scan decided, per record. Empty changes means the upload proceeds untouched.""" + + changes: tuple[_RecordChange, ...] + scanned_records: int + redactions: BinaryIO + """Spool holding every rewritten record, keyed by the offsets on each ``RecordRedacted``.""" + + @property + def submitted_records(self) -> int: + return self.scanned_records - sum(1 for change in self.changes if isinstance(change, RecordDropped)) + + def summary(self) -> str: + """Compact per-record outcome for the server-side log line, capped so one upload cannot flood it.""" + shown: Final = ", ".join( + f"line {change.line_number}{_describe(change.custom_id)} " + f"{'redacted' if isinstance(change, RecordRedacted) else 'dropped'}" + for change in self.changes[:_SUMMARY_LIMIT] + ) + remaining: Final = len(self.changes) - _SUMMARY_LIMIT + return shown if remaining <= 0 else f"{shown}, and {remaining} more" + + def report(self) -> BatchGuardrailReport: + return BatchGuardrailReport( + submitted_records=self.submitted_records, + modified_records=tuple( + BatchGuardrailRecord( + line=change.line_number, + custom_id=change.custom_id, + action="redacted" if isinstance(change, RecordRedacted) else "dropped", + guardrail=change.guardrail if isinstance(change, RecordDropped) else None, + ) + for change in self.changes + ), + ) + + +@dataclass(frozen=True, slots=True) +class _ParsedRecord: + line_number: int + payload: Mapping[str, object] + + +def _rejected(message: str) -> HTTPException: + return HTTPException(status_code=400, detail={"error": message}) # mutable-ok: FastAPI detail shape + + +def raise_public(failure: BatchScanFailure) -> NoReturn: + """Map a scan failure onto the 400 contract the files endpoint already returns.""" + match failure: + case UnparseableRecord(line_number=line_number): + raise _rejected( + f"The 'body' of batch input line {line_number} is not an object, so guardrails cannot be applied to it" + ) + case UnscannableRecord(line_number=line_number, custom_id=custom_id, url=url): + raise _rejected( + f"Batch input line {line_number}{_describe(custom_id)} targets {url or 'no url'} " + "and its body has no messages, prompt or input, so guardrails cannot read it. " + "Give the record a chat, completion, embedding, responses or messages body" + ) + case UnroutableRecord(line_number=line_number, custom_id=custom_id, guardrail=guardrail): + raise _rejected( + f"Batch input line {line_number}{_describe(custom_id)} was routed to a different model by " + f"{guardrail or 'a guardrail'}, and every record of a batch file goes to one provider, so " + "the file cannot be submitted. Send that record outside the batch" + ) + case _: + assert_never(failure) + + +def raise_nothing_to_submit() -> NoReturn: + """Every record was blocked, so there is no batch left to create.""" + raise _rejected( + "Every record in the batch input file was blocked by a guardrail, so there is nothing left to submit" + ) + + +def _is_content_block(exc: BaseException) -> bool: + """ + Whether the guardrail judged the record, as opposed to failing to judge it. + + Stricter than ``is_guardrail_intervention``, which answers a different question and counts + every ``GuardrailRaisedException`` as a block. Several integrations raise that same exception + for an unreachable backend or an unparseable response, and only when the operator configured + the guardrail to fail closed, so treating it as a block would turn "refuse this request" into + "drop this record and submit the rest", which is the silent loss of enforcement this whole + path exists to prevent. A guardrail that does not say it blocked content aborts the upload. + + Guardrails that report a technical failure as an ``HTTPException`` carrying a block status + are caught by ``__cause__``: raising ``from`` the underlying error is a deliberate statement + that something else caused this, which a verdict on content never is. Implicit context is + left alone, since a block raised inside an unrelated ``except`` would read as a failure. + """ + if isinstance(exc, GuardrailRaisedException): + return exc.blocked_content + if exc.__cause__ is not None: + return False + return is_guardrail_intervention(exc) + + +def _naming_guardrail(exc: BaseException) -> str | None: + """The guardrail that raised, from whichever place it recorded its own name.""" + named: Final = getattr(exc, "guardrail_name", None) + if isinstance(named, str): + return named + detail: Final = getattr(exc, "detail", None) + enriched: Final = detail.get("guardrail_name") if isinstance(detail, dict) else None + return enriched if isinstance(enriched, str) else None + + +def _describe(custom_id: str | None) -> str: + if not custom_id: + return "" + safe: Final = _CONTROL_CHARACTERS.sub(" ", custom_id)[:_CUSTOM_ID_LOG_LIMIT] + return f" (custom_id {safe})" + + +def _iter_lines(source: BinaryIO) -> Iterator[tuple[int, str]]: + """Yield every non-blank line with its 1-based number, so both passes number records alike.""" + for line_number, raw_line in enumerate(source, start=1): + text = raw_line.decode("utf-8") + if text.strip(): + yield line_number, text + + +def _iter_records(source: BinaryIO) -> Iterator[_ParsedRecord]: + """Yield one record per line, relying on the upload validation that already ran.""" + for line_number, text in _iter_lines(source): + yield _ParsedRecord(line_number=line_number, payload=json.loads(text)) + + +def _call_type_from_url(url: str) -> CallTypesLiteral | None: + """ + Resolve the route a record names, tolerating how callers actually write it. + + An absolute url has to reduce to its path or nothing matches, and a record naming + ``/v1/responses`` in full would fall through to its body, where ``input`` reads as an + embedding and the record gets scanned as the wrong call type rather than the right one. + """ + path: Final = urlsplit(url).path.split("?")[0].rstrip("/") + call_types: Final = get_call_types_for_route(path) + if call_types is None: + return None + scannable: Final = next((c for c in call_types if c in _SCANNABLE_CALL_TYPES), None) + return None if scannable is None else scannable.value + + +def _call_type_from_body(body: Mapping[str, object]) -> CallTypesLiteral | None: + shape: Final = next((call_type for field, call_type in _BODY_SHAPE_CALL_TYPES if field in body), None) + return None if shape is None else shape.value + + +def _scannable_call_type(url: object, body: Mapping[str, object]) -> CallTypesLiteral | None: + """ + Resolve how to scan a record: its url when we recognize one, otherwise its body shape. + + An unrecognized url falls through to the body rather than rejecting, because a record we can + still read is a record we can still scan, and the provider transformers treat an unknown url + as chat rather than as an error. + """ + from_url: Final = _call_type_from_url(url) if isinstance(url, str) and url else None + return from_url if from_url is not None else _call_type_from_body(body) + + +def _custom_id_of(payload: Mapping[str, object]) -> str | None: + custom_id: Final = payload.get("custom_id") + return custom_id if isinstance(custom_id, str) else None + + +def _fingerprint(body: Mapping[str, object], keys: frozenset[str]) -> str: + """ + Order-insensitive projection, so a guardrail re-serializing a dict does not read as a change. + + An absent key projects to ``null`` while a key holding ``None`` projects to the string + ``"null"``, so adding or dropping a null-valued key still reads as a change. + """ + return json.dumps( + tuple( + (key, json.dumps(body[key], sort_keys=True, default=str) if key in body else None) for key in sorted(keys) + ) + ) + + +def build_scan_metadata(request_metadata: Mapping[str, object]) -> Mapping[str, object]: + """ + Narrow the request metadata to the keys guardrail dispatch reads. + + Passing the whole thing through would carry values that cannot be copied, such as the parent + OTel span, and would hand every record proxy state it has no business seeing. + """ + return MappingProxyType( + {key: value for key, value in request_metadata.items() if key in _SCAN_METADATA_KEYS} + ) # mutable-ok: MappingProxyType freezes the comprehension + + +async def _scan_record( + record: _ParsedRecord, + scan_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> _ScanOutcome | None: + body: Final = record.payload.get("body") + if not isinstance(body, dict): + return UnparseableRecord(line_number=record.line_number) + + custom_id: Final = _custom_id_of(record.payload) + url: Final = record.payload.get("url") + call_type: Final = _scannable_call_type(url, body) + if call_type is None: + return UnscannableRecord( + line_number=record.line_number, + custom_id=custom_id, + url=url if isinstance(url, str) else None, + ) + + scan_input: Final[dict[str, object]] = copy.deepcopy(body) # mutable-ok: pre_call_hook mutates the dict it is given + own_injected: Final = MappingProxyType({key: body[key] for key in _INJECTED_KEYS if key in body}) + for injected in _INJECTED_KEYS: + scan_input.pop(injected, None) + # Both bags, because guardrails read whichever one their own route populates and a record + # scanned as chat reaches ones that only ever look at `metadata`; both are injected keys, so + # neither survives into the record that ships. Deep, and per bag per record, because `headers` + # and `tags` are nested containers otherwise shared with the upload request and with every + # other record in the window. The narrowing above already removed what cannot be copied. + for injected in _SCAN_METADATA_BAGS: + scan_input[injected] = copy.deepcopy(dict(scan_metadata)) # mutable-ok: guardrails write here + + try: + # The chain hands back the body it produced, which may be a replacement for the dict it was + # given rather than that same dict mutated, so this is what gets compared. + scanned: Final[dict] = await proxy_logging_obj.pre_call_hook( # mutable-ok: the guardrails' own dict + user_api_key_dict=user_api_key_dict, + data=scan_input, + call_type=call_type, + guardrails_only=True, + ) + except Exception as exc: + if _is_content_block(exc): + return RecordDropped(line_number=record.line_number, custom_id=custom_id, guardrail=_naming_guardrail(exc)) + raise + + rerouted: Final = scanned.get("metadata") + if isinstance(rerouted, dict) and rerouted.get(_ROUTE_APPLIED_KEY): + return UnroutableRecord( + line_number=record.line_number, + custom_id=custom_id, + guardrail=rerouted.get("sensitive_data_routing_guardrail"), + ) + + compared: Final = (frozenset(body) | frozenset(scanned)) - _INJECTED_KEYS + if _fingerprint(scanned, compared) == _fingerprint(body, compared): + return None + for injected in _INJECTED_KEYS: + scanned.pop(injected, None) + scanned.update(own_injected) + return _Redaction( + line_number=record.line_number, + custom_id=custom_id, + text=json.dumps({**record.payload, "body": scanned}), # mutable-ok: json.dumps needs a plain dict + ) + + +async def _scan_window( + window: tuple[_ParsedRecord, ...], + scan_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> tuple[tuple[int, _ScanOutcome | BaseException], ...]: + """``return_exceptions=True`` so one record raising never leaves its siblings unobserved.""" + outcomes: Final = await asyncio.gather( + *(_scan_record(record, scan_metadata, user_api_key_dict, proxy_logging_obj) for record in window), + return_exceptions=True, + ) + return tuple((record.line_number, outcome) for record, outcome in zip(window, outcomes) if outcome is not None) + + +def _spool(redactions: BinaryIO, redaction: _Redaction) -> RecordRedacted: + """Park the rewritten record on disk so only its location is carried for the rest of the scan.""" + encoded: Final = redaction.text.encode("utf-8") + redactions.seek(0, 2) + offset: Final = redactions.tell() + redactions.write(encoded) + return RecordRedacted( + line_number=redaction.line_number, + custom_id=redaction.custom_id, + offset=offset, + length=len(encoded), + ) + + +def _worst(problems: tuple[tuple[int, BatchScanFailure | BaseException], ...]) -> BatchScanFailure | BaseException: + """A guardrail that blocked outranks a record we merely refused; then earliest line wins.""" + raised: Final = tuple(problem for problem in problems if isinstance(problem[1], BaseException)) + return min(raised or problems, key=lambda problem: problem[0])[1] + + +async def scan_batch_input_file( + *, + file_source: BinaryIO, + request_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> BatchScanFailure | BatchScanResult: + """ + Stream a batch input file and run the pre-call guardrail chain against every record. + + A record a guardrail rewrites is kept in its rewritten form and a record it blocks is dropped, + which is what the online path does per request. Both are returned for reporting. A guardrail + exception that is not a block is re-raised untouched so its status code survives, since dropping + a record that was never inspected is worse than refusing the file. + """ + scan_metadata: Final = build_scan_metadata(request_metadata) + problems: Final[list[tuple[int, BatchScanFailure | BaseException]]] = [] # mutable-ok: spans windows + changes: Final[list[_RecordChange]] = [] # mutable-ok: accumulates across windows + window: Final[list[_ParsedRecord]] = [] # mutable-ok: bounded read-ahead buffer + scanned: Final[list[int]] = [] # mutable-ok: counts records the scan actually reached + redactions: Final = tempfile.SpooledTemporaryFile( # noqa: SIM115 # the rewrite reads this back + max_size=_REWRITE_SPOOL_BYTES + ) + + async def drain() -> None: + if window: + scanned.append(len(window)) + for line_number, outcome in await _scan_window( + tuple(window), scan_metadata, user_api_key_dict, proxy_logging_obj + ): + if isinstance(outcome, _Redaction): + changes.append(_spool(redactions, outcome)) + elif isinstance(outcome, RecordDropped): + changes.append(outcome) + else: + problems.append((line_number, outcome)) + window.clear() + + try: + for item in _iter_records(file_source): + window.append(item) + if len(window) >= _SCAN_WINDOW: + await drain() + if problems: + break + if not problems: + await drain() + except BaseException: + redactions.close() + raise + finally: + file_source.seek(0) + + if problems: + redactions.close() + worst: Final = _worst(tuple(problems)) + if isinstance(worst, BaseException): + raise worst + return worst + if not changes: + redactions.close() + return BatchScanResult( + changes=tuple(sorted(changes, key=lambda change: change.line_number)), + scanned_records=sum(scanned), + redactions=redactions, + ) + + +def _read_spooled(redactions: BinaryIO, change: RecordRedacted) -> str: + redactions.seek(change.offset) + return redactions.read(change.length).decode("utf-8") + + +def rewrite_batch_input_file(file_source: BinaryIO, result: BatchScanResult) -> BinaryIO: + """ + Re-emit the file with redacted records rewritten and dropped records left out. + + Untouched records are copied through as written rather than re-serialized, so enabling the + feature does not reformat records no guardrail objected to. Blank lines between records are + not carried over, since they are not records. Rewritten records are read back from the scan's + spool rather than from memory, so a file whose records are mostly rewritten does not put a + second copy of itself on the heap. + """ + redacted: Final = MappingProxyType( + {change.line_number: change for change in result.changes if isinstance(change, RecordRedacted)} + ) # mutable-ok: MappingProxyType freezes the lookup table + dropped: Final = frozenset(change.line_number for change in result.changes if isinstance(change, RecordDropped)) + + output: Final = tempfile.SpooledTemporaryFile( # noqa: SIM115 # the caller uploads this handle + max_size=_REWRITE_SPOOL_BYTES + ) + wrote_any = False # rebind-ok: tracks whether a separator is needed + try: + for line_number, text in _iter_lines(file_source): + if line_number in dropped: + continue + change = redacted.get(line_number) + line = text.rstrip("\n") if change is None else _read_spooled(result.redactions, change) + output.write((("\n" if wrote_any else "") + line).encode("utf-8")) + wrote_any = True + except BaseException: + output.close() + raise + finally: + file_source.seek(0) + output.seek(0) + return output diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index b7200de8fb6..37cfd9d073d 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,6 +7,7 @@ import asyncio import traceback +from collections.abc import Mapping from typing import Any, BinaryIO, Final, cast, get_args import httpx @@ -29,6 +30,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_cloud_storage_uri, ) +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -46,6 +48,14 @@ from litellm.proxy.openai_files_endpoints.batch_file_validation import ( check_batch_file_upload, raise_batch_file_validation_failure, ) +from litellm.proxy.openai_files_endpoints.batch_guardrails import ( + EMPTY_MAPPING, + BatchScanResult, + raise_nothing_to_submit, + raise_public, + rewrite_batch_input_file, + scan_batch_input_file, +) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, add_internal_model_credentials, @@ -106,6 +116,34 @@ def get_files_provider_config( return None +async def _scan_batch_upload( + *, + file_source: bytes | BinaryIO, + purpose: str, + request_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> BatchScanResult | None: + """Guardrail the records of a batch input file, or None when this upload has nothing to scan.""" + if ( + purpose != "batch" + or isinstance(file_source, bytes) + or not proxy_logging_obj.has_pre_call_guardrails(request_metadata) + ): + return None + outcome: Final = await scan_batch_input_file( + file_source=file_source, + request_metadata=request_metadata, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + if not isinstance(outcome, BatchScanResult): + raise_public(outcome) + if outcome.changes and outcome.submitted_records == 0: + raise_nothing_to_submit() + return outcome + + def get_first_json_object(file_source: bytes | BinaryIO) -> dict | None: try: if isinstance(file_source, (bytes, bytearray)): @@ -333,6 +371,10 @@ async def create_file( ) data: dict = {} + # Spools this request owns. Starlette owns the upload handle; anything the guardrail scan + # opens is ours, and a batch upload that fails after the scan would otherwise hold the + # descriptor and its disk blocks until the collector runs. + spools: Final[list[BinaryIO]] = [] # mutable-ok: filled as the scan opens handles try: # Batch uploads can be gigabytes. Starlette has already spooled the upload # to disk, so stream from that handle instead of reading it into memory. @@ -471,14 +513,44 @@ async def create_file( proxy_config=proxy_config, ) + # /v1/files stores its proxy metadata under litellm_metadata, not metadata + request_metadata: Final = data.get("metadata") or data.get("litellm_metadata") or EMPTY_MAPPING + scan_result: Final = await _scan_batch_upload( + file_source=file_source, + purpose=purpose, + request_metadata=request_metadata, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + if scan_result is not None and scan_result.changes: + # The caller sees this in the response; a proxy admin needs it server side too, + # and it has to land before the post-call hook for logging callbacks to pick it up. + get_or_create_metadata_bucket(data)[1]["batch_guardrail"] = scan_result.report().model_dump() + verbose_proxy_logger.warning( + "batch guardrails changed %s of %s records in %s: %s", + len(scan_result.changes), + scan_result.scanned_records, + file.filename, + scan_result.summary(), + ) + # Prepare the file data according to FileTypes - file_data: Final = (file.filename, file_source, file.content_type) + if scan_result is not None: + spools.append(scan_result.redactions) + upload_source: Final = ( + await asyncio.to_thread(rewrite_batch_input_file, file_source, scan_result) + if scan_result is not None and scan_result.changes + else file_source + ) + if upload_source is not file_source: + spools.append(upload_source) + file_data: Final = (file.filename, upload_source, file.content_type) ## check if model is a loadbalanced model router_model: str | None = None is_router_model = False if litellm.enable_loadbalancing_on_batch_endpoints is True: - json_obj: Final = get_first_json_object(file_source) + json_obj: Final = get_first_json_object(upload_source) if json_obj: router_model = get_model_from_json_obj(json_object=json_obj) is_router_model = is_known_model(model=router_model, llm_router=llm_router) @@ -546,6 +618,9 @@ async def create_file( if _response is not None and isinstance(_response, OpenAIFileObject): response = _response + if scan_result is not None and scan_result.changes: + response.litellm_batch_guardrail = scan_result.report() + ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} model_id: Final = hidden_params.get("model_id", None) or "" @@ -585,6 +660,9 @@ async def create_file( param=getattr(e, "param", "None"), code=getattr(e, "status_code", 500), ) + finally: + for spool in spools: + spool.close() @router.get( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e0546ffb070..12d0f13ebdd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15255,6 +15255,29 @@ def get_logo_url(): return {"logo_url": ""} +def _serve_custom_ui_logo(candidate: str) -> Response | None: + """Serve one admin-configured logo, or None when it is unusable so the caller falls back.""" + from litellm.proxy.common_utils.static_asset_utils import ( + resolve_validated_local_image_path, + ) + + # Remote logo URLs are loaded by the browser. The proxy should not fetch + # arbitrary admin-configured URLs server-side. + if candidate.startswith(("http://", "https://")): + return RedirectResponse(url=candidate) + + safe_logo: Final = resolve_validated_local_image_path(candidate) + if safe_logo is None: + verbose_proxy_logger.warning( + "Custom UI logo %r is not a supported image file or does not exist, falling back", + candidate, + ) + return None + + safe_logo_path, media_type = safe_logo + return FileResponse(safe_logo_path, media_type=media_type) + + @app.get("/get_image", include_in_schema=False) async def get_image(theme: Literal["light", "dark"] | None = None): """Get logo to show on admin UI""" @@ -15293,31 +15316,33 @@ async def get_image(theme: Literal["light", "dark"] | None = None): if assets_dir != current_dir and not os.path.exists(default_logo): default_logo = default_site_logo - logo_path = os.getenv("UI_LOGO_PATH", default_logo) - verbose_proxy_logger.debug("Reading logo from path: %s", logo_path) + custom_logo_candidates: Final = tuple( + candidate.strip() + for candidate in ( + os.getenv("UI_LOGO_PATH_DARK", "") if theme == "dark" else "", + os.getenv("UI_LOGO_PATH", ""), + ) + if candidate.strip() + ) + verbose_proxy_logger.debug("Custom logo candidates, in fallback order: %s", custom_logo_candidates) + + custom_logo_response: Final = next( + ( + response + for response in (_serve_custom_ui_logo(candidate) for candidate in custom_logo_candidates) + if response is not None + ), + None, + ) + if custom_logo_response is not None: + return custom_logo_response from litellm.proxy.common_utils.static_asset_utils import ( resolve_validated_local_image_path, ) - if logo_path != default_logo and not logo_path.startswith(("http://", "https://")): - safe_logo = resolve_validated_local_image_path(logo_path) - if safe_logo is not None: - safe_logo_path, media_type = safe_logo - return FileResponse(safe_logo_path, media_type=media_type) - verbose_proxy_logger.warning( - "UI_LOGO_PATH %r is not a supported image file or does not exist, falling back to default logo", - logo_path, - ) - logo_path = default_logo - - # Remote logo URLs are loaded by the browser. The proxy should not fetch - # arbitrary admin-configured URLs server-side. - if logo_path.startswith(("http://", "https://")): - return RedirectResponse(url=logo_path) - # Default logo (resolved from the bundled asset, not user-controlled). - safe_logo = resolve_validated_local_image_path(logo_path) + safe_logo: Final = resolve_validated_local_image_path(default_logo) if safe_logo is not None: safe_logo_path, media_type = safe_logo return FileResponse(safe_logo_path, media_type=media_type) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 5584dae9e15..66a8c0622fa 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -110,6 +110,7 @@ def _config_param_db(repo: _HasConfigParamTable) -> _PrismaTableActions[_ConfigP # reflect a deployment branded purely through process env. _UI_THEME_FIELD_ENV_VARS: Final[dict[str, str]] = { "logo_url": "UI_LOGO_PATH", + "logo_url_dark": "UI_LOGO_PATH_DARK", "favicon_url": "LITELLM_FAVICON_URL", } @@ -156,6 +157,14 @@ class UIThemeConfig(BaseModel): description="URL or path to custom logo image. Can be a local file path or HTTP/HTTPS URL", ) + logo_url_dark: str | None = Field( + default=None, + description=( + "URL or path to a custom logo image for dark mode. Can be a local file path or HTTP/HTTPS URL. " + "Leave unset to reuse logo_url in dark mode" + ), + ) + # Favicon configuration favicon_url: str | None = Field( default=None, @@ -1184,6 +1193,7 @@ async def update_ui_theme_settings( ) _validate_public_image_url(theme_config.logo_url, "logo_url") + _validate_public_image_url(theme_config.logo_url_dark, "logo_url_dark") _validate_public_image_url(theme_config.favicon_url, "favicon_url") if store_model_in_db is not True: @@ -1204,16 +1214,18 @@ async def update_ui_theme_settings( config["litellm_settings"] = {} config["litellm_settings"]["ui_theme_config"] = theme_data - # UI_LOGO_PATH and LITELLM_FAVICON_URL are the only environment variables - # this endpoint owns. A non-empty value sets the var; an empty or missing - # one clears it back to the default. Apply to the live process immediately, - # then persist only these two keys so an unrelated env var (a YAML/OS value - # merged in by get_config) is never snapshotted into the DB. + # The vars below are the only environment variables this endpoint owns, and + # they must stay in step with _UI_THEME_FIELD_ENV_VARS. A non-empty value + # sets the var; an empty or missing one clears it back to the default. Apply + # to the live process immediately, then persist only those keys so an + # unrelated env var (a YAML/OS value merged in by get_config) is never + # snapshotted into the DB. def _clean(url: str | None) -> str | None: return url if url is not None and url.strip() else None env_updates: Final[dict[str, str | None]] = { "UI_LOGO_PATH": _clean(theme_config.logo_url), + "UI_LOGO_PATH_DARK": _clean(theme_config.logo_url_dark), "LITELLM_FAVICON_URL": _clean(theme_config.favicon_url), } for env_key, env_value in env_updates.items(): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index eb04a862c5a..81a86ebe34d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1527,6 +1527,23 @@ class ProxyLogging: return data + def has_pre_call_guardrails(self, request_metadata: Mapping[str, object]) -> bool: + """ + Whether any guardrail or guardrail pipeline would inspect a request carrying this metadata. + + Evaluated with the same predicate the pre-call loop uses, so a proxy configured only with + post-call guardrails answers False. Callers that must pay a real cost to build the hook's + input, such as streaming a batch input file off disk, use this to skip that work. + """ + if request_metadata.get("_guardrail_pipelines"): + return True + probe: Final = {"metadata": dict(request_metadata)} # mutable-ok: should_run_guardrail takes a dict + return any( + isinstance(callback, CustomGuardrail) + and callback.should_run_guardrail(data=probe, event_type=GuardrailEventHooks.pre_call) + for callback in ProxyLogging._callback_capabilities().resolved_callbacks + ) + # The actual implementation of the function @overload async def pre_call_hook( @@ -1534,6 +1551,7 @@ class ProxyLogging: user_api_key_dict: UserAPIKeyAuth, data: None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ) -> None: pass @@ -1543,6 +1561,7 @@ class ProxyLogging: user_api_key_dict: UserAPIKeyAuth, data: dict, call_type: CallTypesLiteral, + guardrails_only: bool = False, ) -> dict: pass @@ -1551,6 +1570,7 @@ class ProxyLogging: user_api_key_dict: UserAPIKeyAuth, data: dict | None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ) -> dict | None: """ Allows users to modify/reject the incoming request to the proxy, without having to deal with parsing Request body. @@ -1559,10 +1579,15 @@ class ProxyLogging: 1. /chat/completions 2. /embeddings 3. /image/generation + + With ``guardrails_only`` the walk is limited to guardrails and guardrail pipelines: rate + limiting, budget accounting, prompt templates and hanging-request alerting are skipped. + Use it to scan a payload that is not itself a request, such as one record of a batch file. """ verbose_proxy_logger.debug("Inside Proxy Logging Pre-call hook!") - self._init_response_taking_too_long_task(data=data) + if not guardrails_only: + self._init_response_taking_too_long_task(data=data) if data is None: return None @@ -1574,7 +1599,8 @@ class ProxyLogging: ## PROMPT TEMPLATE CHECK ## if ( - litellm_logging_obj is not None + not guardrails_only + and litellm_logging_obj is not None and prompt_id is not None and (call_type == "completion" or call_type == "acompletion") ): @@ -1605,7 +1631,7 @@ class ProxyLogging: # CustomGuardrail is configured. Saves the loop overhead + # ``time.time()`` x2 per registered callback for the common # "callbacks=[]" case on small / dev deployments. - if not caps.has_guardrail and not caps.has_pre_call_override: + if not caps.has_guardrail and (guardrails_only or not caps.has_pre_call_override): if data is not None: self._process_guardrail_metadata(data) return data @@ -1642,7 +1668,8 @@ class ProxyLogging: data = result elif ( - _callback is not None + not guardrails_only + and _callback is not None and isinstance(_callback, CustomLogger) and "async_pre_call_hook" in vars(_callback.__class__) and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 6457b285cb5..1588c650177 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -279,6 +279,40 @@ OpenAIFilesPurpose = Literal[ ] +class BatchGuardrailRecord(BaseModel): + """One batch input record a guardrail acted on.""" + + line: int + """The 1-based line of the uploaded file the record started on.""" + + custom_id: str | None = None + """The record's own `custom_id`, when it carried one.""" + + action: Literal["redacted", "dropped"] + """`redacted` means the record was submitted with the guardrail's rewrite applied. + + `dropped` means the guardrail blocked it and it was left out of the submitted file. + """ + + guardrail: str | None = None + """Which guardrail dropped the record, when it named itself. + + Set for dropped records only. A guardrail refusing content and a guardrail that is + unreachable under a fail-closed setting raise the same way, so this names the guardrail + to check rather than claiming a reason it cannot distinguish. + """ + + +class BatchGuardrailReport(BaseModel): + """What guardrails did to a batch input file, per record.""" + + submitted_records: int + """How many records reached the provider.""" + + modified_records: tuple[BatchGuardrailRecord, ...] + """Every record that was redacted or dropped, in file order.""" + + class OpenAIFileObject(BaseModel): id: str """The file identifier, which can be referenced in the API endpoints.""" @@ -319,6 +353,12 @@ class OpenAIFileObject(BaseModel): `error` field on `fine_tuning.job`. """ + litellm_batch_guardrail: BatchGuardrailReport | None = None + """Set by the proxy when guardrails acted on a `purpose=batch` upload. + + Absent on every other upload, so OpenAI-shaped clients see an unchanged response. + """ + _hidden_params: dict = {"response_cost": 0.0} # no cost for writing a file def __contains__(self, key) -> bool: diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 63c93e0f268..d68d1dc9625 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -27,6 +27,22 @@ class RequestComplexityRouterConfig(ComplexityRouterConfig): ) +class ComplexityRouterConfigValidationRequest(BaseModel): + """A complexity-router config to validate without saving, so a form can surface the + backend's own verdict inline instead of a raw 400 at write time.""" + + complexity_router_config: Mapping[str, object] + team_id: str | None = Field( + default=None, + description="Team the router is being created for. Required for a team admin, who may only validate their own team's routers", + ) + + +class ComplexityRouterConfigValidationResponse(BaseModel): + valid: bool + error: str | None = None + + class AutoRouterRoutingTestRequest(BaseModel): """A single prompt to classify against a complexity-router config that need not be saved yet.""" @@ -60,7 +76,7 @@ class AutoRouterRoutingTestResponse(BaseModel): routed_model: str = Field(description="The model group the router picked") routed_model_configured: bool = Field( - description="Whether routed_model is a model group this proxy actually serves", + description="Whether routed_model is a model group available to the caller, scoped to team_id when given. Never confirms models the caller could not use", ) routing_decision: StandardLoggingRoutingDecision = Field( description="The decision record this request would have written to its log row", diff --git a/ruff-tests.toml b/ruff-tests.toml new file mode 100644 index 00000000000..c1bdcc755a7 --- /dev/null +++ b/ruff-tests.toml @@ -0,0 +1,15 @@ +# Lint config for the test tree, which ruff.toml excludes from `ruff check`. +# +# Deliberately one rule. F821 is the cheapest guard against a test that cannot fail: +# a name that does not exist raises NameError, and a test whose body is wrapped in +# `except Exception: pass` swallows that NameError and reports green. Widening this +# select list means ratcheting thousands of pre-existing findings, so new rules go in +# one at a time, each with its violations already fixed. +# +# No target-version here on purpose: it resolves from requires-python (>=3.10), so +# 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that +# still has to run on 3.10. + +line-length = 120 + +lint.select = ["F821"] diff --git a/tests/documentation_tests/test_router_settings.py b/tests/documentation_tests/test_router_settings.py index 290aa283af4..a1b6f1dac1d 100644 --- a/tests/documentation_tests/test_router_settings.py +++ b/tests/documentation_tests/test_router_settings.py @@ -61,7 +61,7 @@ try: documented_keys.update(doc_key_pattern.findall(table_content)) except Exception as e: raise Exception( - f"Error reading documentation: {e}, \n repo base - {os.listdir(repo_base)}" + f"Error reading documentation: {e}, \n repo base - {os.listdir(_repo_root)}" ) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index 6f58bb2eb35..01fd35cb42d 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -842,6 +842,8 @@ def test_completion_mistral_api_modified_input(): @pytest.mark.skip(reason="this test is flaky") def test_completion_gpt4_vision(): + import openai + try: litellm.set_verbose = True response = completion( @@ -1820,6 +1822,8 @@ def test_completion_openai_litellm_key(): @pytest.mark.skip(reason="Unresponsive endpoint.[TODO] Rehost this somewhere else") def test_completion_ollama_hosted(): + import openai + try: litellm.request_timeout = 20 # give ollama 20 seconds to response litellm.set_verbose = True @@ -2057,17 +2061,12 @@ def test_completion_openrouter_reasoning_effort(): def test_completion_hf_model_no_provider(): - try: - response = completion( + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + completion( model="WizardLM/WizardLM-70B-V1.0", messages=messages, max_tokens=5, ) - # Add any assertions here to check the response - print(response) - pytest.fail(f"Error occurred: {e}") - except Exception as e: - pass # test_completion_hf_model_no_provider() @@ -2546,7 +2545,7 @@ def test_completion_replicate_vicuna(): response_str = response["choices"][0]["message"]["content"] print("RESPONSE STRING\n", response_str) if type(response_str) != str: - pytest.fail(f"Error occurred: {e}") + pytest.fail(f"Expected a string response, got {type(response_str)}: {response_str}") except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index e02d9e21171..8c1df52e28e 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -573,7 +573,7 @@ def test_content_policy_violation_error_streaming(): num_finish_reason += 1 print("finish_reason", chunk["choices"][0].get("finish_reason")) - pytest.fail(f"Expected to return 400 error In streaming{e}") + pytest.fail("Expected a content-policy error in streaming, got a clean stream") except Exception as e: pass diff --git a/tests/local_testing/test_llm_guard.py b/tests/local_testing/test_llm_guard.py index 78bbd1c0af8..86fa80ee944 100644 --- a/tests/local_testing/test_llm_guard.py +++ b/tests/local_testing/test_llm_guard.py @@ -15,6 +15,8 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest +from fastapi import HTTPException + import litellm from litellm_enterprise.enterprise_callbacks.llm_guard import _ENTERPRISE_LLMGuard from litellm import Router, mock_completion @@ -128,7 +130,7 @@ async def test_llm_guard_error_raising(): user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) local_cache = DualCache() - try: + with pytest.raises(HTTPException) as exc_info: await llm_guard.async_moderation_hook( data={ "messages": [ @@ -141,9 +143,9 @@ async def test_llm_guard_error_raising(): user_api_key_dict=user_api_key_dict, call_type="completion", ) - pytest.fail(f"Should have failed - {str(e)}") - except Exception as e: - pass + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "Violated content safety policy"} def test_llm_guard_key_specific_mode(): diff --git a/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py b/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py index 48eb7d85ec1..c1339ce6280 100644 --- a/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py +++ b/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py @@ -148,65 +148,6 @@ async def test_claude_agent_sdk_streaming( f"Test failed for {model_name} ({model_description}) after {MAX_RETRIES} attempts: {last_error}" ) - # Test query - test_query = "Say 'Hello from LiteLLM!' and nothing else." - - # Track streaming - received_chunks = [] - full_response = "" - - try: - async with ClaudeSDKClient(options=options) as client: - await client.query(test_query) - - # Collect streaming response - async for msg in client.receive_response(): - # Handle different message types - if hasattr(msg, "type"): - if msg.type == "content_block_delta": - # Streaming text delta - if hasattr(msg, "delta") and hasattr(msg.delta, "text"): - chunk_text = msg.delta.text - received_chunks.append(chunk_text) - full_response += chunk_text - elif msg.type == "content_block_start": - # Start of content block - if hasattr(msg, "content_block") and hasattr( - msg.content_block, "text" - ): - chunk_text = msg.content_block.text - received_chunks.append(chunk_text) - full_response += chunk_text - - # Fallback to content handling - if hasattr(msg, "content"): - for content_block in msg.content: - if hasattr(content_block, "text"): - chunk_text = content_block.text - received_chunks.append(chunk_text) - full_response += chunk_text - - # Assertions - print(f"\n✅ Received {len(received_chunks)} chunks") - print(f"📝 Full response: {full_response[:100]}...") - - # Verify we got a response - assert len(full_response) > 0, f"No response received from {model_name}" - - # Verify streaming (should have multiple chunks for most responses) - # Note: Very short responses might come in 1 chunk, so we just verify we got content - assert len(received_chunks) > 0, f"No chunks received from {model_name}" - - # Verify response is non-empty (don't assert on specific LLM content — it's non-deterministic) - assert ( - len(full_response.strip()) > 0 - ), f"Empty response received from {model_name}" - - print(f"✅ Test passed for {model_name}") - - except Exception as e: - pytest.fail(f"Test failed for {model_name} ({model_description}): {str(e)}") - if __name__ == "__main__": # Run tests diff --git a/tests/proxy_unit_tests/test_custom_callback_input.py b/tests/proxy_unit_tests/test_custom_callback_input.py index 71a7e94b180..a032b8706bc 100644 --- a/tests/proxy_unit_tests/test_custom_callback_input.py +++ b/tests/proxy_unit_tests/test_custom_callback_input.py @@ -2,6 +2,7 @@ ## This test asserts the type of data passed into each method of the custom callback handler import asyncio import inspect +import json import os import sys import time diff --git a/tests/test_end_users.py b/tests/test_end_users.py index ff3cc4ec94b..bc1fcbb662d 100644 --- a/tests/test_end_users.py +++ b/tests/test_end_users.py @@ -14,47 +14,6 @@ from typing import Optional """ -async def chat_completion_with_headers(session, key, model="gpt-4"): - url = "http://0.0.0.0:4000/chat/completions" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - data = { - "model": model, - "messages": [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello!"}, - ], - } - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - print(response_text) - print() - - if status != 200: - raise Exception(f"Request did not return a 200 status code: {status}") - - response_header_check( - response - ) # calling the function to check response headers - - raw_headers = response.raw_headers - raw_headers_json = {} - - for ( - item - ) in ( - response.raw_headers - ): # ((b'date', b'Fri, 19 Apr 2024 21:17:29 GMT'), (), ) - raw_headers_json[item[0].decode("utf-8")] = item[1].decode("utf-8") - - return raw_headers_json - - async def generate_key( session, i, diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py index b8ac1987453..7e7eee5373f 100644 --- a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py @@ -80,8 +80,28 @@ def _write_metadata_only_file(home): return path -def _blob(base_url=SERVER, key="sk-vault", jwt_token="", timestamp=0.0): - return json.dumps({"base_url": base_url, "key": key, "jwt_token": jwt_token, "timestamp": timestamp}) +def _blob(base_url=SERVER, key="sk-vault", jwt_token="", timestamp=0.0, refresh_token=None): + return json.dumps( + { + "base_url": base_url, + "key": key, + "jwt_token": jwt_token, + "refresh_token": refresh_token, + "timestamp": timestamp, + } + ) + + +def _write_key_only_keychain_file(home, *, refresh_token="rt-live", timestamp=2000.0): + """What the release that kept only the key in the keychain left on disk: metadata, plus the + refresh token in the clear.""" + path = _token_file(home) + path.parent.mkdir(exist_ok=True) + path.write_text( + json.dumps({"base_url": SERVER, "user_id": "u-1", "refresh_token": refresh_token, "timestamp": timestamp}) + ) + path.chmod(0o600) + return path _REAL_MKSTEMP = tempfile.mkstemp @@ -161,6 +181,56 @@ class TestLoadCliToken: assert (record.key, record.jwt_token) == ("sk-a", "jwt-a") + def test_the_refresh_token_round_trips_through_the_vault(self, isolated_home, secret_vault_factory): + """A refresh token mints a fresh key from the proxy on demand, so it is the credential just + as much as the key is, and it has to come back out of the keychain to be usable.""" + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(blob=_blob(key="sk-a", refresh_token="rt-a")) + + record = load_cli_token(vault=vault) + + assert (record.key, record.refresh_token) == ("sk-a", "rt-a") + + def test_a_plaintext_refresh_token_is_moved_off_disk(self, isolated_home, secret_vault_factory): + path = _write_legacy_file(isolated_home, refresh_token="rt-legacy") + vault = secret_vault_factory() + + record = load_cli_token(vault=vault) + + assert record.refresh_token == "rt-legacy" + assert "rt-legacy" not in path.read_text() + assert json.loads(vault.blob)["refresh_token"] == "rt-legacy" + + def test_an_upgrade_that_left_the_refresh_token_on_disk_rejoins_it_with_the_key( + self, isolated_home, secret_vault_factory + ): + """The release before this one took the key into the keychain and left the refresh token + behind, so upgrading finds one sign-in split across both stores. The read has to end with + the whole credential in the keychain, not with whichever half it happened to prefer.""" + path = _write_key_only_keychain_file(isolated_home) + vault = secret_vault_factory(blob=_blob(key="sk-live", timestamp=2000.0)) + + record = load_cli_token(vault=vault) + + assert (record.key, record.refresh_token) == ("sk-live", "rt-live") + assert "rt-live" not in path.read_text() + assert json.loads(vault.blob)["key"] == "sk-live" + assert json.loads(vault.blob)["refresh_token"] == "rt-live" + + def test_a_superseded_refresh_token_on_disk_never_outlives_the_keychain( + self, isolated_home, secret_vault_factory + ): + """Two stores, two sign-ins, and the newer one is in the keychain. Handing back its key with + the older one's refresh token would build a credential neither store ever held, and would + renew the login the user already replaced.""" + path = _write_legacy_file(isolated_home, key="sk-old", refresh_token="rt-old", timestamp=1000.0) + vault = secret_vault_factory(blob=_blob(key="sk-new", refresh_token="rt-new", timestamp=2000.0)) + + record = load_cli_token(vault=vault) + + assert (record.key, record.refresh_token) == ("sk-new", "rt-new") + assert "rt-old" not in path.read_text() + def test_legacy_plaintext_file_still_authenticates_and_is_migrated(self, isolated_home, secret_vault_factory): """A token.json written by an older `lite` keeps working, and reading it moves the secret into the keychain and scrubs it from disk.""" @@ -361,6 +431,36 @@ class TestSaveCliToken: assert json.loads(vault.blob)["key"] == "sk-new" assert load_cli_token(vault=vault).key == "sk-new" + def test_the_refresh_token_goes_to_the_keychain_and_never_to_the_file( + self, isolated_home, secret_vault_factory + ): + vault = secret_vault_factory() + + stored = save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-new", refresh_token="rt-new", timestamp=time.time()), + vault=vault, + ) + + assert stored == SecretStored() + assert "rt-new" not in _token_file(isolated_home).read_text() + assert json.loads(vault.blob)["refresh_token"] == "rt-new" + assert load_cli_token(vault=vault).refresh_token == "rt-new" + + def test_the_refresh_token_falls_back_to_the_owner_only_file_with_the_key( + self, isolated_home, secret_vault_factory + ): + """A machine with no keychain keeps the whole credential in the 0600 file, refresh token + included, because a renewal that cannot be stored logs the user out on the next command.""" + vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) + + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-new", refresh_token="rt-new", timestamp=time.time()), + vault=vault, + ) + + assert json.loads(_token_file(isolated_home).read_text())["refresh_token"] == "rt-new" + assert load_cli_token(vault=vault).refresh_token == "rt-new" + def test_falls_back_to_the_owner_only_file_when_there_is_no_keychain(self, isolated_home, secret_vault_factory): stored = save_cli_token( CliTokenRecord(base_url=SERVER, key="sk-new", timestamp=time.time()), @@ -628,6 +728,26 @@ class TestScrubFailure: assert json.loads(path.read_text()).get("key") is None + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") + def test_a_rejoin_the_file_refuses_never_takes_the_key_with_it( + self, isolated_home, secret_vault_factory, monkeypatch + ): + """Rolling the rejoined entry back would erase a key that was safely in the keychain before + this read began, and the file it would fall back to is the one that has just refused to be + rewritten. The duplicate refresh token stays until a later read can finish the move.""" + path = _write_key_only_keychain_file(isolated_home) + vault = secret_vault_factory(blob=_blob(key="sk-live", timestamp=2000.0)) + monkeypatch.setattr("litellm.litellm_core_utils.private_json.os.replace", _refuse_replace) + path.chmod(0o400) + + record = load_cli_token(vault=vault) + + assert (record.key, record.refresh_token) == ("sk-live", "rt-live") + assert json.loads(vault.blob)["key"] == "sk-live" + assert json.loads(vault.blob)["refresh_token"] == "rt-live" + assert vault.erases == 0 + + class TestClearCliToken: def test_removes_the_credential_from_both_stores(self, isolated_home, secret_vault_factory): vault = secret_vault_factory() @@ -717,12 +837,14 @@ class TestClearCliToken: ): """Keeping a record of the unreachable keychain must never mean keeping the cleartext copy the user just asked to be rid of.""" - _write_legacy_file(isolated_home) + _write_legacy_file(isolated_home, refresh_token="rt-legacy") vault = secret_vault_factory(available=False, failure=KeyringUnreachable()) clear_cli_token(vault=vault) - assert "sk-legacy" not in _token_file(isolated_home).read_text() + left_on_disk = _token_file(isolated_home).read_text() + assert "sk-legacy" not in left_on_disk + assert "rt-legacy" not in left_on_disk def test_a_repeat_logout_never_answers_its_own_warning_with_an_all_clear( self, isolated_home, secret_vault_factory diff --git a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py b/tests/test_litellm/llms/oci/test_oci_coverage_boost.py index 0b7afa3775d..7c91ece70b5 100644 --- a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py +++ b/tests/test_litellm/llms/oci/test_oci_coverage_boost.py @@ -11,11 +11,16 @@ All tests are self-contained and require no real OCI credentials or network acce """ import json +from typing import TYPE_CHECKING + import pytest from unittest.mock import patch, MagicMock, AsyncMock import httpx +if TYPE_CHECKING: + from litellm.llms.oci.chat.transformation import OCIStreamWrapper + from litellm import ModelResponse from litellm.llms.oci.chat.cohere import ( _extract_text_content, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index 64afa52ab55..65e2faee1b2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -2,6 +2,11 @@ to exactly one category, wire values never carry upstream prose, and single-upstream HTTP statuses stay truthful to who failed.""" +import sys + +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + import httpx import pytest from mcp import McpError diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py index a12c02339e6..b8bf4da1dc4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py @@ -2,6 +2,11 @@ links win (the ``raise ... from`` cause subtree, then ExceptionGroup members in raise order, then the incidental ``__context__`` chain last), and adversarial shapes terminate.""" +import sys + +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index bdaf1458fb0..34852850de6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4,6 +4,7 @@ import hashlib import json import time from base64 import urlsafe_b64encode +from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -11,6 +12,11 @@ from fastapi import HTTPException from litellm.types.mcp import MCPAuth +if TYPE_CHECKING: + import httpx + + from litellm.types.mcp_server.mcp_server_manager import MCPServer + # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 7bdd3b36763..feab179570b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -1411,6 +1411,8 @@ async def _run_passthrough_connect( ): """Drive handle_streamable_http_mcp through the preemptive-401 gate and report whether it challenged (raised) or forwarded to the session manager. Returns (challenged, www_authenticate).""" + from fastapi import HTTPException + from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, session_manager_stateless, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 7ba9f463197..d7fb121ef9b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,9 +1,13 @@ import asyncio import json +import sys from datetime import datetime from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + import httpx import pytest from fastapi import HTTPException diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 83e4dcf5677..da7f43c7118 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -12,6 +12,9 @@ from unittest.mock import AsyncMock, Mock, patch import pytest +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + sys.path.insert(0, os.path.abspath("../..")) from mcp.types import Tool as MCPTool diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 1d1bd9ebf8a..7a3288dec37 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -13,7 +13,7 @@ from datetime import datetime, timedelta, timezone import httpx import pytest -from fastapi import status +from fastapi import Request, status import litellm from litellm.proxy._types import ( @@ -2760,11 +2760,9 @@ async def test_common_checks_metadata_route_keeps_key_tags_out_of_provider_metad assert "metadata" not in request_body -def _pass_through_request() -> "Request": +def _pass_through_request() -> Request: """A Request whose FastAPI-resolved endpoint carries the pass-through marker, i.e. the request was dispatched to a user-defined pass-through handler.""" - from fastapi import Request - from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) @@ -2776,10 +2774,9 @@ def _pass_through_request() -> "Request": return Request(scope={"type": "http", "headers": [], "endpoint": pass_through_endpoint}) -def _builtin_request() -> "Request": +def _builtin_request() -> Request: """A Request dispatched to a built-in (non-pass-through) handler, e.g. what a custom path colliding with a core route actually resolves to.""" - from fastapi import Request def chat_completions(): ... diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 21511c74154..aa9e5349c87 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3379,3 +3379,52 @@ def test_organization_daily_activity_not_granted_by_org_admin_request_data_branc route="/organization/daily/activity", allowed_routes=LiteLLMRoutes.org_admin_only_routes.value, ) + + +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + LitellmUserRoles.TEAM.value, + ], +) +@pytest.mark.parametrize( + "dry_run_route", + ["/auto_router/test_routing", "/auto_router/validate_complexity_router_config"], +) +def test_auto_router_dry_runs_share_model_new_audience(user_role, dry_run_route): + """The dry runs serve whoever can draft a save on /model/new, no one else: a role + must get the same allow-or-403 from both layers' route check, or the form's + pre-save call 403s for an operator whose save would have been accepted.""" + + def outcome(route: str) -> str: + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=user_role, + ) + valid_token = UserAPIKeyAuth( + user_id="test_user", + user_role=user_role, + ) + request = MagicMock(spec=Request) + request.query_params = {} + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + return "allowed" + except HTTPException: + return "rejected" + + assert outcome(dry_run_route) == outcome("/model/new") + # Anchor so parity cannot be satisfied by both routes 403ing for everyone + if user_role == LitellmUserRoles.INTERNAL_USER.value: + assert outcome(dry_run_route) == "allowed" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py index 36a2e205ea7..a3d86034f70 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -1067,3 +1067,29 @@ async def test_anthropic_non_streaming_response_reports_usage(): payload = _posted_payload(g) assert payload["usage"] == {"input_tokens": 10, "output_tokens": 5} assert payload["response"]["finish_reason"] == "end_turn" + + +def test_fail_closed_backend_failure_is_not_reported_as_a_content_verdict(): + """A drop-one-record consumer must be able to tell a verdict from an outage; _fail is not a verdict.""" + from litellm.exceptions import GuardrailRaisedException + + guardrail = _make_guardrail() + + with pytest.raises(GuardrailRaisedException) as unreachable: + guardrail._fail( + inputs={}, + request_data={"model": "m"}, + input_type="request", + error="connection refused", + is_unreachable=True, + ) + assert unreachable.value.blocked_content is False + + with pytest.raises(GuardrailRaisedException) as verdict: + guardrail._block( + request_data={"model": "m"}, + input_type="request", + message="blocked", + blocked_content=True, + ) + assert verdict.value.blocked_content is True diff --git a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py index c977223bfab..48f6b3ba2b9 100644 --- a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py @@ -9,7 +9,7 @@ and following LiteLLM testing patterns and best practices. import importlib import os import sys -from typing import Dict +from typing import Any, Dict from unittest.mock import Mock, patch # Add parent directory to path for imports diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 3fd023552f5..c901696e108 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -5,6 +5,7 @@ Unit tests for auto router management endpoints import os import sys from pathlib import Path +from typing import Final import pytest from fastapi import HTTPException @@ -1188,6 +1189,149 @@ async def test_stop_shadow_eval_stops_every_unstopped_leg_and_rejects_non_runnin assert forbidden.value.status_code == 403 +@pytest.mark.asyncio +async def test_validate_config_returns_the_write_gates_verdict_without_saving(): + """The dry-run endpoint must agree with the write gate exactly, so a form showing its + verdict inline can never pass a config the save would then reject.""" + from litellm.proxy.management_endpoints.auto_router_endpoints import ( + validate_complexity_router_config, + ) + from litellm.types.management_endpoints.auto_router_endpoints import ( + ComplexityRouterConfigValidationRequest, + ) + + valid = await validate_complexity_router_config( + ComplexityRouterConfigValidationRequest( + complexity_router_config={ + "tiers": {"CASUAL": "m1", "AUDIT": "m2"}, + "tier_definitions": [ + {"name": "CASUAL", "description": "casual chat"}, + {"name": "AUDIT", "description": "security audits"}, + ], + "fallback_tier": "AUDIT", + "classifier_type": "llm", + "classifier_llm_config": {"model": "clf"}, + } + ), + ADMIN, + ) + assert valid.valid is True + assert valid.error is None + + rejected = await validate_complexity_router_config( + ComplexityRouterConfigValidationRequest( + complexity_router_config={ + "tiers": {"CASUAL": "m1", "AUDIT": "m2"}, + "tier_definitions": [ + {"name": "CASUAL", "description": "casual chat"}, + {"name": "AUDIT", "description": "security\naudits"}, + ], + "fallback_tier": "AUDIT", + "classifier_type": "llm", + "classifier_llm_config": {"model": "clf"}, + } + ), + ADMIN, + ) + assert rejected.valid is False + assert rejected.error is not None and "newline" in rejected.error + + +@pytest.mark.asyncio +async def test_routing_test_never_confirms_models_the_caller_cannot_use(monkeypatch: pytest.MonkeyPatch): + """routed_model_configured must not be an existence oracle for the whole proxy: a team + admin probing a guessed global model name reads False unless the named team could + actually use that model, and True once the team grants it.""" + from litellm.proxy import proxy_server + + def _team_prisma(team_id: str, models: list[str]) -> MagicMock: + row_data = { + "team_id": team_id, + "members_with_roles": [{"role": "admin", "user_id": "team-admin"}], + "models": models, + } + team_row = MagicMock() + team_row.model_dump.return_value = row_data + team_row.dict.return_value = row_data + prisma = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + return prisma + + def _request(team_id: str) -> AutoRouterRoutingTestRequest: + return AutoRouterRoutingTestRequest.model_validate( + { + "prompt": "what is 2+2", + "complexity_router_config": {"tiers": TIERS, "classifier_type": "heuristic"}, + "team_id": team_id, + } + ) + + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + + team_admin: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-team", user_id="team-admin" + ) + + monkeypatch.setattr(proxy_server, "prisma_client", _team_prisma("team-probe", models=["mid-model"])) + probing = await preview_auto_router_routing(data=_request("team-probe"), user_api_key_dict=team_admin) + assert probing.routed_model == "cheap-model" + assert probing.routed_model_configured is False + + monkeypatch.setattr(proxy_server, "prisma_client", _team_prisma("team-grant", models=["cheap-model"])) + granted = await preview_auto_router_routing(data=_request("team-grant"), user_api_key_dict=team_admin) + assert granted.routed_model == "cheap-model" + assert granted.routed_model_configured is True + + +@pytest.mark.asyncio +async def test_validate_config_gates_like_the_write_it_rehearses(monkeypatch: pytest.MonkeyPatch): + """A caller who could not save the router must not get the dry run either: matching + /model/new, a team admin passes only when naming their own team, and a caller who is + neither proxy admin nor team admin is rejected before validation runs.""" + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.auto_router_endpoints import ( + validate_complexity_router_config, + ) + from litellm.types.management_endpoints.auto_router_endpoints import ( + ComplexityRouterConfigValidationRequest, + ) + + config: Final = {"tiers": {"SIMPLE": "m1"}, "classifier_type": "heuristic"} + + with pytest.raises(HTTPException) as forbidden: + await validate_complexity_router_config( + ComplexityRouterConfigValidationRequest(complexity_router_config=config), VIEWER + ) + assert forbidden.value.status_code == 403 + + team_row: Final = MagicMock() + team_row.model_dump.return_value = { + "team_id": "team-1", + "members_with_roles": [{"role": "admin", "user_id": "team-admin"}], + } + prisma: Final = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "premium_user", True) + + team_admin: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-team", user_id="team-admin" + ) + verdict = await validate_complexity_router_config( + ComplexityRouterConfigValidationRequest(complexity_router_config=config, team_id="team-1"), + team_admin, + ) + assert verdict.valid is True + + with pytest.raises(HTTPException) as not_their_team: + await validate_complexity_router_config( + ComplexityRouterConfigValidationRequest(complexity_router_config=config, team_id="team-1"), + UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-other", user_id="someone-else"), + ) + assert not_their_team.value.status_code == 403 + + def test_every_shadow_eval_sql_constant_speaks_naive_utc(): """The tables store naive UTC wall time (prisma's convention), so SQL-side time must be NOW() AT TIME ZONE 'utc' and python-side params must cast ::timestamp; a bare NOW() or a diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py new file mode 100644 index 00000000000..a05b8ae530c --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -0,0 +1,1041 @@ +import io +import json + +import pytest +from fastapi import HTTPException + +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.batch_guardrails import ( + BatchScanResult, + RecordDropped, + RecordRedacted, + rewrite_batch_input_file, + UnparseableRecord, + UnscannableRecord, + raise_public, + scan_batch_input_file, +) + + +def _record(custom_id, content="hello", url="/v1/chat/completions"): + return { + "custom_id": custom_id, + "method": "POST", + "url": url, + "body": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": content}], + }, + } + + +def _jsonl(*records): + return io.BytesIO("\n".join(json.dumps(r) for r in records).encode()) + + +class FakeProxyLogging: + """Stands in for ProxyLogging so the scan can be driven without a live proxy.""" + + def __init__(self, on_record=None): + self.on_record = on_record or (lambda data: None) + self.seen = [] + + async def pre_call_hook(self, user_api_key_dict, data, call_type, guardrails_only=False): + self.seen.append((call_type, json.dumps(data.get("messages"), sort_keys=True))) + self.on_record(data) + return data + + def has_pre_call_guardrails(self, request_metadata): + return True + + +def _redact_containing(needle): + def _hook(data): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and needle in message["content"]: + message["content"] = message["content"].replace(needle, "***") + + return _hook + + +def _raise_on(needle, exc): + def _hook(data): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and needle in message["content"]: + raise exc + + return _hook + + +async def _scan_full(source, logging_obj, metadata=None): + return await scan_batch_input_file( + file_source=source, + request_metadata=metadata if metadata is not None else {}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=logging_obj, + ) + + +async def _scan(source, logging_obj, metadata=None): + """Collapses "the scan found nothing to do" to None so the reject-mode cases read plainly.""" + result = await _scan_full(source, logging_obj, metadata) + if isinstance(result, BatchScanResult): + return None if not result.changes else result + return result + + +@pytest.mark.asyncio +async def test_clean_file_passes_and_rewinds_the_handle(): + source = _jsonl(_record("a"), _record("b"), _record("c")) + logging_obj = FakeProxyLogging() + + assert await _scan(source, logging_obj) is None + assert len(logging_obj.seen) == 3 + assert source.tell() == 0, "handle must be rewound so the upload still sees the whole file" + + +@pytest.mark.asyncio +async def test_every_record_is_scanned_not_just_the_first(): + records = [_record(f"r{i}") for i in range(70)] + logging_obj = FakeProxyLogging() + + assert await _scan(_jsonl(*records), logging_obj) is None + assert len(logging_obj.seen) == 70, "records past the first scan window must still be scanned" + + +@pytest.mark.asyncio +async def test_redaction_is_reported_with_line_and_custom_id(): + source = _jsonl(_record("keep-1"), _record("dirty", content="my secret is here"), _record("keep-2")) + + failure = await _scan(source, FakeProxyLogging(_redact_containing("secret"))) + + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(2, "dirty")] + + +@pytest.mark.asyncio +async def test_body_carrying_its_own_metadata_is_not_reported_as_redacted(): + record = _record("has-meta") + record["body"]["metadata"] = {"team": "finance"} + + failure = await _scan(_jsonl(record), FakeProxyLogging(), metadata={"guardrails": ["x"]}) + + assert failure is None, "the metadata the proxy injects must not be diffed as record content" + + +@pytest.mark.asyncio +async def test_guardrail_writing_bookkeeping_into_metadata_is_not_a_redaction(): + def _touch_metadata(data): + data["litellm_metadata"]["applied_guardrails"] = ["some-guard"] + + assert await _scan(_jsonl(_record("a")), FakeProxyLogging(_touch_metadata)) is None + + +@pytest.mark.asyncio +async def test_records_own_metadata_is_left_out_of_the_scan_and_the_diff(): + """Guardrail dispatch writes bookkeeping into `metadata`; diffing it would reject every such record.""" + record = _record("has-meta") + record["body"]["metadata"] = {"team": "finance"} + seen = [] + + def _write_bookkeeping(data): + seen.append(dict(data.get("metadata") or {})) + data.setdefault("metadata", {})["applied_guardrails"] = ["g"] + + assert await _scan(_jsonl(record), FakeProxyLogging(_write_bookkeeping), metadata={"tags": ["t"]}) is None + assert seen == [{"tags": ["t"]}], "dispatch sees the proxy's metadata, never the record's own" + assert record["body"]["metadata"] == {"team": "finance"} + + +@pytest.mark.asyncio +async def test_the_scan_metadata_reaches_guardrails_that_only_read_the_metadata_bag(): + """noma and aim read `metadata["headers"]`; a record scanned as chat must reach them too.""" + seen = [] + + await _scan( + _jsonl(_record("a")), + FakeProxyLogging(lambda d: seen.append((d.get("metadata") or {}).get("headers"))), + metadata={"guardrails": ["g"], "headers": {"x-noma-application-id": "app-1"}}, + ) + + assert seen == [{"x-noma-application-id": "app-1"}] + + +@pytest.mark.asyncio +async def test_request_metadata_is_narrowed_to_what_guardrails_read(): + """An OTel-enabled proxy puts a lock-bearing span here; a per-record copy of it is a crash.""" + import threading + + seen = [] + metadata = { + "guardrails": ["g"], + "tags": ["t"], + "headers": {"x-noma-application-id": "app-1"}, + "litellm_parent_otel_span": threading.RLock(), + "user_api_key": "sk-secret", + } + + failure = await _scan( + _jsonl(_record("a")), + FakeProxyLogging(lambda d: seen.append(dict(d["litellm_metadata"]))), + metadata=metadata, + ) + + assert failure is None + assert seen == [{"guardrails": ["g"], "tags": ["t"], "headers": {"x-noma-application-id": "app-1"}}] + + +@pytest.mark.asyncio +async def test_one_record_cannot_leak_a_metadata_write_into_the_next_one(): + """`headers` and `tags` are nested and shared; an in-place write must not cross records.""" + seen = [] + + def _tamper(data): + bag = data["litellm_metadata"] + seen.append((dict(bag["headers"]), list(bag["tags"]))) + bag["headers"]["x-injected"] = "from-record-1" + bag["tags"].append("from-record-1") + + metadata = {"guardrails": ["g"], "headers": {"x-real": "yes"}, "tags": ["real"]} + await _scan(_jsonl(_record("a"), _record("b")), FakeProxyLogging(_tamper), metadata=metadata) + + assert seen == [({"x-real": "yes"}, ["real"]), ({"x-real": "yes"}, ["real"])] + assert metadata == {"guardrails": ["g"], "headers": {"x-real": "yes"}, "tags": ["real"]} + + +@pytest.mark.asyncio +async def test_records_are_scanned_under_the_headers_the_upload_carried(): + """Guardrails such as noma pick their application from a header, so dropping it changes the policy.""" + seen = [] + + await _scan( + _jsonl(_record("a")), + FakeProxyLogging(lambda d: seen.append(d["litellm_metadata"].get("headers"))), + metadata={"guardrails": ["g"], "headers": {"x-noma-application-id": "app-1"}}, + ) + + assert seen == [{"x-noma-application-id": "app-1"}] + + +@pytest.mark.asyncio +async def test_guardrail_that_adds_a_key_is_detected(): + def _add_key(data): + data["mock_response"] = "intercepted" + + failure = await _scan(_jsonl(_record("a")), FakeProxyLogging(_add_key)) + + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")] + + +@pytest.mark.asyncio +async def test_guardrail_that_adds_a_null_valued_key_is_detected(): + """A null value must not read the same as a missing key, or dropping one hides a change.""" + + def _add_null_key(data): + data["response_format"] = None + + failure = await _scan(_jsonl(_record("a")), FakeProxyLogging(_add_null_key)) + + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")] + + +@pytest.mark.asyncio +async def test_guardrail_that_drops_a_null_valued_key_is_detected(): + def _drop_null_key(data): + data.pop("response_format") + + record = _record("a") + record["body"]["response_format"] = None + + failure = await _scan(_jsonl(record), FakeProxyLogging(_drop_null_key)) + + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")] + + +@pytest.mark.asyncio +async def test_guardrail_that_only_reorders_a_nested_dict_is_not_a_redaction(): + def _reorder(data): + message = data["messages"][0] + data["messages"][0] = {key: message[key] for key in reversed(list(message))} + + assert await _scan(_jsonl(_record("a")), FakeProxyLogging(_reorder)) is None + + +@pytest.mark.asyncio +async def test_record_without_a_url_falls_back_to_its_body_shape(): + logging_obj = FakeProxyLogging() + record = _record("no-url") + del record["url"] + + assert await _scan(_jsonl(record), logging_obj) is None + assert logging_obj.seen[0][0] == "acompletion" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body, expected_call_type", + [ + ({"messages": [{"role": "user", "content": "x"}]}, "acompletion"), + ({"prompt": "x"}, "atext_completion"), + ({"input": "x"}, "aembedding"), + ], +) +async def test_empty_url_falls_back_to_its_body_shape(body, expected_call_type): + logging_obj = FakeProxyLogging() + record = {"custom_id": "c", "url": "", "body": {"model": "m", **body}} + + assert await _scan(_jsonl(record), logging_obj) is None + assert logging_obj.seen[0][0] == expected_call_type + + +@pytest.mark.asyncio +async def test_handle_is_rewound_even_when_a_record_is_refused(): + source = _jsonl(_record("a", content="secret")) + + await _scan(source, FakeProxyLogging(_redact_containing("secret"))) + + assert source.tell() == 0 + + +@pytest.mark.asyncio +async def test_record_without_a_body_object_is_rejected(): + source = io.BytesIO(b'{"custom_id": "no-body", "url": "/v1/chat/completions"}\n') + + assert await _scan(source, FakeProxyLogging()) == UnparseableRecord(line_number=1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url, expected_call_type", + [ + ("/v1/chat/completions", "acompletion"), + ("/v1/completions", "atext_completion"), + ("/v1/embeddings", "aembedding"), + ("/v1/responses", "aresponses"), + ("/v1/messages", "anthropic_messages"), + ], +) +async def test_supported_urls_scan_under_the_matching_call_type(url, expected_call_type): + logging_obj = FakeProxyLogging() + + assert await _scan(_jsonl(_record("a", url=url)), logging_obj) is None + assert logging_obj.seen[0][0] == expected_call_type + + +@pytest.mark.asyncio +async def test_unrecognized_url_falls_back_to_the_body_shape(): + """A record we can still read is a record we can still scan, so the url alone must not reject it.""" + logging_obj = FakeProxyLogging() + + assert await _scan(_jsonl(_record("img", url="/v1/images/generations")), logging_obj) is None + assert logging_obj.seen[0][0] == "acompletion" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", + ["/chat/completions", "/v1/chat/completions/", "https://api.openai.com/v1/chat/completions"], +) +async def test_url_variants_callers_actually_write_are_accepted(url): + logging_obj = FakeProxyLogging() + + assert await _scan(_jsonl(_record("v", url=url)), logging_obj) is None + assert logging_obj.seen[0][0] == "acompletion" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url, expected_call_type", + [ + ("https://api.openai.com/v1/responses", "aresponses"), + ("https://api.openai.com/v1/embeddings", "aembedding"), + ("https://api.openai.com/v1/messages", "anthropic_messages"), + ("https://api.openai.com/v1/responses?api-version=1", "aresponses"), + ], +) +async def test_an_absolute_url_resolves_by_path_not_by_body_shape(url, expected_call_type): + """A Responses body carries `input`, which reads as an embedding if the host is not stripped first.""" + logging_obj = FakeProxyLogging() + record = {"custom_id": "abs", "method": "POST", "url": url, "body": {"model": "m", "input": "x"}} + + assert await _scan(_jsonl(record), logging_obj) is None + assert logging_obj.seen[0][0] == expected_call_type + + +@pytest.mark.asyncio +async def test_query_string_on_a_known_url_does_not_change_the_call_type(): + """The body carries `messages`, so only stripping the query string can yield aembedding.""" + logging_obj = FakeProxyLogging() + record = { + "custom_id": "q", + "url": "/v1/embeddings?api-version=1", + "body": {"model": "m", "input": "x", "messages": [{"role": "user", "content": "y"}]}, + } + + assert await _scan(_jsonl(record), logging_obj) is None + assert logging_obj.seen[0][0] == "aembedding" + + +@pytest.mark.asyncio +async def test_record_whose_body_cannot_be_read_is_rejected(): + source = _jsonl({"custom_id": "opaque", "url": "/v1/rerank", "body": {"model": "m", "documents": ["a"]}}) + + failure = await _scan(source, FakeProxyLogging()) + + assert failure == UnscannableRecord(line_number=1, custom_id="opaque", url="/v1/rerank") + + +@pytest.mark.asyncio +async def test_url_less_record_whose_body_shape_is_unknown_is_rejected(): + record = {"custom_id": "opaque", "body": {"model": "m", "something_else": 1}} + + assert await _scan(_jsonl(record), FakeProxyLogging()) == UnscannableRecord( + line_number=1, custom_id="opaque", url=None + ) + + +@pytest.mark.asyncio +async def test_blocking_guardrail_exception_propagates_unwrapped(): + blocked = HTTPException(status_code=503, detail={"error": "guardrail service unavailable"}) + source = _jsonl(_record("a"), _record("b", content="tripwire")) + + with pytest.raises(HTTPException) as raised: + await _scan(source, FakeProxyLogging(_raise_on("tripwire", blocked))) + + assert raised.value is blocked, "the guardrail's own exception must survive so its status code does" + assert raised.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_records_are_not_mutated_by_the_scan(): + record = _record("a", content="my secret is here") + payload = json.dumps(record) + source = io.BytesIO(payload.encode()) + + await _scan(source, FakeProxyLogging(_redact_containing("secret"))) + + assert source.getvalue().decode() == payload, "the scan must never rewrite the uploaded bytes" + + +@pytest.mark.parametrize( + "failure, fragment", + [ + (UnparseableRecord(line_number=7), "line 7"), + (UnscannableRecord(line_number=3, custom_id="x", url="/v1/audio/speech"), "custom_id x"), + ], +) +def test_every_failure_maps_to_a_400_naming_the_record(failure, fragment): + with pytest.raises(HTTPException) as raised: + raise_public(failure) + + assert raised.value.status_code == 400 + assert fragment in raised.value.detail["error"] + + +@pytest.mark.asyncio +async def test_scan_does_not_mutate_the_parsed_record(): + """The guardrail must redact a copy. Mutating the record would corrupt what PR 2 writes out.""" + from litellm.proxy.openai_files_endpoints.batch_guardrails import _ParsedRecord, _scan_record + + payload = _record("a", content="my secret is here") + record = _ParsedRecord(line_number=1, payload=payload) + + failure = await _scan_record( + record, + {}, + UserAPIKeyAuth(api_key="sk-test"), + FakeProxyLogging(_redact_containing("secret")), + ) + + assert (failure.line_number, failure.custom_id) == (1, "a") + assert record.payload["body"]["messages"][0]["content"] == "my secret is here", ( + "the guardrail redacted the record itself instead of a copy" + ) + + +@pytest.mark.asyncio +async def test_scan_is_bounded_so_a_huge_file_cannot_fan_out_without_limit(): + import asyncio + + from litellm.proxy.openai_files_endpoints.batch_guardrails import _SCAN_WINDOW + + in_flight = {"now": 0, "peak": 0} + + class CountingLogging(FakeProxyLogging): + async def pre_call_hook(self, user_api_key_dict, data, call_type, guardrails_only=False): + in_flight["now"] += 1 + in_flight["peak"] = max(in_flight["peak"], in_flight["now"]) + await asyncio.sleep(0) + in_flight["now"] -= 1 + return data + + records = [_record(f"r{i}") for i in range(_SCAN_WINDOW * 3)] + + assert await _scan(_jsonl(*records), CountingLogging()) is None + assert in_flight["peak"] <= _SCAN_WINDOW, ( + f"peak {in_flight['peak']} exceeded the scan window; a gigabyte file would fan out unbounded" + ) + + +@pytest.mark.asyncio +async def test_scan_runs_guardrails_only(): + """Rate limiters, budget hooks and the hanging-request alert must not fire once per record.""" + flags = [] + + class FlagCapturingLogging(FakeProxyLogging): + async def pre_call_hook(self, user_api_key_dict, data, call_type, guardrails_only=False): + flags.append(guardrails_only) + return data + + await _scan(_jsonl(_record("a"), _record("b")), FlagCapturingLogging()) + + assert flags == [True, True] + + +@pytest.mark.asyncio +async def test_guardrail_that_returns_a_replacement_dict_is_detected(): + """async_pre_call_hook may return a NEW dict instead of mutating; that result is the real input.""" + + class ReplacingLogging(FakeProxyLogging): + async def pre_call_hook(self, user_api_key_dict, data, call_type, guardrails_only=False): + replacement = json.loads(json.dumps(data)) + replacement["messages"][0]["content"] = "***" + return replacement + + failure = await _scan(_jsonl(_record("a", content="my secret is here")), ReplacingLogging()) + + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")] + + +def _blocking(needle, status_code=400, guardrail_name="block-guard"): + def _hook(data): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and needle in message["content"]: + raise HTTPException( + status_code=status_code, + detail={"error": "Violated guardrail policy", "guardrail_name": guardrail_name}, + ) + + return _hook + + +@pytest.mark.asyncio +async def test_redact_mode_keeps_a_masked_record_instead_of_rejecting(): + source = _jsonl(_record("a"), _record("b", content="my secret is here"), _record("c")) + + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + + assert [(c.line_number, c.custom_id) for c in result.changes] == [(2, "b")] + rewritten = json.loads(rewrite_batch_input_file(source, result).read().decode().splitlines()[1]) + assert rewritten["body"]["messages"][0]["content"] == "my *** is here" + assert "litellm_metadata" not in rewritten["body"], "proxy metadata must not reach the uploaded file" + assert result.submitted_records == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [400, 403, 422], ids=["content_policy", "akto", "llm_as_a_judge"]) +async def test_every_status_litellm_calls_a_block_drops_the_record(status_code): + """Follows CustomGuardrail._is_guardrail_intervention, so drop matches what litellm logs as a block.""" + source = _jsonl(_record("a"), _record("b", content="tripwire")) + + result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire", status_code))) + + assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail="block-guard"),) + + +@pytest.mark.asyncio +async def test_redact_mode_drops_a_blocked_record_and_submits_the_rest(): + source = _jsonl(_record("a"), _record("b", content="tripwire"), _record("c")) + + result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire"))) + + assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail="block-guard"),) + assert result.submitted_records == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [500, 502, 408, 429, 401]) +async def test_redact_mode_does_not_drop_a_record_on_an_infrastructure_failure(status_code): + """A guardrail service that is down must abort the upload, never silently cost the caller records.""" + source = _jsonl(_record("a"), _record("b", content="tripwire")) + + with pytest.raises(HTTPException) as raised: + await _scan_full(source, FakeProxyLogging(_blocking("tripwire", status_code))) + + assert raised.value.status_code == status_code + + +@pytest.mark.asyncio +async def test_every_record_blocked_leaves_nothing_to_submit(): + source = _jsonl(_record("a", content="tripwire"), _record("b", content="tripwire")) + + result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire"))) + + assert result.submitted_records == 0 + assert [change.line_number for change in result.changes] == [1, 2] + + +@pytest.mark.asyncio +async def test_rewrite_drops_blocked_records_and_masks_redacted_ones(): + records = [_record("a"), _record("b", content="my secret is here"), _record("c", content="tripwire"), _record("d")] + source = _jsonl(*records) + + def _hook(data): + _redact_containing("secret")(data) + _blocking("tripwire")(data) + + result = await _scan_full(source, FakeProxyLogging(_hook)) + rewritten = rewrite_batch_input_file(source, result) + + lines = [json.loads(line) for line in (rewritten.seek(0), rewritten.read().decode())[1].splitlines()] + assert [line["custom_id"] for line in lines] == ["a", "b", "d"] + assert lines[1]["body"]["messages"][0]["content"] == "my *** is here" + + +@pytest.mark.asyncio +async def test_rewrite_copies_untouched_records_byte_for_byte(): + """Enabling the feature must not reformat records no guardrail objected to.""" + untouched = '{"custom_id":"keep","url":"/v1/chat/completions","body":{"messages":[{"role":"user","content":"hi"}],"model":"m"}}' + dirty = json.dumps(_record("dirty", content="my secret is here")) + source = io.BytesIO((untouched + "\n" + dirty).encode()) + + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + rewritten = rewrite_batch_input_file(source, result) + + assert (rewritten.seek(0), rewritten.read().decode())[1].splitlines()[0] == untouched + + +@pytest.mark.asyncio +async def test_report_names_every_changed_record_in_file_order(): + records = [_record("a"), _record("b", content="tripwire"), _record("c", content="my secret is here")] + + def _hook(data): + _redact_containing("secret")(data) + _blocking("tripwire")(data) + + result = await _scan_full(_jsonl(*records), FakeProxyLogging(_hook)) + report = result.report() + + assert report.submitted_records == 2 + assert [(r.line, r.custom_id, r.action, r.guardrail) for r in report.modified_records] == [ + (2, "b", "dropped", "block-guard"), + (3, "c", "redacted", None), + ] + + +@pytest.mark.asyncio +async def test_clean_file_needs_no_rewrite(): + """A file nothing objected to keeps streaming off disk rather than being buffered in memory.""" + result = await _scan_full(_jsonl(_record("a"), _record("b")), FakeProxyLogging()) + + assert result.changes == () + assert result.submitted_records == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exc", + [ + GuardrailRaisedException(guardrail_name="g", message="blocked", blocked_content=True), + BlockedPiiEntityError(entity_type="US_SSN", guardrail_name="presidio"), + ], + ids=["guardrail_raised", "blocked_pii_entity"], +) +async def test_litellm_native_block_exceptions_drop_the_record(exc): + """Presidio and friends raise these rather than an HTTPException; they are still policy blocks.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise exc + + source = _jsonl(_record("a"), _record("b", content="tripwire")) + + result = await _scan_full(source, FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail=exc.guardrail_name),) + assert result.submitted_records == 1 + + +@pytest.mark.asyncio +async def test_raising_a_native_block_exception_drops_whatever_status_it_carries(): + """Raising this type IS the block signal in litellm, so the drop set matches what it logs as a block.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise GuardrailRaisedException( + guardrail_name="g", message="refused", status_code=503, blocked_content=True + ) + + result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail="g"),) + + +@pytest.mark.asyncio +async def test_an_unreachable_guardrail_aborts_instead_of_quietly_dropping_the_record(): + """Several integrations raise this same exception when their backend is down and they fail closed.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise GuardrailRaisedException( + guardrail_name="g", message="Singulr API unreachable (block_on_error=True): timed out" + ) + + with pytest.raises(GuardrailRaisedException): + await _scan_full(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook)) + + +@pytest.mark.asyncio +async def test_a_guardrail_subclass_that_blocks_content_drops_only_that_record(): + """A subclass has to opt in too, or a real block takes the whole upload down with it.""" + from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix import OvalixGuardrailBlockedException + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise OvalixGuardrailBlockedException(guardrail_name="ovalix", message="blocked") + + result = await _scan_full(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail="ovalix"),) + assert result.submitted_records == 1 + + +@pytest.mark.asyncio +async def test_a_record_a_guardrail_rerouted_aborts_rather_than_shipping_to_the_original_provider(): + """pre_call_hook honours a reroute by rewriting `model`; a batch file cannot follow it.""" + from litellm.proxy.openai_files_endpoints.batch_guardrails import UnroutableRecord + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + data["model"] = "on-prem-model" + data["metadata"] = { + "sensitive_data_routing_applied": True, + "sensitive_data_routing_guardrail": "router-guard", + } + + failure = await _scan(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert failure == UnroutableRecord(line_number=2, custom_id="b", guardrail="router-guard") + with pytest.raises(HTTPException) as caught: + raise_public(failure) + assert "routed to a different model" in str(caught.value.detail) + + +@pytest.mark.asyncio +async def test_the_scan_spool_is_closed_when_nothing_will_read_it(): + """The spool is opened for every scan, so a clean file must not leave a temp handle behind.""" + result = await _scan_full(_jsonl(_record("a")), FakeProxyLogging()) + + assert result.changes == () + assert result.redactions.closed + + +@pytest.mark.asyncio +async def test_the_scan_spool_is_closed_when_the_upload_is_refused(): + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise RuntimeError("infrastructure is down") + + source = _jsonl(_record("a"), _record("b", content="tripwire")) + spools = [] + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + bg.tempfile.SpooledTemporaryFile = _tracking + try: + with pytest.raises(RuntimeError): + await _scan_full(source, FakeProxyLogging(_hook)) + finally: + bg.tempfile.SpooledTemporaryFile = real + + assert spools and all(handle.closed for handle in spools) + + +@pytest.mark.asyncio +async def test_the_rewrite_closes_its_own_output_when_it_cannot_finish(): + """A half-written rewrite spool has no owner yet, so it has to clean up after itself.""" + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + + source = _jsonl(_record("a"), _record("b", content="my secret is here")) + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + + spools = [] + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + def _boom(*args, **kwargs): + raise OSError("no space left on device") + + bg.tempfile.SpooledTemporaryFile = _tracking + original_read = bg._read_spooled + bg._read_spooled = _boom + try: + with pytest.raises(OSError): + rewrite_batch_input_file(source, result) + finally: + bg.tempfile.SpooledTemporaryFile = real + bg._read_spooled = original_read + + assert spools and all(handle.closed for handle in spools) + + +@pytest.mark.asyncio +async def test_the_scan_spool_is_closed_when_a_record_escapes_the_iterator(): + """A raise from inside the read loop bypasses the per-record outcome path entirely.""" + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + + spools = [] + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + bg.tempfile.SpooledTemporaryFile = _tracking + try: + with pytest.raises(json.JSONDecodeError): + await _scan_full(io.BytesIO(b"{not json at all}\n"), FakeProxyLogging()) + finally: + bg.tempfile.SpooledTemporaryFile = real + + assert spools and all(handle.closed for handle in spools) + + +@pytest.mark.asyncio +async def test_a_technical_failure_dressed_as_a_block_status_still_aborts(): + """xecguard and purview report an unreachable backend as HTTPException(400) under fail-closed.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + try: + raise ConnectionError("backend unreachable") + except ConnectionError as exc: + raise HTTPException( + status_code=400, detail={"error": "XecGuard API unreachable (block_on_error=True)"} + ) from exc + + with pytest.raises(HTTPException): + await _scan_full(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook)) + + +@pytest.mark.asyncio +async def test_a_record_body_cannot_opt_itself_out_of_the_guardrail_chain(): + """Guardrail selection reads a body-level `guardrails` key first; online it can only add.""" + seen = [] + + await _scan_full( + _jsonl({**_record("a"), "body": {**_record("a")["body"], "guardrails": []}}), + FakeProxyLogging(lambda d: seen.append(sorted(d))), + metadata={"guardrails": ["team-guard"]}, + ) + + assert seen and "guardrails" not in seen[0] + + +@pytest.mark.asyncio +async def test_a_redacted_record_keeps_its_own_guardrails_key(): + """Stripping it for the scan must not rewrite what the caller asked the provider to run.""" + record = _record("m", content="my secret is here") + record["body"]["guardrails"] = ["extra-guard"] + + body = await _rewritten_body(record, _redact_containing("secret")) + + assert body["guardrails"] == ["extra-guard"] + + +@pytest.mark.asyncio +async def test_a_400_that_is_not_a_guardrail_decision_still_aborts(): + """A guardrail's own HTTP client can raise a 400 because OUR payload was rejected, not the content.""" + from litellm.exceptions import BadRequestError + + def _hook(data): + raise BadRequestError(message="guardrail service rejected the payload", model="m", llm_provider="p") + + with pytest.raises(BadRequestError): + await _scan_full(_jsonl(_record("a")), FakeProxyLogging(_hook)) + + +async def _rewritten_body(record, hook): + """Scan one record and hand back the body as it lands in the uploaded file.""" + source = _jsonl(record) + result = await _scan_full(source, FakeProxyLogging(hook)) + rewritten = rewrite_batch_input_file(source, result) + return json.loads(rewritten.read().decode())["body"] + + +@pytest.mark.asyncio +async def test_a_redacted_record_keeps_its_own_body_metadata(): + """`metadata` is a real chat-completions parameter; redaction must not silently drop it.""" + record = _record("m", content="my secret is here") + record["body"]["metadata"] = {"team": "finance"} + + body = await _rewritten_body(record, _redact_containing("secret")) + + assert body["metadata"] == {"team": "finance"} + assert body["messages"][0]["content"] == "my *** is here" + assert "litellm_metadata" not in body + + +@pytest.mark.asyncio +async def test_a_redacted_record_keeps_its_own_litellm_metadata(): + """Tags ride in litellm_metadata; a guardrail firing must not change how the record is attributed.""" + record = _record("m", content="my secret is here") + record["body"]["litellm_metadata"] = {"tags": ["cost-center-42"]} + + body = await _rewritten_body(record, _redact_containing("secret")) + + assert body["litellm_metadata"] == {"tags": ["cost-center-42"]} + + +@pytest.mark.asyncio +async def test_a_redacted_record_keeps_an_explicitly_null_metadata(): + """An absent key and a null one are different records, so redaction must not collapse them.""" + record = _record("m", content="my secret is here") + record["body"]["metadata"] = None + + body = await _rewritten_body(record, _redact_containing("secret")) + + assert "metadata" in body and body["metadata"] is None + + +@pytest.mark.asyncio +async def test_the_log_summary_cannot_be_used_to_forge_log_lines(): + """custom_id is caller-supplied and lands in a log line, so control characters must not survive.""" + forged = "a\nWARNING: proxy shutting down" + result = await _scan_full(_jsonl(_record(forged, content="tripwire")), FakeProxyLogging(_blocking("tripwire"))) + + summary = result.summary() + + assert "\n" not in summary + assert "a WARNING: proxy shutting down" in summary + + +@pytest.mark.asyncio +async def test_the_log_summary_is_capped_so_one_upload_cannot_flood_it(): + records = [_record(f"row-{index}", content="tripwire") for index in range(60)] + result = await _scan_full(_jsonl(*records), FakeProxyLogging(_blocking("tripwire"))) + + summary = result.summary() + + assert summary.endswith("and 10 more") + assert "row-49" in summary and "row-50" not in summary + + +@pytest.mark.asyncio +async def test_the_scan_keeps_rewritten_records_off_the_heap(): + """A file whose records are mostly rewritten must not build a second copy of itself in memory.""" + import dataclasses + + bulky = "my secret is here" + ("x" * 50_000) + result = await _scan_full( + _jsonl(*(_record(str(index), content=bulky) for index in range(4))), + FakeProxyLogging(_redact_containing("secret")), + ) + + retained = sum( + len(value) + for change in result.changes + for value in (getattr(change, field.name) for field in dataclasses.fields(change)) + if isinstance(value, str) + ) + assert len(result.changes) == 4 + assert retained < 100, f"{retained} bytes of record text retained per scan" + assert result.redactions.tell() > 200_000 + + +@pytest.mark.asyncio +async def test_the_uploaded_file_is_what_the_loadbalancing_model_sniff_reads(): + """If line 1 is dropped, the router must not pick its model from a record nobody submitted.""" + dropped_first = { + "custom_id": "gone", + "url": "/v1/chat/completions", + "body": {"model": "model-a", "messages": [{"role": "user", "content": "tripwire"}]}, + } + kept = { + "custom_id": "kept", + "url": "/v1/chat/completions", + "body": {"model": "model-b", "messages": [{"role": "user", "content": "fine"}]}, + } + source = _jsonl(dropped_first, kept) + + result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire"))) + rewritten = rewrite_batch_input_file(source, result) + + first_line = json.loads((rewritten.seek(0), rewritten.read().decode())[1].splitlines()[0]) + assert first_line["custom_id"] == "kept" + assert first_line["body"]["model"] == "model-b" + + +@pytest.mark.asyncio +async def test_an_infrastructure_failure_outranks_a_redaction_and_aborts(): + """A record we could not inspect must abort the upload even when an earlier record was rewritten.""" + down = HTTPException(status_code=503, detail={"error": "guardrail service unavailable"}) + + def _hook(data): + content = data["messages"][0]["content"] + if content == "raiser": + raise down + if content == "redact": + data["messages"][0]["content"] = "***" + + source = _jsonl(_record("a", content="redact"), _record("b", content="raiser")) + + with pytest.raises(HTTPException) as raised: + await _scan_full(source, FakeProxyLogging(_hook)) + + assert raised.value is down + + +@pytest.mark.asyncio +async def test_the_earliest_unscannable_record_is_the_one_reported(): + source = _jsonl( + _record("a"), + {"custom_id": "bad-1", "url": "/v1/rerank", "body": {"model": "m"}}, + {"custom_id": "bad-2", "url": "/v1/rerank", "body": {"model": "m"}}, + ) + + failure = await _scan_full(source, FakeProxyLogging()) + + assert failure == UnscannableRecord(line_number=2, custom_id="bad-1", url="/v1/rerank") + + +@pytest.mark.asyncio +async def test_a_dropped_record_names_the_guardrail_from_an_enriched_http_detail(): + """litellm stamps guardrail_name into a block's detail dict; the report should carry it through.""" + blocked = HTTPException( + status_code=400, + detail={"error": "Violated guardrail policy", "guardrail_name": "zscaler"}, + ) + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise blocked + + result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail="zscaler"),) + + +@pytest.mark.asyncio +async def test_a_dropped_record_without_a_named_guardrail_reports_none(): + """An unnamed block still drops; the report just cannot say which guardrail did it.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise HTTPException(status_code=400, detail="blocked") + + result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail=None),) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 99fb19f0d60..bf9323cdc6a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3516,3 +3516,245 @@ def test_create_file_non_batch_purpose_skips_batch_validation(monkeypatch, llm_r assert response.status_code == 200, response.text assert len(forwarded_calls) == 1 + + +def _batch_upload(client_, content: bytes, purpose: str = "batch"): + return client_.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": purpose}, + headers={"Authorization": "Bearer test-key"}, + ) + + +@pytest.mark.parametrize( + "content, purpose, expected_status, expected_fragment", + [ + ( + b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",' + b'"body":{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"hi"}]}}\n', + "batch", + 200, + None, + ), + ( + b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",' + b'"body":{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"leak me"}]}}\n', + "batch", + 200, + None, + ), + ( + b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",' + b'"body":{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"leak me"}]}}\n', + "assistants", + 200, + None, + ), + (b"{ not json\n", "batch", 400, "line 1"), + ], +) +def test_batch_upload_runs_guardrails_on_each_record( + monkeypatch, llm_router: Router, content, purpose, expected_status, expected_fragment +): + """POST /v1/files with purpose=batch must reach the guardrail chain; other purposes must not.""" + import litellm + import litellm.proxy.openai_files_endpoints.files_endpoints as fe + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and "leak" in message["content"]: + message["content"] = "***" + return data + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + async def fake_route_create_file(**kwargs): + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + try: + resp = _batch_upload(client, content, purpose) + assert resp.status_code == expected_status, resp.text + if expected_fragment is not None: + assert expected_fragment in resp.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + ProxyLogging._callback_capabilities_cache.clear() + + +def test_batch_upload_redacts_per_record(monkeypatch, llm_router: Router): + """An offending record is submitted masked, matching what the online path does per request.""" + expected_custom_ids = ["keep-1", "dirty", "keep-2"] + import json as _json + + import litellm + import litellm.proxy.openai_files_endpoints.files_endpoints as fe + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and "leak" in message["content"]: + message["content"] = message["content"].replace("leak", "***") + return data + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + uploaded = {} + + async def fake_route_create_file(**kwargs): + handle = kwargs["_create_file_request"]["file"][1] + uploaded["body"] = handle.read() if hasattr(handle, "read") else handle + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + + def _row(custom_id, content): + return _json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": content}]}, + } + ) + + content = ("\n".join([_row("keep-1", "fine"), _row("dirty", "please leak this"), _row("keep-2", "fine")])).encode() + try: + resp = client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + rows = [_json.loads(line) for line in uploaded["body"].decode().splitlines()] + assert [row["custom_id"] for row in rows] == expected_custom_ids + assert rows[1]["body"]["messages"][0]["content"] == "please *** this" + report = resp.json()["litellm_batch_guardrail"] + assert report["submitted_records"] == 3 + assert report["modified_records"] == [ + {"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None} + ] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + ProxyLogging._callback_capabilities_cache.clear() + + +def test_batch_upload_closes_the_spools_it_opened(monkeypatch, llm_router: Router): + """The scan and the rewrite each open a spool; the request owns both and must not leak them.""" + import json as _json + + import litellm + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + import litellm.proxy.openai_files_endpoints.files_endpoints as fe + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + for message in data.get("messages") or []: + if "leak" in (message.get("content") or ""): + message["content"] = message["content"].replace("leak", "***") + return data + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + spools = [] + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + monkeypatch.setattr(bg.tempfile, "SpooledTemporaryFile", _tracking) + + async def fake_route_create_file(**kwargs): + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + + def _row(custom_id, content): + return _json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": content}]}, + } + ) + + content = ("\n".join([_row("keep", "fine"), _row("dirty", "please leak this")])).encode() + try: + resp = client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + assert len(spools) == 2, f"expected a scan spool and a rewrite spool, saw {len(spools)}" + assert all(handle.closed for handle in spools), "the request must close every spool it opened" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + ProxyLogging._callback_capabilities_cache.clear() diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_misc.py b/tests/test_litellm/proxy/proxy_server/test_routes_misc.py index e0e5a81cb2c..ad9b489b8f0 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_misc.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_misc.py @@ -225,8 +225,8 @@ def test_get_image_without_theme_still_serves_the_light_jpeg(client, monkeypatch def test_get_image_dark_theme_keeps_serving_a_custom_ui_logo(client, monkeypatch, tmp_path): - """A custom UI_LOGO_PATH has no dark variant yet, so dark mode must fall back to the - admin's own logo rather than replacing it with LiteLLM's.""" + """With no UI_LOGO_PATH_DARK set, dark mode falls back to the admin's own light logo + rather than replacing their branding with LiteLLM's.""" custom_logo = tmp_path / "custom.png" custom_logo.write_bytes(PNG_SIGNATURE + b"custom-logo-marker") monkeypatch.setenv("UI_LOGO_PATH", str(custom_logo)) @@ -235,6 +235,64 @@ def test_get_image_dark_theme_keeps_serving_a_custom_ui_logo(client, monkeypatch assert shape == {"status": 200, "body": PNG_SIGNATURE + b"custom-logo-marker"} +def test_get_image_dark_theme_prefers_the_dark_custom_logo(client, monkeypatch, tmp_path): + """UI_LOGO_PATH_DARK outranks UI_LOGO_PATH when the dark logo is requested.""" + light_logo = tmp_path / "light.png" + light_logo.write_bytes(PNG_SIGNATURE + b"light-marker") + dark_logo = tmp_path / "dark.png" + dark_logo.write_bytes(PNG_SIGNATURE + b"dark-marker") + monkeypatch.setenv("UI_LOGO_PATH", str(light_logo)) + monkeypatch.setenv("UI_LOGO_PATH_DARK", str(dark_logo)) + + response = client.get("/get_image", params={"theme": "dark"}) + + shape = {"status": response.status_code, "body": response.content} + assert shape == {"status": 200, "body": PNG_SIGNATURE + b"dark-marker"} + + +def test_get_image_unusable_dark_logo_falls_back_to_the_light_custom_logo(client, monkeypatch, tmp_path): + """A broken UI_LOGO_PATH_DARK must not drop the admin all the way to LiteLLM's own + logo while their light logo is still perfectly serviceable.""" + light_logo = tmp_path / "light.png" + light_logo.write_bytes(PNG_SIGNATURE + b"light-marker") + monkeypatch.setenv("UI_LOGO_PATH", str(light_logo)) + monkeypatch.setenv("UI_LOGO_PATH_DARK", str(tmp_path / "missing.png")) + + response = client.get("/get_image", params={"theme": "dark"}) + + shape = {"status": response.status_code, "body": response.content} + assert shape == {"status": 200, "body": PNG_SIGNATURE + b"light-marker"} + + +def test_get_image_light_theme_ignores_the_dark_custom_logo(client, monkeypatch, tmp_path): + """The dark logo must never leak into a light-mode request.""" + light_logo = tmp_path / "light.png" + light_logo.write_bytes(PNG_SIGNATURE + b"light-marker") + dark_logo = tmp_path / "dark.png" + dark_logo.write_bytes(PNG_SIGNATURE + b"dark-marker") + monkeypatch.setenv("UI_LOGO_PATH", str(light_logo)) + monkeypatch.setenv("UI_LOGO_PATH_DARK", str(dark_logo)) + + response = client.get("/get_image") + + shape = {"status": response.status_code, "body": response.content} + assert shape == {"status": 200, "body": PNG_SIGNATURE + b"light-marker"} + + +def test_get_image_dark_logo_alone_still_serves_the_bundled_light_logo_in_light_mode(client, monkeypatch): + """Setting only UI_LOGO_PATH_DARK leaves light mode on the bundled default.""" + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + monkeypatch.setenv("UI_LOGO_PATH_DARK", "https://cdn.example.invalid/logo-dark.png") + + response = client.get("/get_image") + + shape = { + "status": response.status_code, + "media_type": response.headers.get("content-type", "").split(";")[0], + } + assert shape == {"status": 200, "media_type": "image/jpeg"} + + def test_get_image_redirects_remote_url(client, monkeypatch): """Remote logo URLs are served via redirect — the proxy never fetches them server-side.""" monkeypatch.setenv("UI_LOGO_PATH", "https://example.invalid/logo.png") diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 355c6d27eb2..510fb977a61 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3,7 +3,7 @@ import copy import datetime import json from types import SimpleNamespace -from typing import AsyncGenerator, Callable, Optional +from typing import AsyncGenerator, Callable, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 8ee4e92ca9b..a1a02ec427b 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1064,11 +1064,15 @@ class TestProxySettingEndpoints: assert mock_proxy_config["save_call_count"]() == 1 # env vars are persisted through the dedicated per-key path, and ONLY - # the two keys this endpoint owns are touched. The unrelated SSO env + # the keys this endpoint owns are touched. The unrelated SSO env # vars in the merged config are never snapshotted. env_updates = mock_proxy_config["env_updates"]() assert env_updates == [ - {"UI_LOGO_PATH": "https://example.com/new-logo.png", "LITELLM_FAVICON_URL": None} + { + "UI_LOGO_PATH": "https://example.com/new-logo.png", + "UI_LOGO_PATH_DARK": None, + "LITELLM_FAVICON_URL": None, + } ] def test_update_ui_theme_settings_with_favicon( @@ -1097,14 +1101,90 @@ class TestProxySettingEndpoints: assert os.environ["UI_LOGO_PATH"] == "https://example.com/new-logo.png" assert os.environ["LITELLM_FAVICON_URL"] == "https://example.com/custom-favicon.ico" - # Only the two owned keys are persisted, both with their new values + # Only the owned keys are persisted, each with its new value assert mock_proxy_config["env_updates"]() == [ { "UI_LOGO_PATH": "https://example.com/new-logo.png", + "UI_LOGO_PATH_DARK": None, "LITELLM_FAVICON_URL": "https://example.com/custom-favicon.ico", } ] + def test_update_ui_theme_settings_with_dark_logo( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """A dark-mode logo is stored and applied to the live process like the light one.""" + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + new_theme = { + "logo_url": "https://example.com/logo.png", + "logo_url_dark": "https://example.com/logo-dark.png", + } + + response = client.patch("/update/ui_theme_settings", json=new_theme) + + assert response.status_code == 200 + assert response.json()["theme_config"]["logo_url_dark"] == "https://example.com/logo-dark.png" + assert os.environ["UI_LOGO_PATH_DARK"] == "https://example.com/logo-dark.png" + assert mock_proxy_config["env_updates"]() == [ + { + "UI_LOGO_PATH": "https://example.com/logo.png", + "UI_LOGO_PATH_DARK": "https://example.com/logo-dark.png", + "LITELLM_FAVICON_URL": None, + } + ] + + def test_update_ui_theme_settings_rejects_local_path_dark_logo( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """The dark logo is served by the unauthenticated /get_image, so a local + filesystem path must be refused exactly as it is for the light logo.""" + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + response = client.patch( + "/update/ui_theme_settings", + json={"logo_url_dark": "/etc/passwd"}, + ) + + assert response.status_code == 400 + assert "logo_url_dark" in str(response.json()) + + def test_update_ui_theme_settings_persists_every_env_var_it_resolves( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """Read and write must cover the same env vars. + + /get/ui_theme_settings resolves each field through _UI_THEME_FIELD_ENV_VARS, + so a var missing from the update path would read back from an env value the + save never cleared, and the settings page would show a field it cannot unset. + """ + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + _UI_THEME_FIELD_ENV_VARS, + ) + + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + response = client.patch("/update/ui_theme_settings", json={}) + + assert response.status_code == 200 + persisted = mock_proxy_config["env_updates"]() + assert len(persisted) == 1 + assert set(persisted[0]) == set(_UI_THEME_FIELD_ENV_VARS.values()) + + def test_get_ui_theme_settings_surfaces_dark_logo_from_process_env( + self, mock_proxy_config, monkeypatch + ): + """A dark logo supplied only as a process env var must surface in the read.""" + monkeypatch.setenv("UI_LOGO_PATH_DARK", "https://cdn.example.com/logo-dark.png") + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + assert response.json()["values"]["logo_url_dark"] == "https://cdn.example.com/logo-dark.png" + def test_update_ui_theme_settings_clear_favicon( self, mock_proxy_config, mock_auth, monkeypatch ): diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 05005dae797..12fc9310d48 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -166,3 +166,123 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox ) assert out is data assert invoked["data"] is data + + +@pytest.mark.asyncio +async def test_guardrails_only_skips_non_guardrail_pre_call_callbacks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """Rate limiters and budget hooks ride this same loop; a content scan must not trip them.""" + calls: list[str] = [] + + class _RateLimiterLike(CustomLogger): + async def async_pre_call_hook(self, **kwargs): # type: ignore[override] + calls.append("ran") + return None + + monkeypatch.setattr(litellm, "callbacks", [_RateLimiterLike()]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m"}, + call_type="completion", + guardrails_only=True, + ) + assert calls == [] + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m"}, + call_type="completion", + ) + assert calls == ["ran"], "the default path must still run non-guardrail pre-call callbacks" + + +@pytest.mark.asyncio +async def test_guardrails_only_skips_the_hanging_request_alert(proxy_logging, make_user_api_key_auth, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + alerting = MagicMock(alerting=True) + proxy_logging.slack_alerting_instance = alerting + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m"}, + call_type="completion", + guardrails_only=True, + ) + alerting.response_taking_too_long.assert_not_called() + + +@pytest.mark.asyncio +async def test_guardrails_only_skips_prompt_template_rewriting(proxy_logging, make_user_api_key_auth, monkeypatch): + """A prompt template would rewrite messages, which a per-record content diff would misread.""" + monkeypatch.setattr(litellm, "callbacks", []) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + process = AsyncMock() + monkeypatch.setattr(proxy_logging, "_process_prompt_template", process) + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m", "prompt_id": "p1", "litellm_logging_obj": MagicMock()}, + call_type="acompletion", + guardrails_only=True, + ) + process.assert_not_awaited() + + +@pytest.mark.parametrize( + "event_hook, expected", + [("pre_call", True), ("post_call", False), ("during_call", False)], +) +def test_has_pre_call_guardrails_follows_the_guardrail_event_hook(proxy_logging, monkeypatch, event_hook, expected): + """A post-call-only guardrail must not make callers pay for pre-call work.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + + guardrail = CustomGuardrail(guardrail_name="g", event_hook=event_hook, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + assert proxy_logging.has_pre_call_guardrails({}) is expected + + +def test_has_pre_call_guardrails_is_false_without_callbacks(proxy_logging, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + + assert proxy_logging.has_pre_call_guardrails({}) is False + + +def test_has_pre_call_guardrails_is_true_for_a_configured_pipeline(proxy_logging, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + + assert proxy_logging.has_pre_call_guardrails({"_guardrail_pipelines": ["p1"]}) is True + + +@pytest.mark.asyncio +async def test_default_path_still_arms_the_hanging_request_alert(proxy_logging, make_user_api_key_auth, monkeypatch): + """Pins the other side of the gate: without the flag, the alert must still fire.""" + monkeypatch.setattr(litellm, "callbacks", []) + alerting = MagicMock(alerting=True) + alerting.response_taking_too_long = AsyncMock() + proxy_logging.slack_alerting_instance = alerting + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m"}, + call_type="completion", + ) + alerting.response_taking_too_long.assert_called_once() + + +@pytest.mark.asyncio +async def test_default_path_still_applies_prompt_templates(proxy_logging, make_user_api_key_auth, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + process = AsyncMock() + monkeypatch.setattr(proxy_logging, "_process_prompt_template", process) + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m", "prompt_id": "p1", "litellm_logging_obj": MagicMock()}, + call_type="acompletion", + ) + process.assert_awaited_once() diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 7ea6aa5b2a4..564f32e2573 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -24,6 +24,7 @@ "lucide-react": "0.513.0", "moment": "2.30.1", "next": "16.2.11", + "next-themes": "^0.4.6", "nuqs": "^2.9.4", "openai": "4.104.0", "openapi-fetch": "^0.17.0", @@ -9805,6 +9806,16 @@ } } }, + "node_modules/next-themes": { + "version": "0.4.6", + "resolved": "https://registry.npmjs.org/next-themes/-/next-themes-0.4.6.tgz", + "integrity": "sha512-pZvgD5L0IEvX5/9GWyHMf3m8BKiVQwsCMHfoFosXtXBMnaS0ZnIJ9ST4b4NqLVKDEm8QBxoNNGNaBv2JNF6XNA==", + "license": "MIT", + "peerDependencies": { + "react": "^16.8 || ^17 || ^18 || ^19 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17 || ^18 || ^19 || ^19.0.0-rc" + } + }, "node_modules/node-domexception": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/node-domexception/-/node-domexception-1.0.0.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 6b2d4e6106b..ff6448ad75c 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -40,6 +40,7 @@ "lucide-react": "0.513.0", "moment": "2.30.1", "next": "16.2.11", + "next-themes": "^0.4.6", "nuqs": "^2.9.4", "openai": "4.104.0", "openapi-fetch": "^0.17.0", 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 cbb7fc7aa33..4b58590ffd1 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 @@ -7,10 +7,18 @@ import { toast } from "@/lib/toast"; import UIThemeSettings from "./UIThemeSettings"; const setLogoUrl = vi.fn(); +const setLogoUrlDark = vi.fn(); const setFaviconUrl = vi.fn(); vi.mock("@/contexts/ThemeContext", () => ({ - useTheme: () => ({ logoUrl: null, setLogoUrl, faviconUrl: null, setFaviconUrl }), + useTheme: () => ({ + logoUrl: null, + setLogoUrl, + logoUrlDark: null, + setLogoUrlDark, + faviconUrl: null, + setFaviconUrl, + }), })); vi.mock("@/components/networking", () => ({ @@ -19,6 +27,7 @@ vi.mock("@/components/networking", () => ({ })); const LOGO_PLACEHOLDER = "https://example.com/logo.png"; +const LOGO_DARK_PLACEHOLDER = "https://example.com/logo-dark.png"; const FAVICON_PLACEHOLDER = "https://example.com/favicon.ico"; const okResponse = (values: Record = {}) => @@ -76,11 +85,32 @@ describe("UIThemeSettings", () => { await waitFor(() => expect(patchCalls()).toHaveLength(1)); expect(bodyOf(patchCalls()[0])).toEqual({ logo_url: "https://a.test/logo.png", + logo_url_dark: null, favicon_url: "https://a.test/fav.ico", }); await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Theme settings updated successfully!")); }); + it("should load and save a separate dark-mode logo url", async () => { + const user = userEvent.setup(); + fetchMock.mockImplementation(() => okResponse({ logo_url_dark: "https://cdn.example.com/logo-dark.svg" })); + + render(); + + await waitFor(() => { + expect(screen.getByPlaceholderText(LOGO_DARK_PLACEHOLDER)).toHaveValue("https://cdn.example.com/logo-dark.svg"); + }); + expect(setLogoUrlDark).toHaveBeenCalledWith("https://cdn.example.com/logo-dark.svg"); + + fireEvent.change(screen.getByPlaceholderText(LOGO_DARK_PLACEHOLDER), { + target: { value: "https://a.test/logo-dark.png" }, + }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => expect(patchCalls()).toHaveLength(1)); + expect(bodyOf(patchCalls()[0]).logo_url_dark).toBe("https://a.test/logo-dark.png"); + }); + it("should surface a backend failure when saving fails", async () => { const user = userEvent.setup(); render(); @@ -109,10 +139,12 @@ describe("UIThemeSettings", () => { await user.click(screen.getByRole("button", { name: "Reset to Default" })); await waitFor(() => expect(patchCalls()).toHaveLength(1)); - expect(bodyOf(patchCalls()[0])).toEqual({ logo_url: null, favicon_url: null }); + expect(bodyOf(patchCalls()[0])).toEqual({ logo_url: null, logo_url_dark: null, favicon_url: null }); expect(screen.getByPlaceholderText(LOGO_PLACEHOLDER)).toHaveValue(""); + expect(screen.getByPlaceholderText(LOGO_DARK_PLACEHOLDER)).toHaveValue(""); expect(screen.getByPlaceholderText(FAVICON_PLACEHOLDER)).toHaveValue(""); expect(setLogoUrl).toHaveBeenLastCalledWith(null); + expect(setLogoUrlDark).toHaveBeenLastCalledWith(null); expect(setFaviconUrl).toHaveBeenLastCalledWith(null); await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Theme settings reset to default!")); }); 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 8d7a0fa3de2..b6bc0f28dd5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx @@ -15,8 +15,9 @@ interface UIThemeSettingsProps { } const UIThemeSettings: React.FC = ({ userID, userRole, accessToken }) => { - const { setLogoUrl, setFaviconUrl } = useTheme(); + const { setLogoUrl, setLogoUrlDark, setFaviconUrl } = useTheme(); const [logoUrlInput, setLogoUrlInput] = useState(""); + const [logoUrlDarkInput, setLogoUrlDarkInput] = useState(""); const [faviconUrlInput, setFaviconUrlInput] = useState(""); const [loading, setLoading] = useState(false); @@ -40,8 +41,10 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc if (response.ok) { const data = await response.json(); 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); } } catch (error) { @@ -62,12 +65,14 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc }, body: JSON.stringify({ logo_url: logoUrlInput || null, + logo_url_dark: logoUrlDarkInput || null, favicon_url: faviconUrlInput || null, }), }); if (response.ok) { toast.success("Theme settings updated successfully!"); setLogoUrl(logoUrlInput || null); + setLogoUrlDark(logoUrlDarkInput || null); setFaviconUrl(faviconUrlInput || null); } else { throw new Error("Failed to update settings"); @@ -82,8 +87,10 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc const handleReset = async () => { setLogoUrlInput(""); + setLogoUrlDarkInput(""); setFaviconUrlInput(""); setLogoUrl(null); + setLogoUrlDark(null); setFaviconUrl(null); setLoading(true); try { @@ -95,7 +102,7 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, - body: JSON.stringify({ logo_url: null, favicon_url: null }), + body: JSON.stringify({ logo_url: null, logo_url_dark: null, favicon_url: null }), }); if (response.ok) { toast.success("Theme settings reset to default!"); @@ -141,6 +148,23 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc Enter a URL for your custom logo or leave empty for default

+
+ + { + setLogoUrlDarkInput(event.target.value); + setLogoUrlDark(event.target.value || null); + }} + /> +

+ Enter a URL for a logo suited to dark backgrounds, or leave empty to reuse the logo above +

+
diff --git a/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.test.tsx b/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.test.tsx new file mode 100644 index 00000000000..41efedcd43a --- /dev/null +++ b/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.test.tsx @@ -0,0 +1,77 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { ThemeProvider } from "next-themes"; +import { afterAll, beforeEach, describe, expect, it } from "vitest"; +import ThemeToggle from "./ThemeToggle"; + +const renderToggle = () => + render( + + + , + ); + +const openMenu = async () => { + await userEvent.click(screen.getByRole("button", { name: "Theme" })); + await screen.findByRole("menu"); +}; + +const pick = async (label: string | RegExp) => userEvent.click(screen.getByRole("menuitemradio", { name: label })); + +beforeEach(() => { + localStorage.clear(); + document.documentElement.classList.remove("dark", "light"); +}); + +afterAll(() => { + document.documentElement.classList.remove("dark", "light"); +}); + +describe("ThemeToggle", () => { + it("starts on light rather than following the system preference", async () => { + renderToggle(); + await openMenu(); + + expect(screen.getByRole("menuitemradio", { name: "Light" })).toBeChecked(); + expect(screen.getByRole("menuitemradio", { name: /^Dark/ })).not.toBeChecked(); + expect(screen.getByRole("menuitemradio", { name: "System" })).not.toBeChecked(); + }); + + it("puts the dark class on the document and remembers the choice", async () => { + renderToggle(); + await openMenu(); + + await pick(/^Dark/); + + expect(document.documentElement).toHaveClass("dark"); + expect(localStorage.getItem("theme")).toBe("dark"); + }); + + it("hands control back to the system preference when asked", async () => { + renderToggle(); + await openMenu(); + await pick(/^Dark/); + + await pick("System"); + + expect(localStorage.getItem("theme")).toBe("system"); + expect(document.documentElement).not.toHaveClass("dark"); + }); + + it("marks dark as beta in the menu, and leaves the other choices unmarked", async () => { + renderToggle(); + await openMenu(); + + expect(screen.getByRole("menuitemradio", { name: /^Dark/ })).toHaveTextContent("Beta"); + expect(screen.getByRole("menuitemradio", { name: "Light" })).not.toHaveTextContent("Beta"); + expect(screen.getByRole("menuitemradio", { name: "System" })).not.toHaveTextContent("Beta"); + }); + + it("keeps the beta marker inside the menu rather than in the toolbar", async () => { + renderToggle(); + await openMenu(); + await pick(/^Dark/); + + expect(screen.getByRole("button", { name: "Theme" })).not.toHaveTextContent("Beta"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.tsx b/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.tsx new file mode 100644 index 00000000000..3fbb3d2eb5b --- /dev/null +++ b/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.tsx @@ -0,0 +1,58 @@ +"use client"; + +import { Monitor, Moon, Sun } from "lucide-react"; +import { useTheme } from "next-themes"; +import React from "react"; + +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuRadioGroup, + DropdownMenuRadioItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; + +const THEMES = [ + { value: "system", label: "System", Icon: Monitor, beta: false }, + { value: "light", label: "Light", Icon: Sun, beta: false }, + { value: "dark", label: "Dark", Icon: Moon, beta: true }, +] as const; + +const ThemeToggle: React.FC = () => { + const { theme, setTheme, resolvedTheme } = useTheme(); + + return ( + + + } + > + {resolvedTheme === "dark" ? : } + + + + {THEMES.map(({ value, label, Icon, beta }) => ( + + + {label} + {beta && ( + + Beta + + )} + + ))} + + + + ); +}; + +export default ThemeToggle; diff --git a/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx b/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx index b5291a34cb6..0b47bc8500a 100644 --- a/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx +++ b/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx @@ -14,7 +14,8 @@ export type ComplexityTier = "SIMPLE" | "MEDIUM" | "COMPLEX" | "REASONING"; export interface KeywordTierRule { id: string; keywords: string[]; - tier: ComplexityTier; + /** A built-in tier name, or with a custom tier set, one of the defined tier names. */ + tier: string; } interface KeywordTierRulesProps { @@ -99,7 +100,7 @@ const KeywordTierRules: React.FC = ({ rules, onChange, ti