From 2d82915084f063010771865e4e5bcd948bcbe880 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 5 Oct 2026 10:25:47 -0700 Subject: [PATCH 001/191] fix(mcp): bind OAuth clients to their upstream issuer (#37777) * feat(mcp): advertise the SDK's latest spec revision and validate the RFC 9207 iss MCPSpecVersion stopped at 2025-06-18 while the pinned SDK negotiates 2025-11-25, and the version LiteLLM puts on its own outbound initialize was a hardcoded historical member. Add the missing revision, name the highest revision we speak once, and pin it to the SDK's LATEST_PROTOCOL_VERSION with a test so the two cannot drift apart silently. /authorize now seals the issuer it sent the user to into the OAuth state, and /callback holds the authorization response's RFC 9207 iss against it, refusing to forward a code that came back from an authorization server we never sent the user to. An absent iss, an unanchored server row and a state minted before the seal all keep their current behavior. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): keep params, query and fragment significant in issuer comparison The shared canonicalizer drops all three, so two issuers differing only outside the path compared equal and a response from another tenant's authorization server would have continued through the flow. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): refresh generated API snapshots Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): cover OAuth client isolation and lint checks * fix(mcp): preserve registered clients in the existing save payload * test(mcp): cover optional OAuth registration metadata * fix(mcp): preserve compatible OAuth registrations across edits * fix(mcp): retain OAuth state through pending authorization * fix(mcp): guard pending OAuth at form submission * fix(mcp): discard canceled OAuth edit snapshots * test(mcp): preserve complete OAuth registration assertions * refactor(mcp): construct OAuth credential updates without mutation * fix(mcp): simplify issuer binding and reject unverifiable callbacks * fix(mcp): preserve replacement clients and pending redirect bindings * fix(mcp): retain clients with replacement authentication methods * fix(mcp): preserve cached clients and pin manual OAuth issuers --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/proxy/_experimental/mcp_server/db.py | 168 ++++++- .../mcp_server/discoverable_endpoints.py | 202 ++++++-- .../mcp_server/mcp_server_manager.py | 45 +- .../_experimental/mcp_server/oauth_utils.py | 11 + litellm/proxy/_lazy_openapi_snapshot.json | 106 +++- .../mcp_management_endpoints.py | 72 ++- litellm/types/mcp.py | 5 +- .../types/mcp_server/mcp_server_manager.py | 4 + .../e2e/mcp/test_mcp_oauth_happy_path_e2e.py | 2 + tests/e2e/models.py | 1 + .../mcp_server/test_discoverable_endpoints.py | 463 +++++++++++++++++- .../mcp_server/test_mcp_partial_update.py | 184 +++++++ .../mcp_server/test_mcp_server_manager.py | 16 +- .../test_mcp_management_endpoints.py | 181 ++++++- .../CreateMCPServer.integration.test.tsx | 15 +- .../_components/CreateMCPServer.tsx | 11 +- .../_components/createServerPayload.ts | 2 +- .../editServerPayload.differential.test.ts | 28 ++ .../_components/editServerPayload.ts | 11 + .../_components/mcp_server_edit.test.tsx | 144 +++++- .../_components/mcp_server_edit.tsx | 39 +- .../src/hooks/useMcpOAuthFlow.test.tsx | 130 ++++- .../src/hooks/useMcpOAuthFlow.tsx | 64 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 9 +- 24 files changed, 1739 insertions(+), 174 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index bf1f6fa90c1..b85ff2ce50e 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -20,7 +20,9 @@ from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( credential_binding_matches, enforce_oauth_identity_binding, ) -from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + build_upstream_oauth2_token_request, +) from litellm.proxy._types import ( LiteLLM_MCPServerTable, MCPApprovalStatus, @@ -105,7 +107,31 @@ _AUTH_FLOW_SCOPED_FIELDS: Final["frozenset[str]"] = frozenset( ) -def _blank_to_none(value: str | None) -> str | None: +_OAUTH_CLIENT_CREDENTIAL_FIELDS: Final = frozenset( + { + "client_id", + "client_secret", + "token_endpoint_auth_method", + "redirect_uris", + "dcr_issuer", + "dcr_server_url", + "access_token", + "refresh_token", + "expires_in", + "scope", + } +) + + +def stale_mcp_auth_fields(submitted: Mapping[str, object], previous_value: Callable[[str], object]) -> dict[str, None]: + return { + field: None + for field in _AUTH_FLOW_SCOPED_FIELDS + if field not in submitted or submitted[field] == previous_value(field) + } + + +def _blank_to_none(value: object) -> str | None: if not isinstance(value, str): return None return value.strip() or None @@ -137,6 +163,56 @@ _CLIENT_FORWARDED_AUTH_TYPES: Final["frozenset[str]"] = frozenset({"true_passthr _MINTED_TOKEN_CREDENTIAL_FIELDS: Final["frozenset[str]"] = frozenset({"access_token", "refresh_token", "expires_in"}) +def _bind_submitted_oauth_client( + credentials: Mapping[str, object], issuer: str | None, url: str | None +) -> dict[str, object]: + if not credentials.get("client_id") or credentials.get("dcr_issuer") or credentials.get("dcr_server_url"): + return dict(credentials) + return {**credentials, "dcr_issuer": issuer, "dcr_server_url": url} + + +def is_resubmitted_oauth_client(supplied: dict[str, object], existing: dict[str, object]) -> bool: + """Recognize the saved client without conflating same-ID clients from different issuers.""" + return bool( + supplied.get("client_id") + and _decrypted_credential_field(supplied, "client_id") == _decrypted_credential_field(existing, "client_id") + and all( + _decrypted_credential_field(supplied, field) == _decrypted_credential_field(existing, field) + for field in ("client_secret", "token_endpoint_auth_method", "dcr_issuer", "dcr_server_url") + if field in supplied + ) + ) + + +def oauth_credentials_for_upstream_edit( + credentials: Mapping[str, object], + previous_issuer: str | None, + previous_url: str | None, + *, + issuer_changed: bool, +) -> dict[str, object]: + """Bind an existing client on an explicit edit, so discovery can verify reuse at the new URL. + + No first-use backfill: unchanged legacy rows never enter this path. Without a known previous + issuer, or after a known issuer change, the old client cannot be carried to the new resource. + """ + registered_issuer: Final = _blank_to_none(credentials.get("dcr_issuer")) or previous_issuer + keep_client: Final = bool(registered_issuer and credentials.get("client_id") and not issuer_changed) + removed: Final = ( + _MINTED_TOKEN_CREDENTIAL_FIELDS | {"auth_value"} + if keep_client + else _OAUTH_CLIENT_CREDENTIAL_FIELDS | {"auth_value"} + ) + retained: Final = {key: value for key, value in credentials.items() if key not in removed} + if keep_client: + return { + **retained, + "dcr_issuer": registered_issuer, + "dcr_server_url": credentials.get("dcr_server_url") or previous_url, + } + return retained + + class _OAuthCredentialAccessToken(TypedDict): access_token: str @@ -386,7 +462,11 @@ def _prepare_mcp_server_data( if blob_value is not None and te_field not in data_dict: data_dict[te_field] = blob_value data_dict["credentials"] = encrypt_credentials(credentials=credentials, encryption_key=_get_salt_key()) - data_dict["credentials"] = safe_dumps(data_dict["credentials"]) + data_dict["credentials"] = safe_dumps( + _bind_submitted_oauth_client(data_dict["credentials"], data.issuer, data.url) + if not exclude_unset and data.auth_type == "oauth2" + else data_dict["credentials"] + ) # Serialize JSON fields from ``data_dict`` (not ``data``) so the # exclude_unset filter is respected. Reading back from ``data`` would @@ -1132,6 +1212,7 @@ async def _update_mcp_server_row( *, server_id: str, data_dict: Mapping[str, object], + expected_updated_at: datetime | None = None, ) -> "prisma_db_models.LiteLLM_MCPServerTable | McpIdentifierConflict | None": identifier_write: Final = any(field in data_dict for field in ("server_name", "alias")) protocol_write: Final = bool({"transport", "mcp_info"}.intersection(data_dict)) @@ -1144,6 +1225,11 @@ async def _update_mcp_server_row( if stored is None: return None _validate_mcp_protocol_write(stored, data_dict) + if expected_updated_at is not None: + changed: Final = await table.update_many( + where={"server_id": server_id, "updated_at": expected_updated_at}, data=data_dict + ) + return await table.find_unique(where={"server_id": server_id}) if changed else None return await table.update( where={"server_id": server_id}, data=data_dict, @@ -1180,6 +1266,7 @@ async def update_mcp_server( data: UpdateMCPServerRequest, touched_by: str, fields_set: set[str] | None = None, + expected_updated_at: datetime | None = None, ) -> LiteLLM_MCPServerTable | McpIdentifierConflict | None: """ Update a new mcp server record in the db @@ -1193,16 +1280,16 @@ async def update_mcp_server( # exclude_unset=True makes this a true partial update: fields the caller did # not provide are not written, so they keep their existing DB value instead # of being reset to a schema default (transport=sse, allow_all_keys=False...). - data_dict: Final = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set) + prepared_data: Final = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set) # Pre-fetch existing record once if we need it for auth_type, url, or credential logic existing = None - has_credentials: Final = "credentials" in data_dict and data_dict["credentials"] is not None + has_credentials: Final = "credentials" in prepared_data and prepared_data["credentials"] is not None # An explicit token-exchange column write (set or clear) also migrates the # legacy blob copies below, so the existing row is needed for those updates. - explicit_te_write: Final = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & data_dict.keys()) - url_provided: Final = "url" in data_dict and data_dict["url"] is not None - issuer_provided: Final = "issuer" in data_dict + explicit_te_write: Final = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & prepared_data.keys()) + url_provided: Final = "url" in prepared_data and prepared_data["url"] is not None + issuer_provided: Final = "issuer" in prepared_data if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided: existing = await _db_find_mcp_server_row(prisma_client, data.server_id) @@ -1213,29 +1300,60 @@ async def update_mcp_server( ) # A url change re-points the server at a potentially different upstream, so any discovered or # trust-on-first-use OAuth endpoints/issuer belong to the old upstream and must re-discover. - url_changed: Final = bool(url_provided and existing and existing.url != data_dict["url"]) + url_changed: Final = bool(url_provided and existing and existing.url != prepared_data["url"]) old_issuer: Final = _blank_to_none(getattr(existing, "issuer", None)) if existing else None issuer_changed: Final = bool( - issuer_provided and old_issuer is not None and _blank_to_none(data_dict.get("issuer")) != old_issuer + issuer_provided and existing is not None and _blank_to_none(prepared_data.get("issuer")) != old_issuer ) + oauth_upstream_edited: Final = bool(existing and existing.auth_type == "oauth2" and (url_changed or issuer_changed)) + existing_credentials: Final = _credentials_blob_to_mutable_dict((existing.credentials or {}) if existing else {}) + retained_credentials: Final = ( + oauth_credentials_for_upstream_edit( + existing_credentials, + old_issuer, + existing.url if existing else None, + issuer_changed=issuer_changed or auth_type_changed, + ) + if oauth_upstream_edited + else existing_credentials + ) + edited_credentials: Final = {"credentials": safe_dumps(retained_credentials)} if oauth_upstream_edited else {} + cleared_auth_fields: Final = ( + stale_mcp_auth_fields(prepared_data, lambda field: getattr(existing, field, None)) + if auth_type_changed or url_changed or issuer_changed + else {} + ) + supplied: Final = _credentials_blob_to_mutable_dict(prepared_data.get("credentials") or {}) + safe_submitted: Final = ( + oauth_credentials_for_upstream_edit( + supplied, + old_issuer, + existing.url if existing else None, + issuer_changed=issuer_changed or auth_type_changed, + ) + if oauth_upstream_edited and is_resubmitted_oauth_client(supplied, existing_credentials) + else supplied + ) + submitted_credentials: Final = ( + { + "credentials": safe_dumps( + _bind_submitted_oauth_client( + safe_submitted, + _blank_to_none(cleared_auth_fields.get("issuer", prepared_data.get("issuer", old_issuer))), + _blank_to_none(prepared_data.get("url", existing.url if existing else None)), + ) + ) + } + if has_credentials and (data.auth_type or (existing.auth_type if existing else None)) == "oauth2" + else {} + ) + data_dict: Final = {**edited_credentials, **prepared_data, **submitted_credentials, **cleared_auth_fields} + # Clear stale credentials when auth_type changes but no new credentials provided if auth_type_changed and "credentials" not in data_dict: data_dict["credentials"] = None - if auth_type_changed or url_changed or issuer_changed: - # Clear each auth-flow-scoped field that the caller either omitted (partial update) or - # resubmitted unchanged. The edit form re-sends every field, so a stale issuer/endpoint - # belonging to the old upstream would otherwise survive a url/auth_type change and win in the - # resolution merge; only a genuinely new submitted value is kept. - data_dict.update( - { - field: None - for field in _AUTH_FLOW_SCOPED_FIELDS - if field not in data_dict or data_dict[field] == getattr(existing, field, None) - } - ) - # An explicit column write that does not touch credentials must still migrate # the row's legacy blob copies: lift values for columns the caller left # untouched, strip every copy from the blob. Without this, clearing a column @@ -1263,11 +1381,10 @@ async def update_mcp_server( # within the client-forwarded class (true_passthrough ↔ oauth_delegate) keeps # the same declared app and so must merge, not replace. if not auth_type_changed: - existing_creds = _credentials_blob_to_mutable_dict(existing.credentials) new_creds: Final = _credentials_blob_to_mutable_dict(data_dict["credentials"]) # New values override existing; existing keys not in update are preserved. A client # rotation additionally drops the previous app's stale minted token keys. - merged: Final = _drop_stale_minted_on_client_rotation({**existing_creds, **new_creds}, new_creds) + merged: Final = _drop_stale_minted_on_client_rotation({**retained_credentials, **new_creds}, new_creds) # Migrate-on-write for legacy rows: token-exchange settings the # old blob shape carried move to their dedicated columns (unless # the caller set the column this update, or the row already has @@ -1298,6 +1415,7 @@ async def update_mcp_server( prisma_client, server_id=data.server_id, data_dict=data_dict, + expected_updated_at=expected_updated_at, ) if isinstance(updated_mcp_server, McpIdentifierConflict): diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 7a0f59c3c2b..342c650d2a6 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,7 +12,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Depends, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response -from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError +from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -75,6 +75,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, build_upstream_oauth2_token_request, get_request_base_url, + oauth_client_registration_matches, resolve_upstream_resource, validate_trusted_redirect_uri, well_known_root_suffix, @@ -200,6 +201,8 @@ def encode_state_with_base_url( dcr_client_id: str | None = None, dcr_client_secret: str | None = None, dcr_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None, + expected_issuer: str | None = None, + authorization_response_iss_parameter_supported: bool = False, oauth_nonce: str | None = None, ) -> str: """ @@ -225,6 +228,8 @@ def encode_state_with_base_url( response granted the minted client, sealed alongside the credentials so the exchange authenticates the way the upstream expects instead of falling back to the server row's configured method + expected_issuer: Issuer identifier of the authorization server this flow is being sent to, + sealed so /callback can hold the RFC 9207 ``iss`` of the response against it Returns: An encrypted string that encodes all values @@ -241,6 +246,8 @@ def encode_state_with_base_url( "dcr_client_id": dcr_client_id, "dcr_client_secret": dcr_client_secret, "dcr_token_endpoint_auth_method": dcr_token_endpoint_auth_method, + "expected_issuer": expected_issuer, + "authorization_response_iss_parameter_supported": authorization_response_iss_parameter_supported, } state_json: Final = json.dumps(state_data, sort_keys=True) encrypted_state: Final = encrypt_value_helper(state_json) @@ -807,7 +814,16 @@ def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool: returns directly to the client's redirect URI without transiting the gateway. Gateway-side redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit arm, where the upstream only knows the gateway's own callback.""" - return mcp_server.is_dcr_bridge and bool(mcp_server.effective_registration_url) and not mcp_server.client_id + return bool( + mcp_server.is_dcr_bridge + and mcp_server.effective_registration_url + and not ( + mcp_server.client_id + and oauth_client_registration_matches( + mcp_server.dcr_issuer, mcp_server.dcr_server_url, mcp_server.issuer, mcp_server.url + ) + ) + ) def _require_s256_pkce( @@ -932,6 +948,10 @@ async def authorize_with_server( ): _raise_if_not_oauth2(mcp_server) resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint) + if not oauth_client_registration_matches( + resolved_server.dcr_issuer, resolved_server.dcr_server_url, resolved_server.issuer, resolved_server.url + ): + raise HTTPException(status_code=400, detail="OAuth client belongs to a different issuer; register a new client") if resolved_server.effective_authorization_url is None: raise HTTPException( status_code=400, @@ -1003,6 +1023,8 @@ async def authorize_with_server( dcr_token_endpoint_auth_method=ephemeral_dcr_client.token_endpoint_auth_method if ephemeral_dcr_client else None, + expected_issuer=resolved_server.issuer, + authorization_response_iss_parameter_supported=resolved_server.authorization_response_iss_parameter_supported, ) relay_state: Final = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES) @@ -1065,6 +1087,10 @@ async def exchange_token_with_server( raise HTTPException(status_code=400, detail="Unsupported grant_type") resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _token_flow_needed_endpoint) + if not oauth_client_registration_matches( + resolved_server.dcr_issuer, resolved_server.dcr_server_url, resolved_server.issuer, resolved_server.url + ): + raise HTTPException(status_code=400, detail="OAuth client belongs to a different issuer; register a new client") token_url: Final = resolved_server.effective_token_url if token_url is None: raise HTTPException( @@ -1364,6 +1390,8 @@ class _DcrClientRegistration(BaseModel): class _PersistedDcrCredentials(BaseModel): + dcr_issuer: str | None = None + dcr_server_url: str | None = None client_id: str | None = None client_secret: str | None = None token_endpoint_auth_method: str | None = None @@ -1410,19 +1438,28 @@ def _decrypt_persisted_dcr_credential(value: str | None, key: str) -> str | None def _apply_persisted_dcr_credentials(mcp_server: MCPServer, credentials: _PersistedDcrCredentials) -> bool: + if not oauth_client_registration_matches( + credentials.dcr_issuer, credentials.dcr_server_url, mcp_server.issuer, mcp_server.url + ): + return False client_id: Final = _decrypt_persisted_dcr_credential(credentials.client_id, "client_id") if not client_id: return False + mcp_server.dcr_issuer = credentials.dcr_issuer # rebind-ok: publish through the existing bool hydration contract + mcp_server.dcr_server_url = credentials.dcr_server_url # rebind-ok: retain binding in database-free shared state mcp_server.client_id = client_id mcp_server.client_secret = _decrypt_persisted_dcr_credential(credentials.client_secret, "client_secret") mcp_server.token_endpoint_auth_method = credentials.token_endpoint_auth_method return True -async def _load_store_dcr_credentials(mcp_server: MCPServer) -> _PersistedDcrCredentials | None: +async def _load_store_dcr_credentials( + mcp_server: MCPServer, *, raise_on_error: bool = False +) -> _PersistedDcrCredentials | None: """DCR client persisted in the server-scoped OAuth-client store for a config-declared server (which has no LiteLLM_MCPServerTable row). Returns None when the store has no usable client_id - or the DB is unreachable.""" + or the DB is unreachable. Registration writes request strict reads so a failed lookup cannot + be mistaken for an absent client.""" from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # avoids circular import get_mcp_server_oauth_client_credentials, ) @@ -1434,6 +1471,8 @@ async def _load_store_dcr_credentials(mcp_server: MCPServer) -> _PersistedDcrCre prisma_client=prisma_client, server_id=mcp_server.server_id ) except Exception as exc: # noqa: BLE001 # best-effort read; DB may be unreachable + if raise_on_error: + raise verbose_logger.debug( "register_client_with_server: failed to read stored DCR client for server_id=%s: %s", mcp_server.server_id, @@ -1463,6 +1502,8 @@ async def hydrate_config_server_dcr_client(mcp_server: MCPServer) -> bool: async def _resolve_persisted_dcr_client( mcp_server: MCPServer, + *, + raise_on_error: bool = False, ) -> tuple[Optional["LiteLLM_MCPServerTable"], _PersistedDcrCredentials | None]: """Resolve a server's persisted DCR client using the same two-level rule the write path uses, so read and write always agree. First, whether the server HAS a LiteLLM_MCPServerTable row: a row is @@ -1483,6 +1524,8 @@ async def _resolve_persisted_dcr_client( prisma_client = get_prisma_client_or_throw("Database not connected. Cannot read MCP OAuth client registration.") row: Final = await get_mcp_server(prisma_client=prisma_client, server_id=mcp_server.server_id) except Exception as exc: # noqa: BLE001 # best-effort read; DB may be unreachable + if raise_on_error: + raise verbose_logger.debug( "register_client_with_server: failed to read persisted DCR client for server_id=%s: %s", mcp_server.server_id, @@ -1491,12 +1534,14 @@ async def _resolve_persisted_dcr_client( return None, None if row is not None: + if row.url != mcp_server.url or (row.issuer and mcp_server.issuer and row.issuer != mcp_server.issuer): + return row, None credentials: Final = _get_persisted_dcr_credentials(row.credentials) if credentials is not None and credentials.client_id: return row, credentials return row, None if global_mcp_server_manager.is_config_declared_server(mcp_server.server_id): - return None, await _load_store_dcr_credentials(mcp_server) + return None, await _load_store_dcr_credentials(mcp_server, raise_on_error=raise_on_error) return None, None @@ -1519,20 +1564,24 @@ async def _reuse_persisted_dcr_client_if_available( if not _apply_persisted_dcr_credentials(mcp_server, credentials): return False - if persisted_mcp_server is not None: + await _refresh_persisted_dcr_server(persisted_mcp_server) + return bool(mcp_server.client_id) + + +async def _refresh_persisted_dcr_server(persisted_server: Optional["LiteLLM_MCPServerTable"]) -> None: + if persisted_server is not None and persisted_server.approval_status != "draft": from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids circular import global_mcp_server_manager, ) try: - await global_mcp_server_manager.update_server(persisted_mcp_server) + await global_mcp_server_manager.update_server(persisted_server) except Exception as exc: # noqa: BLE001 # best-effort registry refresh verbose_logger.warning( "register_client_with_server: failed to refresh persisted DCR client registration for server_id=%s: %s", - mcp_server.server_id, + persisted_server.server_id, exc, ) - return bool(mcp_server.client_id) async def _persisted_dcr_redirect_uri_is_stale(mcp_server: MCPServer, current_redirect_uri: str) -> bool: @@ -1575,7 +1624,7 @@ async def _persist_dcr_client_registration( ``refresh_token`` grant has no client identity, so an expired access token forces a full re-authorization instead of a silent refresh. Mirrors the ``encrypt_credentials`` write that ``client_credentials`` and token exchange already use. Failures are logged, - never raised: registration still returns to the caller even when persistence fails. + never raised here: the registration endpoint rejects a failed persistence result. The client-forwarded token modes (``true_passthrough`` / ``oauth_delegate``) are skipped unconditionally: the caller holds the upstream token and the gateway must hold no OAuth @@ -1604,9 +1653,6 @@ async def _persist_dcr_client_registration( ) return "failed" - if await _reuse_persisted_dcr_client_if_available(mcp_server, current_redirect_uri=current_redirect_uri): - return "reused" - token_endpoint_auth_method: Final = ( "client_secret_basic" if registration.token_endpoint_auth_method == "client_secret_basic" else None ) @@ -1615,6 +1661,8 @@ async def _persist_dcr_client_registration( "client_secret": registration.client_secret, "token_endpoint_auth_method": token_endpoint_auth_method, "redirect_uris": [current_redirect_uri], + "dcr_issuer": mcp_server.issuer, + "dcr_server_url": mcp_server.url, } from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # avoids circular import @@ -1632,6 +1680,27 @@ async def _persist_dcr_client_registration( prisma_client: Final = get_prisma_client_or_throw( "Database not connected. Cannot persist MCP OAuth client registration." ) + except HTTPException: + # This getter only raises when no database is configured. The existing single-process + # temporary-session mode keeps registrations in memory; database write errors below fail. + _apply_persisted_dcr_credentials(mcp_server, _PersistedDcrCredentials.model_validate(credentials)) + return "persisted" + + try: + stored, latest_credentials = await _resolve_persisted_dcr_client(mcp_server, raise_on_error=True) + if stored is not None and ( + stored.url != mcp_server.url + or stored.auth_type != mcp_server.auth_type + or (stored.issuer and mcp_server.issuer and stored.issuer != mcp_server.issuer) + ): + return "failed" + if ( + latest_credentials is not None + and not _redirect_uri_not_registered(latest_credentials, current_redirect_uri) + and _apply_persisted_dcr_credentials(mcp_server, latest_credentials) + ): + await _refresh_persisted_dcr_server(stored) + return "reused" updated_row: Final = await update_mcp_server( prisma_client=prisma_client, data=( @@ -1649,19 +1718,25 @@ async def _persist_dcr_client_registration( ) ), touched_by="mcp_oauth_dcr", + expected_updated_at=stored.updated_at if stored else None, ) if updated_row is not None and not isinstance(updated_row, McpIdentifierConflict): - await global_mcp_server_manager.update_server(updated_row) + await _refresh_persisted_dcr_server(updated_row) + _apply_persisted_dcr_credentials(mcp_server, _PersistedDcrCredentials.model_validate(credentials)) return "persisted" + if stored is not None: + return ( + "reused" + if await _reuse_persisted_dcr_client_if_available(mcp_server, current_redirect_uri) + else "failed" + ) if global_mcp_server_manager.is_config_declared_server(mcp_server.server_id): await upsert_mcp_server_oauth_client_credentials( prisma_client=prisma_client, server_id=mcp_server.server_id, credentials=credentials, ) - mcp_server.client_id = registration.client_id - mcp_server.client_secret = registration.client_secret - mcp_server.token_endpoint_auth_method = token_endpoint_auth_method + _apply_persisted_dcr_credentials(mcp_server, _PersistedDcrCredentials.model_validate(credentials)) return "persisted" except Exception as exc: # noqa: BLE001 verbose_logger.warning( @@ -1859,20 +1934,27 @@ async def register_client_with_server( "redirect_uris": client_facing_redirect_uris, } - if mcp_server.client_id and not ( - persist_credentials - and mcp_server.registration_url - and await _persisted_dcr_redirect_uri_is_stale(mcp_server, current_redirect_uri) + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint) + client_matches: Final = oauth_client_registration_matches( + resolved_server.dcr_issuer, resolved_server.dcr_server_url, resolved_server.issuer, resolved_server.url + ) + if ( + resolved_server.client_id + and client_matches + and not ( + persist_credentials + and resolved_server.registration_url + and await _persisted_dcr_redirect_uri_is_stale(resolved_server, current_redirect_uri) + ) ): return dummy_return if await _reuse_persisted_dcr_client_if_available( - mcp_server, + resolved_server, current_redirect_uri=current_redirect_uri if persist_credentials else None, ): return dummy_return - resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint) if resolved_server.effective_authorization_url is None: raise HTTPException( status_code=400, @@ -1908,19 +1990,32 @@ async def register_client_with_server( server_id=resolved_server.server_id, ) - token_response = response.json() - - if persist_credentials and not bridge_relay: - persistence_result = await _persist_dcr_client_registration( - resolved_server, token_response, current_redirect_uri - ) - if persistence_result == "reused": - return dummy_return - - if client_redirect_uris and not bridge_relay and isinstance(token_response, dict): - token_response = {**token_response, "redirect_uris": client_facing_redirect_uris} - - return JSONResponse(token_response) + token_response: Final = response.json() + persistence_result: Final = ( + await _persist_dcr_client_registration(resolved_server, token_response, current_redirect_uri) + if persist_credentials and not bridge_relay + else None + ) + if persistence_result == "reused": + return dummy_return + if persistence_result == "failed": + raise HTTPException(status_code=503, detail="OAuth client registration could not be saved; retry authorization") + bound_response: Final = ( + { + **token_response, + "dcr_issuer": resolved_server.issuer, + "dcr_server_url": resolved_server.url, + "dcr_redirect_uris": [current_redirect_uri], + } + if persistence_result == "persisted" + else token_response + ) + client_response: Final = ( + {**bound_response, "redirect_uris": client_facing_redirect_uris} + if client_redirect_uris and not bridge_relay and isinstance(bound_response, dict) + else bound_response + ) + return JSONResponse(client_response) @router.get("/authorize/mcp-session") @@ -2229,11 +2324,19 @@ def _render_oauth_error_html(error: str, description: str | None) -> HTMLRespons return HTMLResponse(body, status_code=400) +def _authorization_response_issuer_is_trusted(response_issuer: str | None, state_data: Mapping[str, object]) -> bool: + expected_issuer: Final = state_data.get("expected_issuer") + if response_issuer is None: + return state_data.get("authorization_response_iss_parameter_supported") is not True + return isinstance(expected_issuer, str) and bool(expected_issuer) and response_issuer == expected_issuer + + @router.get("/callback") async def callback( request: Request, code: str | None = None, state: str | None = None, + iss: str | None = None, error: str | None = None, error_description: str | None = None, error_uri: str | None = None, @@ -2244,7 +2347,9 @@ async def callback( - A successful authorization response (``code`` + ``state``), which is forwarded back to the validated client ``redirect_uri`` with the - original (un-wrapped) ``state``. + original (un-wrapped) ``state``, once the RFC 9207 ``iss`` (when the + authorization server sent one) matches the issuer /authorize sealed + into the state. - An error response (``error``[+``error_description``/``error_uri``]), per RFC 6749 §4.1.2.1. When ``state`` is present and decodes to a trusted ``redirect_uri``, the error params are propagated back to the client so @@ -2262,6 +2367,13 @@ async def callback( encoded_state = _resolve_encoded_oauth_state(request, state) try: state_data = decode_state_hash(encoded_state) + error_issuer_state: Final = TypeAdapter(dict[str, object]).validate_python(state_data) + if not _authorization_response_issuer_is_trusted(iss, error_issuer_state): + rejected_error: Final = _render_oauth_error_html( + "invalid_issuer", "Unexpected authorization issuer" + ) + _clear_oauth_state_cookie(rejected_error, request, state) + return rejected_error original_state = state_data.get("original_state") redirect_uri = _get_validated_client_redirect_uri(request, state_data) except Exception: @@ -2311,6 +2423,22 @@ async def callback( # states while permitting same-origin / allowlisted clients. redirect_uri = _get_validated_client_redirect_uri(request, state_data) + issuer_state: Final = TypeAdapter(dict[str, object]).validate_python(state_data) + if not _authorization_response_issuer_is_trusted(iss, issuer_state): + verbose_logger.warning( + "MCP /callback rejected an authorization response: RFC 9207 iss=%r does not match the " + "issuer this flow was sent to (%r)", + iss, + issuer_state.get("expected_issuer"), + ) + issuer_error_response: Final = _render_oauth_error_html( + "invalid_issuer", + "This authorization response came from a different identity provider than the one this " + "MCP server is configured to use.", + ) + _clear_oauth_state_cookie(issuer_error_response, request, state) + return issuer_error_response + # Interactive dcr_bridge oauth_delegate: the state carries the litellm user the authorize step # captured. Instead of forwarding the raw upstream code (which the client would present at the # token endpoint with no way to prove who signed in), seal the user and the upstream code into a diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b8099c97861..abfb9dc56fe 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -719,13 +719,12 @@ def _normalized_authorize_endpoint(url: str) -> str: def _issuer_matches(claimed_issuer: object, configured_issuer: str) -> bool: """RFC 8414 §3.3 issuer equality between the metadata document's self-attested ``issuer`` and the - admin-configured issuer, tolerant only of URL-insignificant differences (scheme/host case, the - default port, a trailing slash). A non-string or empty claimed issuer never matches, so a + admin-configured issuer. A non-string or empty claimed issuer never matches, so a document that omits ``issuer`` fails closed under issuer-anchored discovery. """ if not isinstance(claimed_issuer, str) or not claimed_issuer: return False - return _normalized_authorize_endpoint(claimed_issuer) == _normalized_authorize_endpoint(configured_issuer) + return claimed_issuer == configured_issuer def _flow_endpoints_missing( @@ -856,6 +855,11 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv may_carry: Final = _endpoints_corroborate_authorization_url( previous_server.authorization_url, new_server.authorization_url ) + if may_carry and new_server.issuer is None: + new_server.issuer = previous_server.issuer + new_server.authorization_response_iss_parameter_supported = ( # rebind-ok: publish on the existing rebuild object + previous_server.authorization_response_iss_parameter_supported + ) if new_server.authorization_url is None and previous_server.authorization_url: new_server.authorization_url = previous_server.authorization_url if may_carry and new_server.token_url is None and previous_server.token_url: @@ -2110,13 +2114,20 @@ class MCPServerManager: if metadata is None: return server discovered_issuer: Final = metadata.discovered_issuer if not metadata.from_origin_fallback else None - resolved: Final = server.model_copy() - resolved.scopes = server.scopes or metadata.scopes - resolved.issuer = server.issuer or discovered_issuer - resolved.authorization_url = server.authorization_url or metadata.authorization_url - resolved.token_url = server.token_url or metadata.token_url - resolved.registration_url = server.registration_url or metadata.registration_url - return resolved + return server.model_copy( + update={ + "scopes": server.scopes or metadata.scopes, + "issuer": server.issuer or discovered_issuer, + "authorization_response_iss_parameter_supported": ( + metadata.authorization_response_iss_parameter_supported + if discovered_issuer is not None + else server.authorization_response_iss_parameter_supported + ), + "authorization_url": server.authorization_url or metadata.authorization_url, + "token_url": server.token_url or metadata.token_url, + "registration_url": server.registration_url or metadata.registration_url, + } + ) def _oauth_discovery_slot_is_current(self, server_id: str, generation: int) -> bool: slot: Final = self._oauth_discovery_slot(server_id) @@ -2641,6 +2652,11 @@ class MCPServerManager: scopes=resolved_scopes, configured_scopes=tuple(configured_scopes) if configured_scopes else None, issuer=effective_issuer, + authorization_response_iss_parameter_supported=( + gated_oauth_metadata.authorization_response_iss_parameter_supported + if gated_oauth_metadata + else False + ), issuer_is_anchored=use_issuer_anchor, authorization_url=resolved_authorization_url, token_url=resolved_token_url, @@ -3202,12 +3218,17 @@ class MCPServerManager: extra_headers=getattr(mcp_server, "extra_headers", None), static_headers=static_headers_dict, env_vars=env_vars_list, + dcr_issuer=credentials_dict.get("dcr_issuer") if credentials_dict else None, + dcr_server_url=credentials_dict.get("dcr_server_url") if credentials_dict else None, client_id=client_id_value or getattr(mcp_server, "client_id", None), client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)), scopes=resolved_scopes, configured_scopes=configured_scopes, issuer=effective_issuer, + authorization_response_iss_parameter_supported=( + gated_oauth_metadata.authorization_response_iss_parameter_supported if gated_oauth_metadata else False + ), issuer_is_anchored=use_issuer_anchor, authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), token_url=manual_token_url or getattr(gated_oauth_metadata, "token_url", None), @@ -5174,6 +5195,10 @@ class MCPServerManager: token_url=data.get("token_endpoint"), registration_url=data.get("registration_endpoint"), discovered_issuer=claimed_issuer if isinstance(claimed_issuer, str) and claimed_issuer else None, + authorization_response_iss_parameter_supported=data.get( + "authorization_response_iss_parameter_supported" + ) + is True, ) if any( diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 39865a35ec6..3b5eda6cff2 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -658,6 +658,17 @@ def canonicalize_url_identity(url: str) -> str: return urlunparse((scheme, netloc, parsed.path.rstrip("/"), "", "", "")) +def oauth_client_registration_matches( + registered_issuer: str | None, + registered_url: str | None, + current_issuer: str | None, + current_url: str | None, +) -> bool: + if registered_issuer and current_issuer: + return registered_issuer == current_issuer + return not registered_url or registered_url == current_url + + def canonical_resource_uri(url: str) -> str | None: """Canonicalize an upstream MCP server URL into an RFC 8707 resource identifier. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 8530a4d31b4..03408043f8e 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -31307,6 +31307,28 @@ ], "title": "Client Secret" }, + "dcr_issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Dcr Issuer" + }, + "dcr_server_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Dcr Server Url" + }, "id_jag_resource": { "anyOf": [ { @@ -34023,6 +34045,28 @@ ], "title": "Client Secret" }, + "dcr_issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Dcr Issuer" + }, + "dcr_server_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Dcr Server Url" + }, "id_jag_resource": { "anyOf": [ { @@ -35696,7 +35740,7 @@ }, "/callback": { "get": { - "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.", + "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``, once the RFC 9207 ``iss`` (when the\n authorization server sent one) matches the issuer /authorize sealed\n into the state.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.", "operationId": "callback_callback_get", "parameters": [ { @@ -35731,6 +35775,22 @@ "title": "State" } }, + { + "in": "query", + "name": "iss", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Iss" + } + }, { "in": "query", "name": "error", @@ -37486,6 +37546,28 @@ ], "title": "Client Secret" }, + "dcr_issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Dcr Issuer" + }, + "dcr_server_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Dcr Server Url" + }, "id_jag_resource": { "anyOf": [ { @@ -41480,6 +41562,28 @@ ], "title": "Client Secret" }, + "dcr_issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Dcr Issuer" + }, + "dcr_server_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Dcr Server Url" + }, "id_jag_resource": { "anyOf": [ { diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e1504172d9e..2aed0e2afb8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -157,13 +157,16 @@ if MCP_AVAILABLE: get_user_env_vars, get_user_env_vars_bulk, get_user_oauth_credential, + is_resubmitted_oauth_client, list_server_user_credentials, list_user_oauth_credentials, mcp_oauth_token_identity, merge_user_env_vars, + oauth_credentials_for_upstream_edit, purge_user_oauth_credentials_for_server, reject_mcp_server, set_mcp_server_pinned_tools, + stale_mcp_auth_fields, store_user_credential, store_user_oauth_credential, update_mcp_server, @@ -870,6 +873,9 @@ if MCP_AVAILABLE: ("authentication_token", "auth_value"), ("client_id", "client_id"), ("client_secret", "client_secret"), + ("token_endpoint_auth_method", "token_endpoint_auth_method"), + ("dcr_issuer", "dcr_issuer"), + ("dcr_server_url", "dcr_server_url"), ("scopes", "scopes"), ("aws_access_key_id", "aws_access_key_id"), ("aws_secret_access_key", "aws_secret_access_key"), @@ -890,30 +896,73 @@ if MCP_AVAILABLE: value for key, value in as_dict.items() if key not in MCP_ADMIN_CONFIG_CREDENTIAL_KEYS and key != "scopes" ) + def _oauth_session_changes_upstream( + payload: NewMCPServerRequest, existing: MCPServer | LiteLLM_MCPServerTable + ) -> bool: + return existing.auth_type == MCPAuth.oauth2 and ( + payload.url != existing.url + or ("issuer" in payload.model_fields_set and (payload.issuer or None) != existing.issuer) + or (payload.auth_type is not None and payload.auth_type != existing.auth_type) + ) + def _inherit_credentials_from_existing_server( payload: NewMCPServerRequest, ) -> NewMCPServerRequest: - if not payload.server_id or _has_non_admin_config_credentials(payload.credentials): + if not payload.server_id: return payload existing_server: Final = global_mcp_server_manager.get_mcp_server_by_id(payload.server_id) if existing_server is None: return payload - - inherited_credentials: dict[str, object] = { + upstream_changed: Final = _oauth_session_changes_upstream(payload, existing_server) + issuer_changed: Final = ( + "issuer" in payload.model_fields_set and (payload.issuer or None) != existing_server.issuer + ) + cleared: Final = ( + stale_mcp_auth_fields( + payload.model_dump(exclude_unset=True), + lambda field: ( + getattr(existing_server, f"configured_{field}", None) or getattr(existing_server, field, None) + ), + ) + if upstream_changed + else {} + ) + resolved_payload: Final = ( + payload.model_copy(update={**cleared, "oauth2_flow": payload.oauth2_flow}) if upstream_changed else payload + ) + supplied: Final = dict(resolved_payload.credentials or {}) + existing_credentials: Final[dict[str, object]] = { credential_key: value for server_attr, credential_key in _INHERITED_CREDENTIAL_FIELDS if (value := getattr(existing_server, server_attr, None)) } + resubmitted_client: Final = is_resubmitted_oauth_client(supplied, existing_credentials) + if _has_non_admin_config_credentials(resolved_payload.credentials) and not ( + upstream_changed and resubmitted_client + ): + return resolved_payload + + bound_credentials: Final = ( + oauth_credentials_for_upstream_edit( + {**existing_credentials, **supplied}, + existing_server.issuer, + existing_server.url, + issuer_changed=issuer_changed + or (resolved_payload.auth_type or existing_server.auth_type) != existing_server.auth_type, + ) + if upstream_changed + else existing_credentials + ) # The gate above guarantees anything still supplied is admin config, which the admin just # typed, so it wins over the stored value. - inherited_credentials = {**inherited_credentials, **dict(payload.credentials or {})} + inherited_credentials: Final = bound_credentials if upstream_changed else {**bound_credentials, **supplied} if not inherited_credentials: - return payload + return resolved_payload.model_copy(update={"credentials": {}}) if upstream_changed else resolved_payload try: - return payload.model_copy(update={"credentials": inherited_credentials}) + return resolved_payload.model_copy(update={"credentials": inherited_credentials}) except AttributeError: pass @@ -937,15 +986,20 @@ if MCP_AVAILABLE: supplied: Final = payload.server_id if not supplied: return str(uuid.uuid4()) - if global_mcp_server_manager.get_mcp_server_by_id(supplied) is not None: - return supplied + registered: Final = global_mcp_server_manager.get_mcp_server_by_id(supplied) + if registered is not None: + return str(uuid.uuid4()) if _oauth_session_changes_upstream(payload, registered) else supplied prisma_client: Final = _get_prisma_client_or_none() if prisma_client is None: return supplied # A draft is another session's row, not a saved server, so re-supplying an id this # endpoint previously handed back must not let a later session adopt its configuration. existing: Final = await get_mcp_server(prisma_client, supplied) - if existing is None or existing.approval_status == MCPApprovalStatus.draft: + if ( + existing is None + or existing.approval_status == MCPApprovalStatus.draft + or _oauth_session_changes_upstream(payload, existing) + ): return str(uuid.uuid4()) return supplied diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index ce86e79a6f0..949506898b1 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -9,7 +9,7 @@ from urllib.parse import urlsplit import httpx from pydantic import BaseModel, ConfigDict, Field -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.base import HiddenParams @@ -161,6 +161,9 @@ MCPTokenEndpointAuthMethod = Literal["client_secret_basic", "client_secret_post" class MCPCredentials(TypedDict, total=False): + dcr_issuer: ReadOnly[str | None] + dcr_server_url: ReadOnly[str | None] + auth_value: str | None """ Authentication value diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 2ee19b3e59a..343fb731b06 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -38,6 +38,7 @@ class MCPOAuthMetadata(BaseModel): authorization_url: str | None = None token_url: str | None = None registration_url: str | None = None + authorization_response_iss_parameter_supported: bool = False discovered_issuer: str | None = None """The ``issuer`` the authorization-server metadata document self-attests (RFC 8414). Persisted trust-on-first-use as the server's ``issuer`` when none is configured, so that later rebuilds @@ -118,6 +119,9 @@ class MCPServer(BaseModel): client_secret: str | None = None issuer: str | None = None issuer_is_anchored: bool = False + authorization_response_iss_parameter_supported: bool = False + dcr_issuer: str | None = None + dcr_server_url: str | None = None scopes: list[str] | None = None authorization_url: str | None = None token_url: str | None = None diff --git a/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py b/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py index c20b73c0d63..81470c21d51 100644 --- a/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py +++ b/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py @@ -38,6 +38,7 @@ pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_ class OAuthMetadata(BaseModel): + issuer: str authorization_endpoint: str token_endpoint: str registration_endpoint: str @@ -135,6 +136,7 @@ class TestMcpOauthHappyPath: auth_type="oauth2", oauth2_flow="authorization_code", per_server_oauth_discovery=route == "explicit_header_jwt", + issuer=metadata.issuer if metadata else None, authorization_url=metadata.authorization_endpoint if metadata else None, token_url=metadata.token_endpoint if metadata else None, registration_url=metadata.registration_endpoint if metadata else None, diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 90fe04ae896..ccb5cd35280 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -658,6 +658,7 @@ class McpServerCreateBody(BaseModel): auth_type: str | None = None oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None per_server_oauth_discovery: bool | None = None + issuer: str | None = None authorization_url: str | None = None token_url: str | None = None registration_url: str | None = None diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 4e27ec134d4..6d71256ff18 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from litellm.proxy._types import LiteLLM_MCPServerTable from litellm.types.mcp import MCPAuth if TYPE_CHECKING: @@ -1365,6 +1366,8 @@ async def test_register_client_persists_dcr_client_identity(): mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) mock_update = AsyncMock(return_value=MagicMock()) mock_update_server = AsyncMock() @@ -1373,7 +1376,7 @@ async def test_register_client_persists_dcr_client_identity(): "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", return_value=mock_async_client, ), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma), patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), ): @@ -1390,7 +1393,12 @@ async def test_register_client_persists_dcr_client_identity(): import json assert response.status_code == 200 - assert json.loads(response.body.decode("utf-8")) == mock_response.json.return_value + assert json.loads(response.body.decode("utf-8")) == { + **mock_response.json.return_value, + "dcr_issuer": None, + "dcr_server_url": None, + "dcr_redirect_uris": ["https://proxy.litellm.example/callback"], + } mock_update.assert_called_once() update_data = mock_update.call_args.kwargs["data"] @@ -1452,6 +1460,8 @@ async def _register_persistence_attempted_for_auth_type(auth_type: MCPAuth) -> b mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) mock_update = AsyncMock(return_value=MagicMock()) with ( @@ -1459,7 +1469,7 @@ async def _register_persistence_attempted_for_auth_type(auth_type: MCPAuth) -> b "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", return_value=mock_async_client, ), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma), patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), ): @@ -1473,7 +1483,12 @@ async def _register_persistence_attempted_for_auth_type(auth_type: MCPAuth) -> b persist_credentials=True, ) - assert json.loads(response.body.decode("utf-8")) == mock_response.json.return_value + expected_binding = { + "dcr_issuer": None, + "dcr_server_url": None, + "dcr_redirect_uris": ["https://proxy.litellm.example/callback"], + } if auth_type == MCPAuth.oauth2 else {} + assert json.loads(response.body.decode("utf-8")) == {**mock_response.json.return_value, **expected_binding} return mock_update.await_count > 0 @@ -1529,7 +1544,7 @@ async def test_register_client_persists_only_to_its_own_row_when_another_server_ ) sibling_row_with_client = MagicMock(server_id="server-a", url=shared_url) sibling_row_with_client.credentials = {"client_id": "client-a-do-not-adopt"} - own_row_without_client = MagicMock(server_id="server-b", url=shared_url) + own_row_without_client = LiteLLM_MCPServerTable.model_validate(fresh_server.model_dump(exclude_none=True)) own_row_without_client.credentials = {} rows_by_server_id = {"server-a": sibling_row_with_client, "server-b": own_row_without_client} @@ -1622,6 +1637,8 @@ async def test_register_client_does_not_clobber_token_url_when_absent(): mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) mock_update = AsyncMock(return_value=MagicMock()) mock_update_server = AsyncMock() @@ -1630,7 +1647,7 @@ async def test_register_client_does_not_clobber_token_url_when_absent(): "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", return_value=mock_async_client, ), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma), patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), ): @@ -1685,7 +1702,7 @@ async def test_register_client_reuses_persisted_client_id_for_non_admin_when_reg mock_request.base_url = "https://proxy.litellm.example/" mock_request.headers = {} - persisted_server = MagicMock() + persisted_server = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) persisted_server.credentials = {"client_id": "persisted-client"} mock_get_mcp_server = AsyncMock(return_value=persisted_server) mock_update_mcp_server = AsyncMock() @@ -1761,7 +1778,7 @@ async def test_register_client_reuse_refreshes_request_server_when_manager_updat mock_request.base_url = "https://proxy.litellm.example/" mock_request.headers = {} - persisted_server = MagicMock() + persisted_server = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) persisted_server.credentials = { "client_id": "persisted-client", "client_secret": "persisted-secret", @@ -1803,7 +1820,8 @@ async def test_register_client_reuse_refreshes_request_server_when_manager_updat @pytest.mark.asyncio -async def test_register_client_returns_reused_client_when_concurrent_persist_wins(): +@pytest.mark.parametrize("late_winner", [False, True]) +async def test_register_client_returns_reused_client_when_concurrent_persist_wins(late_winner): try: from fastapi import Request @@ -1843,10 +1861,11 @@ async def test_register_client_returns_reused_client_when_concurrent_persist_win mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) - persisted_server = MagicMock() + persisted_server = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) persisted_server.credentials = {"client_id": "persisted-client"} - mock_get_mcp_server = AsyncMock(side_effect=[None, persisted_server]) - mock_update_mcp_server = AsyncMock() + empty_server = persisted_server.model_copy(update={"credentials": None}) + mock_get_mcp_server = AsyncMock(side_effect=[empty_server, empty_server, persisted_server] if late_winner else [empty_server, persisted_server]) + mock_update_mcp_server = AsyncMock(return_value=None) mock_update_server = AsyncMock() with ( @@ -1878,7 +1897,10 @@ async def test_register_client_returns_reused_client_when_concurrent_persist_win assert response["client_id"] == "remote_server" assert oauth2_server.client_id == "persisted-client" mock_async_client.post.assert_called_once() - mock_update_mcp_server.assert_not_called() + if late_winner: + mock_update_mcp_server.assert_awaited_once() + else: + mock_update_mcp_server.assert_not_called() mock_update_server.assert_called_once_with(persisted_server) @@ -1937,7 +1959,7 @@ async def test_register_client_re_registers_when_persisted_redirect_uri_no_longe mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) - persisted_server = MagicMock() + persisted_server = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) persisted_server.credentials = { "client_id": "stale-client", "client_secret": "stale-secret", @@ -2014,7 +2036,7 @@ async def test_register_client_grandfathers_persisted_client_without_recorded_re mock_async_client = MagicMock() mock_async_client.post = AsyncMock() - persisted_server = MagicMock() + persisted_server = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) persisted_server.credentials = {"client_id": "legacy-client"} mock_get_mcp_server = AsyncMock(return_value=persisted_server) mock_update_mcp_server = AsyncMock() @@ -2072,7 +2094,7 @@ async def test_register_client_keeps_persisted_client_when_recorded_redirect_uri mock_async_client = MagicMock() mock_async_client.post = AsyncMock() - persisted_server = MagicMock() + persisted_server = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) persisted_server.credentials = { "client_id": "kept-client", "redirect_uris": ["https://proxy.litellm.example/callback"], @@ -2137,7 +2159,7 @@ async def test_register_client_non_admin_reuses_persisted_client_despite_redirec mock_async_client = MagicMock() mock_async_client.post = AsyncMock() - persisted_server = MagicMock() + persisted_server = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) persisted_server.credentials = { "client_id": "persisted-client", "redirect_uris": ["https://old-origin.example/callback"], @@ -4465,6 +4487,173 @@ async def test_callback_error_path_reads_cookie_and_clears_it(monkeypatch): assert cleared[cookie_name]["max-age"] == "0" +def _issuer_anchored_oauth_server(issuer: str | None = "https://idp.example.com"): + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id="rfc9207_server", + name="rfc9207", + server_name="rfc9207", + alias="rfc9207", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="upstream-client-id", + issuer=issuer, + authorization_url="https://idp.example.com/oauth/authorize", + token_url="https://idp.example.com/oauth/token", + ) + + +async def _authorize_then_callback(server, iss, monkeypatch, expected_issuer_override=..., error=None): + """Run /authorize for ``server``, then feed the resulting flow back through /callback with the + RFC 9207 ``iss`` the authorization server supposedly returned. Returns (callback_response, + sealed_state_data).""" + from http.cookies import SimpleCookie + from urllib.parse import parse_qs, urlparse + + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _oauth_state_cookie_name, + authorize_with_server, + callback, + decode_state_hash, + encode_state_with_base_url, + ) + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-LIT-5940") + client_redirect_uri = "http://127.0.0.1:6274/oauth/callback/debug" + + authorize_request = MagicMock(spec=Request) + authorize_request.base_url = "https://proxy.example.com/" + authorize_request.headers = {} + + authorize_response = await authorize_with_server( + request=authorize_request, + mcp_server=server, + client_id="upstream-client-id", + redirect_uri=client_redirect_uri, + state="client-original-state-9207", + code_challenge="challenge", + code_challenge_method="S256", + ) + handle = parse_qs(urlparse(authorize_response.headers["location"]).query)["state"][0] + jar = SimpleCookie() + jar.load(authorize_response.headers["set-cookie"]) + cookie_name = _oauth_state_cookie_name(handle) + sealed_state = jar[cookie_name].value + if expected_issuer_override is not ...: + # A state minted before the issuer was sealed into it: same shape, key absent. + sealed_state = encode_state_with_base_url( + base_url=client_redirect_uri, + original_state="client-original-state-9207", + client_redirect_uri=client_redirect_uri, + expected_issuer=expected_issuer_override, + ) + + callback_request = MagicMock(spec=Request) + callback_request.base_url = "https://proxy.example.com/" + callback_request.headers = {} + callback_request.cookies = {cookie_name: sealed_state} + + response = await callback( + request=callback_request, + code="upstream-auth-code" if error is None else None, + error=error, + state=handle, + iss=iss, + ) + return response, decode_state_hash(sealed_state) + + +@pytest.mark.asyncio +async def test_authorize_seals_the_issuer_and_callback_accepts_a_matching_rfc9207_iss(monkeypatch): + """The callback accepts the exact issuer sealed during authorization.""" + from urllib.parse import parse_qs, urlparse + + response, state_data = await _authorize_then_callback( + _issuer_anchored_oauth_server(), + iss="https://idp.example.com", + monkeypatch=monkeypatch, + ) + + assert state_data["expected_issuer"] == "https://idp.example.com" + assert response.status_code == 302 + query = parse_qs(urlparse(response.headers["location"]).query) + assert query["code"] == ["upstream-auth-code"] + assert query["state"] == ["client-original-state-9207"] + + +@pytest.mark.asyncio +async def test_callback_rejects_authorization_response_from_a_different_issuer(monkeypatch): + """LIT-5940 / RFC 9207 §2.4: an ``iss`` naming an authorization server we never sent the user to + is a mix-up attack, so the code must not reach the client's redirect_uri.""" + response, _ = await _authorize_then_callback( + _issuer_anchored_oauth_server(), + iss="https://attacker-idp.example.com", + monkeypatch=monkeypatch, + ) + + assert response.status_code == 400 + assert "location" not in response.headers + assert b"upstream-auth-code" not in response.body + assert b"invalid_issuer" in response.body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [None, "access_denied"]) +async def test_callback_rejects_supplied_issuer_without_expected_identity(monkeypatch, error): + response, state_data = await _authorize_then_callback( + _issuer_anchored_oauth_server(issuer=None), + iss="https://unknown-idp.example.com", + monkeypatch=monkeypatch, + error=error, + ) + assert state_data["expected_issuer"] is None + assert response.status_code == 400 + assert "location" not in response.headers + assert b"upstream-auth-code" not in response.body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("issuer", [None, "https://idp.example.com"]) +async def test_callback_accepts_missing_unadvertised_issuer(monkeypatch, issuer): + response, _ = await _authorize_then_callback( + _issuer_anchored_oauth_server(issuer), iss=None, monkeypatch=monkeypatch, + ) + assert response.status_code == 302 + + +@pytest.mark.asyncio +async def test_callback_rejects_an_issuer_differing_only_outside_the_path(monkeypatch): + """A tenant a deployment encoded in a query string is part of that issuer's identity, so the + comparison must not canonicalize it away and let another tenant's response through.""" + response, state_data = await _authorize_then_callback( + _issuer_anchored_oauth_server(issuer="https://idp.example.com/?tenant=a"), + iss="https://idp.example.com/?tenant=b", + monkeypatch=monkeypatch, + ) + + assert state_data["expected_issuer"] == "https://idp.example.com/?tenant=a" + assert response.status_code == 400 + assert "location" not in response.headers + assert b"upstream-auth-code" not in response.body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("iss,expected_status", [(None, 302), ("https://some-idp.example.com", 400)]) +async def test_callback_legacy_state_requires_absent_issuer(monkeypatch, iss, expected_status): + response, _ = await _authorize_then_callback( + _issuer_anchored_oauth_server(), + iss=iss, + monkeypatch=monkeypatch, + expected_issuer_override=None, + ) + assert response.status_code == expected_status + + @pytest.mark.asyncio async def test_oauth_authorize_includes_scopes_from_server_config(): """Test that authorize endpoint includes scopes from server configuration.""" @@ -9070,7 +9259,8 @@ async def test_persist_dcr_client_for_config_server_uses_side_store(): @pytest.mark.asyncio -async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch): +@pytest.mark.parametrize("legacy", [False, True]) +async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch, legacy): """On restart a config server's in-memory object has no client_id; hydration overlays the persisted DCR client from the server-scoped store, decrypting the encrypted-at-rest blob, so the refresh_token grant can authenticate as the registered client instead of re-authenticating.""" @@ -9090,12 +9280,14 @@ async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch): transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id=None, + url="https://resource.example/mcp", + issuer="https://idp.example", ) monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-hydrate-key") stored_blob = safe_dumps( encrypt_credentials( - credentials={ + credentials={**({} if legacy else {"dcr_issuer": "https://idp.example", "dcr_server_url": "https://resource.example/mcp"}), "client_id": "stored-client", "client_secret": "stored-secret", "token_endpoint_auth_method": "client_secret_basic", @@ -9145,12 +9337,14 @@ async def test_reuse_config_server_reads_store_with_real_crypto(monkeypatch): transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id=None, + url="https://resource.example/mcp", + issuer="https://idp.example", ) monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-reuse-key") blob = safe_dumps( encrypt_credentials( - credentials={"client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, + credentials={"dcr_issuer": "https://idp.example", "dcr_server_url": "https://resource.example/mcp","client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, encryption_key="salt-reuse-key", ) ) @@ -12759,3 +12953,232 @@ def test_invalidating_an_idle_server_leaves_no_generation_behind(): finally: for server_id in server_ids: discoverable_endpoints._OAUTH_METADATA_GENERATIONS.pop(server_id, None) + + +@pytest.mark.asyncio +async def test_callback_rejects_missing_advertised_issuer(monkeypatch): + server = _issuer_anchored_oauth_server().model_copy( + update={"authorization_response_iss_parameter_supported": True} + ) + response, _ = await _authorize_then_callback(server, iss=None, monkeypatch=monkeypatch) + assert response.status_code == 400 + assert "location" not in response.headers + assert b"upstream-auth-code" not in response.body + + +@pytest.mark.parametrize("issuer,url", [ + ("https://other.example", "https://resource.example/mcp"), + (None, "https://other.example/mcp"), +]) +def test_persisted_client_is_not_applied_to_another_upstream(issuer, url): + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _apply_persisted_dcr_credentials, _PersistedDcrCredentials, + ) + server = _issuer_anchored_oauth_server(issuer).model_copy(update={"url": url, "client_id": None}) + stored = _PersistedDcrCredentials.model_validate({ + "client_id": "old-client", "client_secret": "old-secret", + "dcr_issuer": "https://idp.example.com", "dcr_server_url": "https://resource.example/mcp", + }) + assert _apply_persisted_dcr_credentials(server, stored) is False + assert server.client_id is None + assert server.client_secret is None + + +@pytest.mark.asyncio +async def test_registration_does_not_write_client_after_upstream_edit(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "test-mcp-registration-salt") + from prisma import models + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _persist_dcr_client_registration + + server = _issuer_anchored_oauth_server().model_copy(update={"url": "https://resource.example/mcp", "client_id": None}) + current = models.LiteLLM_MCPServerTable.model_construct( + server_id=server.server_id, url="https://other.example/mcp", issuer="https://other.example", + auth_type="oauth2", transport="http", credentials=None, updated_at=datetime.now(timezone.utc), + ) + prisma = MagicMock() + table = AsyncMock() + table.find_unique.return_value = current + table.update.return_value = current + prisma.db.litellm_mcpservertable = table + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + result = await _persist_dcr_client_registration(server, {"client_id": "old-issuer-client"}, "https://proxy.example/callback") + assert result == "failed" + table.update.assert_not_called() + table.update_many.assert_not_called() + assert server.client_id is None + + +@pytest.mark.parametrize("response_issuer", ["https://IDP.example.com", "https://idp.example.com/", "https://idp.example.com:443"]) +def test_issuer_validation_uses_exact_identifier(response_issuer): + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _authorization_response_issuer_is_trusted + assert _authorization_response_issuer_is_trusted(response_issuer, {"expected_issuer": "https://idp.example.com"}) is False + + +@pytest.mark.parametrize("issuer,url", [ + ("https://idp.example.com", "https://resource.example/mcp"), + ("https://idp.example.com", "https://other-resource.example/mcp"), +]) +def test_persisted_client_remains_usable_for_its_issuer(issuer, url, monkeypatch): + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _apply_persisted_dcr_credentials, _PersistedDcrCredentials + monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding") + server = _issuer_anchored_oauth_server(issuer).model_copy(update={"url": url, "client_id": None}) + credentials = _PersistedDcrCredentials( + client_id="registered-client", client_secret="registered-secret", dcr_issuer=issuer, + dcr_server_url="https://resource.example/mcp", + ) + assert _apply_persisted_dcr_credentials(server, credentials) is True + assert server.client_id == "registered-client" + assert server.client_secret == "registered-secret" + assert server.dcr_issuer == issuer + + +@pytest.mark.asyncio +async def test_callback_does_not_forward_error_from_another_issuer(monkeypatch): + response, _ = await _authorize_then_callback( + _issuer_anchored_oauth_server(), iss="https://other.example", monkeypatch=monkeypatch, error="access_denied", + ) + assert response.status_code == 400 + assert "location" not in response.headers + assert b"invalid_issuer" in response.body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["authorize", "token"]) +async def test_bound_client_cannot_be_sent_to_a_different_issuer(monkeypatch, operation): + from fastapi import Request + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="bound-client", name="bound-client", transport="http", auth_type="oauth2", + url="https://new.example/mcp", issuer="https://new.example", client_id="old-client", + dcr_issuer="https://old.example", dcr_server_url="https://old.example/mcp", + ) + monkeypatch.setattr(endpoints, "_server_with_oauth_endpoints", AsyncMock(return_value=server)) + http_client = MagicMock() + monkeypatch.setattr(endpoints, "get_async_httpx_client", http_client) + request = ( + endpoints.authorize_with_server( + mcp_server=server, request=MagicMock(spec=Request), redirect_uri="http://localhost/callback", client_id="old-client", + ) + if operation == "authorize" + else endpoints.exchange_token_with_server( + mcp_server=server, request=MagicMock(spec=Request), grant_type="authorization_code", + code="test-code", redirect_uri="http://localhost/callback", client_id="old-client", + client_secret=None, code_verifier="verifier", + ) + ) + with pytest.raises(HTTPException) as error: + await request + assert error.value.status_code == 400 + assert "different issuer" in error.value.detail + http_client.assert_not_called() + + +@pytest.mark.asyncio +async def test_failed_registration_preserves_cached_credentials(monkeypatch): + from fastapi import Request + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + + server = _dcr_redirect_test_server("old-client").model_copy(update={ + "url": "https://new.example/mcp", "issuer": "https://new.example", + "dcr_issuer": "https://old.example", "client_secret": "old-secret", + }) + monkeypatch.setattr(endpoints, "_server_with_oauth_endpoints", AsyncMock(return_value=server)) + monkeypatch.setattr(endpoints, "_reuse_persisted_dcr_client_if_available", AsyncMock(return_value=False)) + response = MagicMock() + response.json.return_value = {"client_id": "new-client"} + register = AsyncMock(return_value=response) + monkeypatch.setattr(endpoints, "_post_dcr_registration", register) + persist = AsyncMock(return_value="failed") + monkeypatch.setattr(endpoints, "_persist_dcr_client_registration", persist) + request = MagicMock(spec=Request) + request.base_url = "https://gateway.example/" + request.headers = {} + with pytest.raises(HTTPException) as error: + await endpoints.register_client_with_server( + request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"], + response_types=["code"], token_endpoint_auth_method="none", persist_credentials=True, + ) + assert error.value.status_code == 503 + assert "could not be saved" in error.value.detail + assert server.client_id == "old-client" + assert server.client_secret == "old-secret" + register.assert_awaited_once() + persist.assert_awaited_once_with(server, {"client_id": "new-client"}, "https://gateway.example/callback") + + +@pytest.mark.asyncio +async def test_registration_without_database_keeps_client_in_temporary_server(monkeypatch): + from fastapi import Request + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + + monkeypatch.setattr(proxy_server, "prisma_client", None) + server = _dcr_redirect_test_server(None).model_copy(update={"issuer": "https://idp.example", "url": "https://resource.example/mcp"}) + monkeypatch.setattr(endpoints, "_server_with_oauth_endpoints", AsyncMock(return_value=server)) + upstream = MagicMock() + upstream.json.return_value = {"client_id": "temporary-client", "client_secret": "temporary-secret"} + monkeypatch.setattr(endpoints, "_post_dcr_registration", AsyncMock(return_value=upstream)) + request = MagicMock(spec=Request) + request.base_url = "https://gateway.example/" + request.headers = {} + response = await endpoints.register_client_with_server( + request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"], + response_types=["code"], token_endpoint_auth_method="none", persist_credentials=True, + ) + assert response.status_code == 200 + assert json.loads(response.body)["client_id"] == "temporary-client" + assert server.client_id == "temporary-client" + assert server.client_secret == "temporary-secret" + assert server.dcr_issuer == server.issuer + assert server.dcr_server_url == server.url + + +@pytest.mark.asyncio +@pytest.mark.parametrize("config_store", [False, True]) +async def test_registration_does_not_overwrite_credentials_after_failed_identity_read(monkeypatch, config_store): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints, db + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy import utils + + server = _dcr_redirect_test_server(None) + monkeypatch.setattr(utils, "get_prisma_client_or_throw", lambda _: MagicMock()) + monkeypatch.setattr(db, "get_mcp_server", AsyncMock(return_value=None, side_effect=None if config_store else RuntimeError("unavailable"))) + monkeypatch.setattr(db, "get_mcp_server_oauth_client_credentials", AsyncMock(side_effect=RuntimeError("unavailable"))) + monkeypatch.setattr(global_mcp_server_manager, "is_config_declared_server", lambda _: config_store) + update = AsyncMock() + upsert = AsyncMock() + monkeypatch.setattr(db, "update_mcp_server", update) + monkeypatch.setattr(db, "upsert_mcp_server_oauth_client_credentials", upsert) + result = await endpoints._persist_dcr_client_registration(server, {"client_id": "new-client"}, "https://gateway.example/callback") + assert result == "failed" + assert server.client_id is None + update.assert_not_awaited() + upsert.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("winner_available", [False, True]) +async def test_registration_losing_conditional_write_reuses_only_a_matching_winner(monkeypatch, winner_available): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints, db + from litellm.proxy._types import LiteLLM_MCPServerTable + from litellm.proxy import utils + + server = _dcr_redirect_test_server(None).model_copy(update={"url": "https://upstream.example/mcp"}) + row = LiteLLM_MCPServerTable.model_validate(server.model_dump(exclude_none=True)) + winner = row.model_copy(update={"credentials": { + "client_id": "winner-client", "redirect_uris": ["https://gateway.example/callback"], + "dcr_server_url": server.url, + }}) if winner_available else row + monkeypatch.setattr(utils, "get_prisma_client_or_throw", lambda _: MagicMock()) + read = AsyncMock(side_effect=[row, winner]) + monkeypatch.setattr(db, "get_mcp_server", read) + monkeypatch.setattr(endpoints, "_refresh_persisted_dcr_server", AsyncMock()) + update = AsyncMock(return_value=None) + monkeypatch.setattr(db, "update_mcp_server", update) + result = await endpoints._persist_dcr_client_registration(server, {"client_id": "losing-client"}, "https://gateway.example/callback") + assert result == ("reused" if winner_available else "failed") + assert update.await_args.kwargs["expected_updated_at"] == row.updated_at + assert server.client_id == ("winner-client" if winner_available else None) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_partial_update.py index a3c52dc16b7..122b22f6273 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -17,6 +17,7 @@ from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.db import ( create_mcp_server, + decrypt_credentials, set_mcp_server_pinned_tools, update_mcp_server, ) @@ -1181,3 +1182,186 @@ async def test_protocol_update_preserves_missing_server_without_writing(clear_al result = await update_mcp_server(prisma, payload, "admin") assert result is None table.update.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("changed,previous_issuer", [ + ({"url": "https://new.example/mcp"}, None), + ({"issuer": "https://new.example"}, "https://old.example"), + ({"url": "https://new.example/mcp", "auth_type": "api_key"}, "https://old.example"), +]) +async def test_upstream_identity_edit_drops_previous_oauth_client(changed, previous_issuer): + prisma = _mock_prisma() + existing = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", auth_type="oauth2", + url="https://old.example/mcp", issuer=previous_issuer, + credentials=json.dumps({"client_id": "old-client", "client_secret": "old-secret", "upstream_resource": "api://resource"}), + ) + prisma.db.litellm_mcpservertable.find_unique.return_value = existing + await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="test-server", **changed), "test-user") + written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + assert "credentials" in written + credentials = json.loads(written["credentials"]) + assert "client_id" not in credentials + assert "client_secret" not in credentials + assert credentials["upstream_resource"] == "api://resource" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("issuer", ["https://old.example/", "https://OLD.example:443"]) +async def test_distinct_issuer_identifier_edit_discards_registered_client(issuer): + prisma = _mock_prisma() + existing = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", auth_type="oauth2", + url="https://resource.example/mcp", issuer="https://old.example", + credentials=json.dumps({"client_id": "old-client", "dcr_issuer": "https://old.example"}), + ) + prisma.db.litellm_mcpservertable.find_unique.return_value = existing + await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="test-server", issuer=issuer), "test-user") + written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + assert "client_id" not in json.loads(written["credentials"]) + assert written["issuer"] == issuer + + +@pytest.mark.asyncio +@pytest.mark.parametrize("previous_issuer,binding", [("https://idp.example", None), (None, "https://idp.example")]) +@pytest.mark.parametrize("submitted_tokens", [None, {"access_token": "old-token", "refresh_token": "old-refresh", "expires_in": 3600}, {"access_token": "fresh-token", "refresh_token": "fresh-refresh", "expires_in": 3600}]) +async def test_url_edit_preserves_client_bound_to_previous_known_issuer_without_old_tokens(previous_issuer, binding, submitted_tokens): + prisma = _mock_prisma() + existing = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", auth_type="oauth2", + url="https://resource.example/mcp", issuer=previous_issuer, + credentials=json.dumps({"client_id": "static-client", "client_secret": "static-secret", "dcr_issuer": binding, + "access_token": "old-token", "refresh_token": "old-refresh", "expires_in": 3600, + "token_endpoint_auth_method": "client_secret_basic"}), + ) + prisma.db.litellm_mcpservertable.find_unique.return_value = existing + await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="test-server", url=existing.url + "?v=2", **({"credentials": submitted_tokens} if submitted_tokens is not None else {})), "test-user") + written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + credentials = json.loads(written["credentials"]) + assert credentials["client_id"] == "static-client" + assert credentials["client_secret"] == "static-secret" + assert credentials["dcr_issuer"] == "https://idp.example" + assert credentials["dcr_server_url"] == existing.url + assert "access_token" not in credentials + assert "refresh_token" not in credentials + assert "expires_in" not in credentials + assert credentials["token_endpoint_auth_method"] == "client_secret_basic" + assert written["issuer"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("changed", [0, 1]) +async def test_registration_write_requires_unchanged_server_revision(changed): + from datetime import datetime, timezone + from litellm.proxy._experimental.mcp_server.db import _update_mcp_server_row + + prisma = _mock_prisma() + table = prisma.db.litellm_mcpservertable + revision = datetime.now(timezone.utc) + table.update_many.return_value = changed + result = await _update_mcp_server_row( + prisma, server_id="test-server", data_dict={"credentials": "{}"}, expected_updated_at=revision, + ) + table.update_many.assert_awaited_once_with( + where={"server_id": "test-server", "updated_at": revision}, data={"credentials": "{}"}, + ) + table.update.assert_not_awaited() + assert (result is not None) is bool(changed) + if changed: + assert result.server_id == "test-server" + else: + table.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["create", "rotate"]) +async def test_explicit_oauth_client_write_binds_to_server_identity(operation): + prisma = _mock_prisma() + existing = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", auth_type="oauth2", + url="https://resource.example/mcp", issuer="https://idp.example", + credentials=json.dumps({"client_id": "old-client"}), + ) + prisma.db.litellm_mcpservertable.find_unique.return_value = existing + if operation == "create": + await create_mcp_server(prisma, NewMCPServerRequest( + server_name="bound-client", transport="http", auth_type="oauth2", + url=existing.url, issuer=existing.issuer, credentials={"client_id": "new-client"}, + ), "admin") + written = prisma.db.litellm_mcpservertable.create.call_args.kwargs["data"] + else: + await update_mcp_server(prisma, UpdateMCPServerRequest( + server_id=existing.server_id, credentials={"client_id": "new-client"}, + ), "admin") + written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + credentials = json.loads(written["credentials"]) + assert credentials["dcr_issuer"] == existing.issuer + assert credentials["dcr_server_url"] == existing.url + + +@pytest.mark.asyncio +async def test_unrelated_edit_does_not_backfill_legacy_oauth_binding(): + prisma = _mock_prisma() + existing = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", auth_type="oauth2", + url="https://resource.example/mcp", issuer="https://idp.example", + credentials=json.dumps({"client_id": "legacy-client"}), + ) + prisma.db.litellm_mcpservertable.find_unique.return_value = existing + await update_mcp_server(prisma, UpdateMCPServerRequest( + server_id=existing.server_id, credentials={"scopes": ["tools.read"]}, + ), "admin") + written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + credentials = json.loads(written["credentials"]) + assert credentials["client_id"] == "legacy-client" + assert "dcr_issuer" not in credentials + assert "dcr_server_url" not in credentials + + +@pytest.mark.asyncio +async def test_issuer_edit_does_not_rebind_resubmitted_saved_client(): + prisma = _mock_prisma() + existing = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", auth_type="oauth2", + url="https://old.example/mcp", issuer="https://old.example", + credentials=json.dumps({"client_id": "saved-client", "client_secret": "saved-secret"}), + ) + prisma.db.litellm_mcpservertable.find_unique.return_value = existing + await update_mcp_server(prisma, UpdateMCPServerRequest( + server_id=existing.server_id, url="https://new.example/mcp", issuer="https://new.example", + credentials={"client_id": "saved-client", "client_secret": "saved-secret"}, + ), "admin") + written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + credentials = json.loads(written["credentials"]) + assert "client_id" not in credentials + assert "client_secret" not in credentials + + +@pytest.mark.asyncio +@pytest.mark.parametrize("replacement", [ + {"client_secret": "replacement-secret"}, + {"client_secret": "old-secret", "token_endpoint_auth_method": "client_secret_basic"}, + {"client_secret": None}, + {"dcr_issuer": "https://new.example", "dcr_server_url": "https://new.example/mcp"}, +]) +async def test_issuer_edit_preserves_replacement_with_same_client_id(replacement): + prisma = _mock_prisma() + existing = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", auth_type="oauth2", + url="https://old.example/mcp", issuer="https://old.example", + credentials=json.dumps({"client_id": "shared-client", "client_secret": "old-secret"}), + ) + prisma.db.litellm_mcpservertable.find_unique.return_value = existing + submitted = {"client_id": "shared-client", **replacement} + await update_mcp_server(prisma, UpdateMCPServerRequest( + server_id=existing.server_id, url="https://new.example/mcp", issuer="https://new.example", + credentials=submitted, + ), "admin") + written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + credentials = decrypt_credentials(json.loads(written["credentials"])) + assert credentials["client_id"] == "shared-client" + assert credentials["dcr_issuer"] == "https://new.example" + assert credentials["dcr_server_url"] == "https://new.example/mcp" + assert credentials.get("client_secret") == replacement.get("client_secret") + assert credentials.get("token_endpoint_auth_method") == replacement.get("token_endpoint_auth_method") diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e7602d7f825..123a5505953 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -9186,6 +9186,8 @@ class TestMCPServerTimestamps: authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", + issuer="https://idp.example.com", + authorization_response_iss_parameter_supported=True, ) blipped = MCPServer( @@ -9200,6 +9202,16 @@ class TestMCPServerTimestamps: assert blipped.authorization_url == "https://idp.example.com/authorize" assert blipped.token_url == "https://idp.example.com/token" assert blipped.registration_url == "https://idp.example.com/register" + assert blipped.issuer == "https://idp.example.com" + assert blipped.authorization_response_iss_parameter_supported is True + fallback = MCPServerManager._merge_discovered_oauth_metadata( + blipped, MCPOAuthMetadata(from_origin_fallback=True), + ) + assert fallback.authorization_response_iss_parameter_supported is True + refreshed = MCPServerManager._merge_discovered_oauth_metadata( + fallback, MCPOAuthMetadata(discovered_issuer="https://idp.example.com"), + ) + assert refreshed.authorization_response_iss_parameter_supported is False same_authorize = MCPServer( server_id="s1", @@ -9316,8 +9328,8 @@ class TestMCPServerTimestamps: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _issuer_matches assert _issuer_matches("https://mcp.slack.com", "https://mcp.slack.com") - assert _issuer_matches("https://MCP.slack.com/", "https://mcp.slack.com") - assert _issuer_matches("https://mcp.slack.com:443", "https://mcp.slack.com") + assert not _issuer_matches("https://MCP.slack.com/", "https://mcp.slack.com") + assert not _issuer_matches("https://mcp.slack.com:443", "https://mcp.slack.com") assert _issuer_matches("https://login.example.com/tenant/v2.0", "https://login.example.com/tenant/v2.0") assert not _issuer_matches("https://attacker.example.com", "https://mcp.slack.com") assert not _issuer_matches("https://login.example.com/other/v2.0", "https://login.example.com/tenant/v2.0") diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 2cbf1d578b2..081a2d8ce73 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1876,7 +1876,7 @@ class TestTemporaryMCPSessionEndpoints: url="https://temp.example.com", transport=MCPTransport.http, ) - existing_server = MagicMock() + existing_server = MagicMock(dcr_issuer=None, dcr_server_url=None, token_endpoint_auth_method=None) existing_server.authentication_token = "token-abc" existing_server.client_id = "client-123" existing_server.client_secret = "secret-xyz" @@ -1912,7 +1912,7 @@ class TestTemporaryMCPSessionEndpoints: @staticmethod def _inherit_with(payload_credentials, **server_overrides): - existing_server = MagicMock() + existing_server = MagicMock(dcr_issuer=None, dcr_server_url=None, token_endpoint_auth_method=None) existing_server.authentication_token = None existing_server.client_id = "client-123" existing_server.client_secret = "secret-xyz" @@ -2616,6 +2616,9 @@ class TestTemporaryMCPSessionEndpoints: user_id="admin-user", ) inherited_server = MagicMock( + dcr_issuer=None, + dcr_server_url=None, + token_endpoint_auth_method=None, authentication_token="token-abc", client_id="client-id", client_secret="client-secret", @@ -11145,3 +11148,177 @@ async def test_protocol_update_on_missing_server_preserves_not_found() -> None: user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) assert error.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_changed_upstream_session_gets_an_isolated_id_without_saved_client(monkeypatch): + from litellm.proxy.management_endpoints import mcp_management_endpoints as management + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + saved = MCPServer( + server_id="saved-upstream", name="saved", transport="http", auth_type="oauth2", + url="https://old.example/mcp", issuer="https://old.example", + client_id="old-client", client_secret="old-secret", + ) + monkeypatch.setitem(management.global_mcp_server_manager.registry, saved.server_id, saved) + payload = NewMCPServerRequest( + server_id=saved.server_id, server_name="saved", transport="http", auth_type="oauth2", + oauth2_flow="authorization_code", url="https://new.example/mcp", issuer="https://new.example", + ) + staged = management._inherit_credentials_from_existing_server(payload) + assert not (staged.credentials or {}).get("client_id") + assert not (staged.credentials or {}).get("client_secret") + assert await management._resolve_session_server_id(staged) != saved.server_id + assert saved.client_id == "old-client" + + +def test_staged_url_edit_clears_resubmitted_issuer_and_endpoints(monkeypatch): + from litellm.proxy.management_endpoints import mcp_management_endpoints as management + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + saved = MCPServer( + server_id="saved-oauth", name="saved", transport="http", auth_type="oauth2", + url="https://old.example/mcp", issuer="https://old.example", + authorization_url="https://old.example/authorize", token_url="https://old.example/token", + registration_url="https://old.example/register", client_id="old-client", + ) + monkeypatch.setitem(management.global_mcp_server_manager.registry, saved.server_id, saved) + payload = NewMCPServerRequest( + server_id=saved.server_id, server_name="saved", transport="http", auth_type="oauth2", + oauth2_flow="authorization_code", url="https://new.example/mcp", issuer=saved.issuer, + authorization_url=saved.authorization_url, token_url=saved.token_url, registration_url=saved.registration_url, + ) + staged = management._inherit_credentials_from_existing_server(payload) + assert staged.issuer is None + assert staged.authorization_url is None + assert staged.token_url is None + assert staged.registration_url is None + assert staged.oauth2_flow == "authorization_code" + assert saved.issuer == "https://old.example" + + +@pytest.mark.asyncio +async def test_distinct_issuer_identifier_edit_isolates_session_client(monkeypatch): + from litellm.proxy.management_endpoints import mcp_management_endpoints as management + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + saved = MCPServer( + server_id="saved-oauth", name="saved", transport="http", auth_type="oauth2", + url="https://resource.example/mcp", issuer="https://idp.example", + client_id="saved-client", client_secret="saved-secret", + ) + monkeypatch.setitem(management.global_mcp_server_manager.registry, saved.server_id, saved) + payload = NewMCPServerRequest( + server_id=saved.server_id, server_name="saved", transport="http", auth_type="oauth2", + url=saved.url, issuer="https://IDP.example:443/", + ) + staged = management._inherit_credentials_from_existing_server(payload) + assert not (staged.credentials or {}).get("client_id") + assert not (staged.credentials or {}).get("client_secret") + assert await management._resolve_session_server_id(staged) != saved.server_id + + +@pytest.mark.asyncio +async def test_url_edit_stages_existing_client_with_previous_issuer_binding(monkeypatch): + from litellm.proxy.management_endpoints import mcp_management_endpoints as management + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + saved = MCPServer( + server_id="saved-oauth", name="saved", transport="http", auth_type="oauth2", + url="https://resource.example/mcp", issuer="https://idp.example", + client_id="static-client", client_secret="static-secret", authentication_token="old-token", + token_endpoint_auth_method="client_secret_basic", + ) + monkeypatch.setitem(management.global_mcp_server_manager.registry, saved.server_id, saved) + payload = NewMCPServerRequest( + server_id=saved.server_id, server_name="saved", transport="http", auth_type="oauth2", + url=saved.url + "?v=2", issuer=saved.issuer, + ) + staged = management._inherit_credentials_from_existing_server(payload) + assert staged.credentials["client_id"] == "static-client" + assert staged.credentials["client_secret"] == "static-secret" + assert staged.credentials["dcr_issuer"] == saved.issuer + assert staged.credentials["dcr_server_url"] == saved.url + assert staged.credentials["token_endpoint_auth_method"] == "client_secret_basic" + assert "auth_value" not in staged.credentials + assert staged.issuer is None + assert await management._resolve_session_server_id(staged) != saved.server_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("upstream_changed", [False, True]) +async def test_session_identity_checks_saved_row_when_registry_is_empty(monkeypatch, upstream_changed): + from litellm.proxy.management_endpoints import mcp_management_endpoints as management + + saved = generate_mock_mcp_server_db_record(server_id="saved-oauth-row") + saved.auth_type = "oauth2" + saved.url = "https://old.example/mcp" + saved.issuer = "https://old.example" + saved.approval_status = "approved" + monkeypatch.setattr(management.global_mcp_server_manager, "get_mcp_server_by_id", lambda _: None) + monkeypatch.setattr(management, "_get_prisma_client_or_none", lambda: MagicMock()) + monkeypatch.setattr(management, "get_mcp_server", AsyncMock(return_value=saved)) + payload = NewMCPServerRequest( + server_id=saved.server_id, auth_type="oauth2", transport="http", + url="https://new.example/mcp" if upstream_changed else saved.url, + ) + resolved = await management._resolve_session_server_id(payload) + assert (resolved != saved.server_id) is upstream_changed + + +def test_new_oauth_session_does_not_look_up_saved_credentials(monkeypatch): + from litellm.proxy.management_endpoints import mcp_management_endpoints as management + + lookup = MagicMock() + monkeypatch.setattr(management.global_mcp_server_manager, "get_mcp_server_by_id", lookup) + payload = NewMCPServerRequest(url="https://new.example/mcp", auth_type="oauth2", transport="http") + assert management._inherit_credentials_from_existing_server(payload) is payload + assert not payload.credentials + lookup.assert_not_called() + + +def test_edit_does_not_rebind_resubmitted_saved_client_to_new_issuer(monkeypatch): + from litellm.proxy.management_endpoints import mcp_management_endpoints as management + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + saved = MCPServer( + server_id="saved-static", name="saved", transport="http", auth_type="oauth2", + url="https://old.example/mcp", issuer="https://old.example", + client_id="saved-client", client_secret="saved-secret", + ) + monkeypatch.setitem(management.global_mcp_server_manager.registry, saved.server_id, saved) + payload = NewMCPServerRequest( + server_id=saved.server_id, transport="http", auth_type="oauth2", + url="https://new.example/mcp", issuer="https://new.example", + credentials={"client_id": saved.client_id, "client_secret": saved.client_secret}, + ) + staged = management._inherit_credentials_from_existing_server(payload) + assert not (staged.credentials or {}).get("client_id") + assert not (staged.credentials or {}).get("client_secret") + assert saved.client_id == "saved-client" + + +@pytest.mark.parametrize("replacement", [ + {"client_secret": "replacement-secret"}, + {"client_secret": "old-secret", "token_endpoint_auth_method": "client_secret_basic"}, + {"client_secret": None}, + {"dcr_issuer": "https://new.example", "dcr_server_url": "https://new.example/mcp"}, +]) +def test_staged_issuer_edit_preserves_replacement_with_same_client_id(monkeypatch, replacement): + from litellm.proxy.management_endpoints import mcp_management_endpoints as management + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + saved = MCPServer( + server_id="saved-static", name="saved", transport="http", auth_type="oauth2", + url="https://old.example/mcp", issuer="https://old.example", + client_id="shared-client", client_secret="old-secret", + ) + monkeypatch.setitem(management.global_mcp_server_manager.registry, saved.server_id, saved) + submitted = {"client_id": "shared-client", **replacement} + payload = NewMCPServerRequest( + server_id=saved.server_id, transport="http", auth_type="oauth2", + url="https://new.example/mcp", issuer="https://new.example", credentials=submitted, + ) + staged = management._inherit_credentials_from_existing_server(payload) + assert staged.credentials == submitted + assert saved.client_secret == "old-secret" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx index 3cf2bcdf238..f6ad2e38658 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx @@ -28,7 +28,10 @@ const oauthHook = vi.hoisted(() => ({ tokenResponse: null as Record | null, reset: vi.fn(), onTokenReceived: null as - | ((token: Record | null, registeredClient?: { clientId?: string; clientSecret?: string }) => void) + | (( + token: Record | null, + registeredClient?: { client_id: string; client_secret?: string }, + ) => void) | null, getCredentials: null as (() => Record | undefined) | null, getTemporaryPayload: null as (() => Record | null) | null, @@ -37,7 +40,7 @@ vi.mock("@/hooks/useMcpOAuthFlow", () => ({ useMcpOAuthFlow: (opts: { onTokenReceived: ( token: Record | null, - registeredClient?: { clientId?: string; clientSecret?: string }, + registeredClient?: { client_id: string; client_secret?: string }, ) => void; getCredentials?: () => Record | undefined; getTemporaryPayload?: () => Record | null; @@ -628,7 +631,7 @@ describe("CreateMCPServer", () => { await act(async () => { oauthHook.onTokenReceived!( { access_token: "oauth2-minted-tok", refresh_token: "oauth2-minted-refresh", token_type: "Bearer" }, - { clientId: "dcr-minted-client", clientSecret: "dcr-minted-secret" }, + { client_id: "dcr-minted-client", client_secret: "dcr-minted-secret" }, ); }); @@ -675,7 +678,7 @@ describe("CreateMCPServer", () => { await act(async () => { oauthHook.onTokenReceived!( { access_token: "oauth2-tok", token_type: "Bearer" }, - { clientId: "dcr-client", clientSecret: "dcr-secret" }, + { client_id: "dcr-client", client_secret: "dcr-secret" }, ); }); @@ -702,7 +705,7 @@ describe("CreateMCPServer", () => { await act(async () => { oauthHook.onTokenReceived!( { access_token: "oauth2-tok", token_type: "Bearer" }, - { clientId: "leak-client", clientSecret: "leak-secret" }, + { client_id: "leak-client", client_secret: "leak-secret" }, ); }); // Ref is held while the modal is open. @@ -731,7 +734,7 @@ describe("CreateMCPServer", () => { await act(async () => { oauthHook.onTokenReceived!( { access_token: "oauth2-tok", token_type: "Bearer" }, - { clientId: "dcr-client", clientSecret: "dcr-secret" }, + { client_id: "dcr-client", client_secret: "dcr-secret" }, ); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx index 2d6e83aff9b..d53d11ddcfb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx @@ -56,7 +56,7 @@ import EnvVarsSection from "./EnvVarsSection"; import { isAdminRole } from "@/utils/roles"; import { validateMCPServerUrl, validateMCPServerName } from "./utils"; import { toast } from "@/lib/toast"; -import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; +import { useMcpOAuthFlow, type McpDcrCredentials } from "@/hooks/useMcpOAuthFlow"; import { useTestMCPConnection } from "@/hooks/useTestMCPConnection"; import { MountedFormField, @@ -144,7 +144,7 @@ const CreateMCPServer: React.FC = ({ // it can never be collected as a client-forwarded server's declared app; injected into the payload // only on an oauth2 submit (where persisting the registered client is correct), and cleared on any // invalidation or modal close. An abandoned authorize leaves it null, which is the desired asymmetry. - const dcrClientRef = React.useRef<{ client_id: string; client_secret?: string } | null>(null); + const dcrClientRef = React.useRef(null); // Set when the upstream identity (url/endpoints) changed while a declared app is present, so the // section can warn that the saved app may not match the new upstream (the app is kept, not wiped). const [appMayNotMatchUpstream, setAppMayNotMatchUpstream] = useState(false); @@ -264,12 +264,7 @@ const CreateMCPServer: React.FC = ({ // The DCR-minted client is held in a ref, NOT written into form.credentials, so it can never be // collected as a client-forwarded server's declared app; it is injected into the payload only on // an oauth2 submit. An admin-typed client already lives in form.credentials and is left untouched. - dcrClientRef.current = registeredClient?.clientId - ? { - client_id: registeredClient.clientId, - ...(registeredClient.clientSecret && { client_secret: registeredClient.clientSecret }), - } - : null; + dcrClientRef.current = registeredClient ?? null; const current = (allFieldsValue(form).credentials as Record | undefined) ?? {}; const nextCredentials = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts index f45857fcc8b..30c4c502c42 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts @@ -29,7 +29,7 @@ export const AUTH_TYPES_REQUIRING_CREDENTIALS = [ export interface DcrClient { readonly client_id: string; - readonly client_secret?: string; + readonly client_secret?: string | null; } export interface CreateServerUiState { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts index 66786fac1c3..8a9bac582d6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts @@ -370,3 +370,31 @@ void isClientForwardedTokenMode; void normalizeEnvVars; void preservedAdminCredentials; export type { MCPServer }; + +it.each([ + [AUTH_TYPE.OAUTH2, OAUTH_FLOW.INTERACTIVE, true], + [AUTH_TYPE.OAUTH2, OAUTH_FLOW.M2M, false], + [AUTH_TYPE.TRUE_PASSTHROUGH, OAUTH_FLOW.INTERACTIVE, false], +])("applies a pending DCR client only to an interactive gateway OAuth save (%s, %s)", (authType, flow, useDcr) => { + const values: EditServerFormValues = { + auth_type: authType, + oauth_flow_type: flow, + transport: "http", + credentials: { client_id: "configured-client", client_secret: "configured-secret" }, + }; + const ui: EditServerUiState = { + ...baseUi, + dcrClient: { + client_id: "new-client", + client_secret: null, + dcr_issuer: "https://new.example", + dcr_server_url: "https://new.example/mcp", + }, + }; + const result = buildEditServerPayload(values, ui); + expect(result.kind).toBe("ok"); + if (result.kind !== "ok") throw new Error("Expected valid MCP server payload"); + expect(result.payload.credentials?.client_id).toBe(useDcr ? "new-client" : "configured-client"); + expect(result.payload.credentials?.client_secret).toBe(useDcr ? null : "configured-secret"); + expect(result.payload.credentials?.dcr_issuer).toBe(useDcr ? "https://new.example" : undefined); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.ts index 3930f1edb9b..f166fa20d1f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.ts @@ -1,3 +1,4 @@ +import type { McpDcrCredentials } from "@/hooks/useMcpOAuthFlow"; import { ADMIN_CONFIG_CREDENTIAL_KEYS, AUTH_TYPE, @@ -73,6 +74,7 @@ export interface EditServerPayload { } export interface EditServerUiState { + readonly dcrClient?: McpDcrCredentials | null; readonly mcpServer: MCPServer; readonly logoUrl: string | undefined; readonly costConfig: MCPServerCostInfo; @@ -213,6 +215,8 @@ const buildCredentials = (credentialValues: unknown): Readonly> | undefined; readonly includeCredentials: boolean; @@ -224,11 +228,16 @@ interface CredentialsEntryInput { // updates), so removal must be an explicit-null write: encrypt skips nulls and the merge overrides // the stored keys, returning the server to dynamic client registration. const resolveCredentialsEntry = ({ + dcrClient, + oauthFlow, authType, credentials, includeCredentials, removeStoredApp, }: CredentialsEntryInput): { readonly credentials?: Readonly> } => { + if (dcrClient && authType === AUTH_TYPE.OAUTH2 && oauthFlow !== OAUTH_FLOW.M2M) { + return { credentials: { ...credentials, ...dcrClient } }; + } if (removeStoredApp && isClientForwardedTokenMode(authType)) { return { credentials: { client_id: null, client_secret: null } }; } @@ -330,6 +339,8 @@ export const buildEditServerPayload = (values: EditServerFormValues, ui: EditSer const includeCredentials = restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type); const credentialsEntryInput: CredentialsEntryInput = { + dcrClient: ui.dcrClient, + oauthFlow: restValues.oauth_flow_type, authType: restValues.auth_type, credentials: submitCredentials, includeCredentials: Boolean(includeCredentials), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx index f79dd178571..91db38db389 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx @@ -16,21 +16,40 @@ vi.mock("@/components/networking", () => ({ })); const mockOauth: { + status: string; tokenResponse: any; getTemporaryPayload: (() => Record | null) | null; - onTokenReceived: ((token: Record | null) => void) | null; + onTokenReceived: + | (( + token: Record | null, + registeredClient?: { + client_id: string; + client_secret?: string; + dcr_issuer?: string; + dcr_server_url?: string; + }, + ) => void) + | null; reset: ReturnType; -} = { tokenResponse: null, getTemporaryPayload: null, onTokenReceived: null, reset: vi.fn() }; +} = { status: "idle", tokenResponse: null, getTemporaryPayload: null, onTokenReceived: null, reset: vi.fn() }; vi.mock("@/hooks/useMcpOAuthFlow", () => ({ useMcpOAuthFlow: (opts: { getTemporaryPayload?: () => Record | null; - onTokenReceived?: (token: Record | null) => void; + onTokenReceived?: ( + token: Record | null, + registeredClient?: { + client_id: string; + client_secret?: string; + dcr_issuer?: string; + dcr_server_url?: string; + }, + ) => void; }) => { mockOauth.getTemporaryPayload = opts?.getTemporaryPayload ?? null; mockOauth.onTokenReceived = opts?.onTokenReceived ?? null; return { startOAuthFlow: vi.fn(), - status: "idle", + status: mockOauth.status, error: null, tokenResponse: mockOauth.tokenResponse, reset: mockOauth.reset, @@ -547,6 +566,7 @@ describe("MCPServerEdit (auth type switch)", () => { describe("MCPServerEdit OAuth token invalidation", () => { beforeEach(() => { vi.clearAllMocks(); + mockOauth.status = "idle"; }); const renderOAuthEdit = () => @@ -560,6 +580,77 @@ describe("MCPServerEdit OAuth token invalidation", () => { />, ); + it.each(["authorizing", "exchanging"])("blocks Save while OAuth is %s", async (status) => { + mockOauth.status = status; + vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer }); + const view = renderOAuthEdit(); + const save = screen.getAllByRole("button", { name: "Save Changes" })[0]; + expect(save).toBeDisabled(); + await act(async () => { + fireEvent.submit(screen.getByRole("form", { name: "Edit MCP server" })); + }); + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + act(() => { + mockOauth.onTokenReceived?.( + { access_token: "new-token" }, + { + client_id: "new-client", + dcr_issuer: "https://new.example", + dcr_server_url: "https://new.example/mcp", + }, + ); + }); + mockOauth.status = "success"; + view.rerender( + , + ); + expect(screen.getAllByRole("button", { name: "Save Changes" })[0]).toBeEnabled(); + await act(async () => { + fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]); + }); + await waitFor(() => expect(networking.updateMCPServer).toHaveBeenCalledTimes(1)); + expect(vi.mocked(networking.updateMCPServer).mock.calls[0][1].credentials).toMatchObject({ + client_id: "new-client", + dcr_issuer: "https://new.example", + }); + }); + + it.each(["Save Changes", "Cancel"])("keeps a newly registered client isolated until %s", async (action) => { + vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer }); + renderOAuthEdit(); + act(() => { + mockOauth.onTokenReceived?.( + { access_token: "new-token" }, + { + client_id: "new-client", + client_secret: "new-secret", + dcr_issuer: "https://new.example", + dcr_server_url: "https://new.example/mcp", + }, + ); + }); + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + await act(async () => { + fireEvent.click(screen.getAllByRole("button", { name: action })[0]); + }); + if (action === "Cancel") { + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + return; + } + await waitFor(() => expect(networking.updateMCPServer).toHaveBeenCalledTimes(1)); + expect(vi.mocked(networking.updateMCPServer).mock.calls[0][1].credentials).toMatchObject({ + client_id: "new-client", + client_secret: "new-secret", + dcr_issuer: "https://new.example", + }); + }); + it("invalidates a session-authorized token when the transport switches to stdio", async () => { // Switching to stdio clears url/auth_type via programmatic form.setFieldsValue, which antd does // not report through onValuesChange; the explicit recheck in handleTransportChange must catch it. @@ -1687,6 +1778,51 @@ describe("MCPServerEdit (OAuth token persistence on save)", () => { expect(screen.getByText(/registered for the previous upstream/)).toBeInTheDocument(); }); + it("discards a canceled OAuth snapshot before saved server data loads", async () => { + setSecureItem( + EDIT_OAUTH_UI_STATE_KEY, + JSON.stringify({ + serverId: interactiveOAuthServer.server_id, + formValues: { ...interactiveOAuthServer, url: "https://new.example/mcp" }, + }), + ); + const props = { accessToken: "access-token", onCancel: vi.fn(), onSuccess: vi.fn(), availableAccessGroups: [] }; + const view = render(); + await act(async () => { + fireEvent.click(screen.getAllByRole("button", { name: "Cancel" })[0]); + }); + expect(props.onCancel).toHaveBeenCalledOnce(); + expect(window.sessionStorage.getItem(EDIT_OAUTH_UI_STATE_KEY)).toBeNull(); + expect(mockOauth.reset).toHaveBeenCalled(); + view.unmount(); + render(); + await waitFor(() => expect(screen.getByLabelText("MCP Server URL")).toHaveValue(interactiveOAuthServer.url)); + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + }); + + it("restores the edited upstream after OAuth when saved server data loads later", async () => { + setSecureItem( + EDIT_OAUTH_UI_STATE_KEY, + JSON.stringify({ + serverId: interactiveOAuthServer.server_id, + formValues: { ...interactiveOAuthServer, url: "https://new.example/mcp", issuer: "https://new.example" }, + }), + ); + const props = { + accessToken: "access-token", + userID: "user-1", + onCancel: vi.fn(), + onSuccess: vi.fn(), + availableAccessGroups: [], + }; + const { rerender } = render( + , + ); + rerender(); + await waitFor(() => expect(screen.getByLabelText("MCP Server URL")).toHaveValue("https://new.example/mcp")); + expect(screen.getByLabelText("Issuer (optional)")).toHaveValue("https://new.example"); + }); + it("preserves a stored client_id on OAuth-resume restore even when the saved snapshot is token-only", async () => { // Post-redirect restore: the sessionStorage snapshot carries only a minted token (no client keys), // while the loaded server has a stored client_id. The restore must merge the server's declared app diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 3894cb6cd0b..1b1e4c27bde 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -55,7 +55,7 @@ import { EditServerFormValues, buildEditServerPayload, editPayloadErrorMessage } import { DUPLICATE_IDENTIFIER_MESSAGE, findDuplicateMcpServer, mcpSubmitErrorReason } from "./duplicateServerCheck"; import { toast } from "@/lib/toast"; import { getEditToolPreview } from "./editToolPreview"; -import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; +import { useMcpOAuthFlow, type McpDcrCredentials } from "@/hooks/useMcpOAuthFlow"; import { MountedFormField, MountedFormProvider, @@ -250,6 +250,7 @@ const MCPServerEdit: React.FC = ({ // in this edit session; undefined when none is held. If a mint-relevant field later diverges from it, // the held token (hook response + sessionStorage) is discarded so the admin must re-authorize. const authorizedIdentityRef = React.useRef(undefined); + const dcrClientRef = React.useRef(null); const { startOAuthFlow, @@ -300,7 +301,7 @@ const MCPServerEdit: React.FC = ({ env: values.env, }; }, - onTokenReceived: (token) => { + onTokenReceived: (token, registeredClient) => { if (!token?.access_token) { return; } @@ -319,6 +320,7 @@ const MCPServerEdit: React.FC = ({ return; } + dcrClientRef.current = registeredClient?.dcr_server_url ? registeredClient : null; const current = (allFieldsValue(form).credentials as Record | undefined) ?? {}; const nextCredentials = { ...(preservedAdminCredentials(current) ?? {}), @@ -393,6 +395,9 @@ const MCPServerEdit: React.FC = ({ if (!parsed || parsed.serverId !== mcpServer.server_id) { return; } + // The saved server may still be loading on the first render after the redirect. + // Consume this snapshot only after the matching server can restore it. + window.sessionStorage.removeItem(EDIT_OAUTH_UI_STATE_KEY); if (parsed.formValues) { // Rebuild credentials from the declared app in EITHER the loaded server or the saved snapshot, // then strip minted token material. Merging the two (server under snapshot) before stripping is @@ -426,7 +431,6 @@ const MCPServerEdit: React.FC = ({ } } catch (err) { console.error("Failed to restore MCP edit state", err); - } finally { window.sessionStorage.removeItem(EDIT_OAUTH_UI_STATE_KEY); } }, [form, mcpServer]); @@ -497,6 +501,7 @@ const MCPServerEdit: React.FC = ({ removeToken(mcpServer.server_id, userID); } setTools([]); + dcrClientRef.current = null; resetOAuthFlow(); // The admin-typed app is upstream-scoped config, not minted material, so it survives every // invalidation; only the held token is discarded. Token-shaped keys are excluded by the filter. @@ -720,7 +725,16 @@ const MCPServerEdit: React.FC = ({ return () => subscription.unsubscribe(); }, [form]); + const isOAuthPending = ["authorizing", "exchanging"].includes(oauthStatus); + + const handleCancel = () => { + window.sessionStorage.removeItem(EDIT_OAUTH_UI_STATE_KEY); + resetOAuthFlow(); + onCancel(); + }; + const submitForm = async () => { + if (isOAuthPending) return; const isValid = await form.trigger(mountedPaths(registry) as string[]); if (!isValid) { return; @@ -742,7 +756,8 @@ const MCPServerEdit: React.FC = ({ return; } try { - const built = buildEditServerPayload(values, { + const uiState = { + dcrClient: dcrClientRef.current, mcpServer, logoUrl, costConfig, @@ -752,7 +767,8 @@ const MCPServerEdit: React.FC = ({ toolNameToDisplayName, toolNameToDescription, removeStoredApp, - }); + }; + const built = buildEditServerPayload(values, uiState); if (built.kind !== "ok") { toast.fromError(editPayloadErrorMessage(built)); return; @@ -820,6 +836,7 @@ const MCPServerEdit: React.FC = ({
{ event.preventDefault(); void submitForm(); @@ -1308,10 +1325,12 @@ const MCPServerEdit: React.FC = ({
- - +
@@ -1323,10 +1342,12 @@ const MCPServerEdit: React.FC = ({
- - +
diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx index 2e809cfd728..2dc404a9218 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx @@ -1,7 +1,7 @@ import { act, renderHook, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import * as networking from "@/components/networking"; -import { setSecureItem } from "@/utils/secureStorage"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; import { useMcpOAuthFlow } from "./useMcpOAuthFlow"; vi.mock("@/components/networking", () => ({ @@ -73,7 +73,7 @@ describe("useMcpOAuthFlow reset", () => { await waitFor(() => expect(result.current.status).toBe("success")); expect(result.current.tokenResponse).toEqual(token); - expect(onTokenReceived).toHaveBeenCalledWith(token, expect.objectContaining({ clientId: "client-1" })); + expect(onTokenReceived).toHaveBeenCalledWith(token, expect.objectContaining({ client_id: "client-1" })); act(() => { result.current.reset(); @@ -84,6 +84,51 @@ describe("useMcpOAuthFlow reset", () => { expect(result.current.error).toBeNull(); }); + it("retains client binding from a redirect started before the state format changed", async () => { + seedCompletedRedirect(); + const client = { + client_id: "client-1", + client_secret: "registered-secret", + dcr_issuer: "https://issuer.example.com", + dcr_server_url: "https://server-1.example.com/mcp", + redirect_uris: ["https://app.example.com/ui/mcp/oauth/callback"], + }; + const state = JSON.parse(getSecureItem(FLOW_STATE_KEY)!); + setSecureItem( + FLOW_STATE_KEY, + JSON.stringify({ ...state, clientSecret: client.client_secret, dcrCredentials: client }), + ); + const token = { access_token: "registered-token" }; + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue(token); + const onTokenReceived = vi.fn(); + const { result } = renderFlow(onTokenReceived); + await waitFor(() => expect(result.current.status).toBe("success")); + expect(onTokenReceived).toHaveBeenCalledWith(token, client); + expect(networking.exchangeMcpOAuthToken).toHaveBeenCalledWith( + expect.objectContaining({ clientId: client.client_id, clientSecret: client.client_secret }), + ); + }); + + it("resumes a legacy server-managed flow without exposing a registered client", async () => { + seedCompletedRedirect(); + const state = JSON.parse(getSecureItem(FLOW_STATE_KEY)!); + delete state.clientId; + setSecureItem(FLOW_STATE_KEY, JSON.stringify(state)); + const token = { access_token: "server-managed-token" }; + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue(token); + const onTokenReceived = vi.fn(); + const { result } = renderFlow(onTokenReceived); + await waitFor(() => expect(result.current.status).toBe("success")); + expect(onTokenReceived).toHaveBeenCalledWith(token, undefined); + expect(networking.exchangeMcpOAuthToken).toHaveBeenCalledWith( + expect.objectContaining({ + serverId: "server-1", + clientId: undefined, + clientSecret: undefined, + }), + ); + }); + it("ignores an in-flight exchange result after reset", async () => { const token = { access_token: "stale-token" }; let resolveExchange: (value: typeof token) => void = () => undefined; @@ -137,7 +182,7 @@ describe("useMcpOAuthFlow reset", () => { rerender({ onTokenReceived: onTokenReceived2 }); await waitFor(() => - expect(onTokenReceived2).toHaveBeenCalledWith(token, expect.objectContaining({ clientId: "client-1" })), + expect(onTokenReceived2).toHaveBeenCalledWith(token, expect.objectContaining({ client_id: "client-1" })), ); }); @@ -163,8 +208,8 @@ describe("useMcpOAuthFlow reset", () => { await waitFor(() => expect(result.current.status).toBe("success")); expect(onTokenReceived).toHaveBeenCalledWith(token, { - clientId: "dcr-client-xyz", - clientSecret: "dcr-secret-abc", + client_id: "dcr-client-xyz", + client_secret: "dcr-secret-abc", }); }); @@ -196,13 +241,27 @@ describe("useMcpOAuthFlow reset", () => { ); }); - it("registers a fresh client when no client_id is present (new URL after the derived client is cleared)", async () => { - vi.mocked(networking.cacheTemporaryMcpServer).mockResolvedValue({ server_id: "server-2" }); - vi.mocked(networking.registerMcpOAuthClient).mockResolvedValue({ client_id: "fresh-client" }); - vi.mocked(networking.buildMcpOAuthAuthorizeUrl).mockReturnValue("https://idp.example.com/authorize"); + it.each([ + { method: "none", secret: undefined, issuer: "https://idp.example.com", bound: true }, + { method: "client_secret_basic", secret: "registered-secret", issuer: "https://idp.example.com", bound: true }, + { method: "none", secret: undefined, issuer: undefined, bound: true }, + { method: "none", secret: undefined, issuer: undefined, bound: false }, + ])( + "retains a fresh $method client with issuer $issuer and binding $bound", + async ({ method, secret, issuer, bound }) => { + vi.mocked(networking.cacheTemporaryMcpServer).mockResolvedValue({ server_id: "server-2" }); + const registration = { + client_id: "fresh-client", + client_secret: secret, + token_endpoint_auth_method: method, + dcr_issuer: issuer, + dcr_server_url: bound ? "https://server-2.example.com/mcp" : undefined, + dcr_redirect_uris: ["https://gateway.example.com/callback"], + }; + vi.mocked(networking.registerMcpOAuthClient).mockResolvedValue(registration); + vi.mocked(networking.buildMcpOAuthAuthorizeUrl).mockReturnValue("https://idp.example.com/authorize"); - const { result } = renderHook(() => - useMcpOAuthFlow({ + const options = { accessToken: "admin-token", getCredentials: () => ({}), getTemporaryPayload: () => ({ @@ -212,16 +271,45 @@ describe("useMcpOAuthFlow reset", () => { }), onTokenReceived: vi.fn(), flowSource: "create", - }), - ); + }; + const { result } = renderHook(() => useMcpOAuthFlow(options)); - await act(async () => { - await result.current.startOAuthFlow(); - }); + await act(async () => { + await result.current.startOAuthFlow(); + }); - expect(networking.registerMcpOAuthClient).toHaveBeenCalledTimes(1); - expect(networking.buildMcpOAuthAuthorizeUrl).toHaveBeenCalledWith( - expect.objectContaining({ clientId: "fresh-client" }), - ); - }); + expect(JSON.parse(getSecureItem(FLOW_STATE_KEY)!)).toEqual( + expect.objectContaining({ + client: { + client_id: "fresh-client", + client_secret: secret ?? null, + ...(bound + ? { + token_endpoint_auth_method: method === "client_secret_basic" ? method : null, + dcr_issuer: issuer ?? null, + dcr_server_url: "https://server-2.example.com/mcp", + redirect_uris: ["https://gateway.example.com/callback"], + } + : {}), + }, + }), + ); + const stored = JSON.parse(getSecureItem(FLOW_STATE_KEY)!); + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue({ access_token: "new-token" }); + setSecureItem(RESULT_KEY, JSON.stringify({ state: stored.state, code: "new-code" })); + const resumed = renderHook(() => useMcpOAuthFlow(options)); + await waitFor(() => expect(resumed.result.current.status).toBe("success")); + expect(options.onTokenReceived).toHaveBeenCalledWith({ access_token: "new-token" }, stored.client); + expect(networking.exchangeMcpOAuthToken).toHaveBeenCalledWith( + expect.objectContaining({ + clientId: "fresh-client", + clientSecret: secret, + }), + ); + expect(networking.registerMcpOAuthClient).toHaveBeenCalledTimes(1); + expect(networking.buildMcpOAuthAuthorizeUrl).toHaveBeenCalledWith( + expect.objectContaining({ clientId: "fresh-client" }), + ); + }, + ); }); diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx index 2d6f0344eb9..25e8790e367 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx @@ -16,20 +16,48 @@ import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; export type McpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; +export interface McpDcrCredentials { + client_id: string; + client_secret?: string | null; + dcr_issuer?: string | null; + dcr_server_url?: string | null; + token_endpoint_auth_method?: string | null; + redirect_uris?: string[]; +} + +const getRegisteredOAuthClient = ( + registration: (Partial & { dcr_redirect_uris?: string[] }) | undefined, +): McpDcrCredentials | undefined => + registration?.client_id + ? { + client_id: registration.client_id, + client_secret: registration.client_secret ?? null, + ...(registration.dcr_server_url && { + dcr_issuer: registration.dcr_issuer ?? null, + dcr_server_url: registration.dcr_server_url, + token_endpoint_auth_method: + registration.token_endpoint_auth_method === "client_secret_basic" ? "client_secret_basic" : null, + redirect_uris: registration.dcr_redirect_uris, + }), + } + : undefined; + +const oauthClientRequest = (client: McpDcrCredentials | undefined) => ({ + clientId: client?.client_id, + clientSecret: client?.client_secret ?? undefined, +}); + interface UseMcpOAuthFlowOptions { accessToken: string | null; getCredentials: () => | { client_id?: string; - client_secret?: string; + client_secret?: string | null; scopes?: string[]; } | undefined; getTemporaryPayload: () => Record | null; - onTokenReceived: ( - tokenResponse: Record, - registeredClient?: { clientId?: string; clientSecret?: string }, - ) => void; + onTokenReceived: (tokenResponse: Record, registeredClient?: McpDcrCredentials) => void; onBeforeRedirect?: () => void; // Distinguishes which form started the flow (e.g. "create" vs "edit"). Both forms // mount this hook with shared storage keys, so the return handler only processes a @@ -69,6 +97,8 @@ export const useMcpOAuthFlow = ({ codeVerifier: string; clientId?: string; clientSecret?: string; + client?: McpDcrCredentials; + dcrCredentials?: McpDcrCredentials; serverId: string; redirectUri: string; flowSource?: string; @@ -147,7 +177,7 @@ export const useMcpOAuthFlow = ({ throw new Error("Temporary MCP server identifier missing. Please retry."); } - let registeredClient: { clientId?: string; clientSecret?: string } = {}; + let registeredClient: McpDcrCredentials | undefined; const hasPreconfiguredCredentials = Boolean(temporaryPayload.credentials?.client_id); if (!hasPreconfiguredCredentials) { @@ -162,24 +192,22 @@ export const useMcpOAuthFlow = ({ // rejects the registration and the admin authorize dead-ends. redirect_uris: [callbackUrl()], }); - registeredClient = { - clientId: registration?.client_id, - clientSecret: registration?.client_secret, - }; + registeredClient = getRegisteredOAuthClient(registration); } const verifier = generateCodeVerifier(); const challenge = await generateCodeChallenge(verifier); const state = crypto.randomUUID(); - const clientId = registeredClient.clientId || credentials.client_id; + const client = registeredClient ?? getRegisteredOAuthClient(credentials); const scopeString = Array.isArray(credentials.scopes) ? credentials.scopes.filter((s) => s && s.trim().length > 0).join(" ") : undefined; + const { clientId } = oauthClientRequest(client); const authorizeUrl = buildMcpOAuthAuthorizeUrl({ serverId, - clientId: clientId, + clientId, redirectUri: callbackUrl(), state, codeChallenge: challenge, @@ -189,8 +217,7 @@ export const useMcpOAuthFlow = ({ const flowState: StoredFlowState = { state, codeVerifier: verifier, - clientId, - clientSecret: registeredClient.clientSecret || credentials.client_secret, + client, serverId, redirectUri: callbackUrl(), flowSource, @@ -311,12 +338,15 @@ export const useMcpOAuthFlow = ({ throw new Error("Authorization code missing in callback."); } + const client = + flowState.client ?? + flowState.dcrCredentials ?? + getRegisteredOAuthClient({ client_id: flowState.clientId, client_secret: flowState.clientSecret }); setStatus("exchanging"); const token = await exchangeMcpOAuthToken({ serverId: flowState.serverId, code: payload.code, - clientId: flowState.clientId, - clientSecret: flowState.clientSecret, + ...oauthClientRequest(client), codeVerifier: flowState.codeVerifier, redirectUri: flowState.redirectUri, accessToken, @@ -326,7 +356,7 @@ export const useMcpOAuthFlow = ({ return; } - onTokenReceived(token, { clientId: flowState.clientId, clientSecret: flowState.clientSecret }); + onTokenReceived(token, client); setTokenResponse(token); setStatus("success"); setError(null); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ae8b99448b0..73336ca1823 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -2273,7 +2273,9 @@ export interface paths { * * - A successful authorization response (``code`` + ``state``), which is * forwarded back to the validated client ``redirect_uri`` with the - * original (un-wrapped) ``state``. + * original (un-wrapped) ``state``, once the RFC 9207 ``iss`` (when the + * authorization server sent one) matches the issuer /authorize sealed + * into the state. * - An error response (``error``[+``error_description``/``error_uri``]), per * RFC 6749 §4.1.2.1. When ``state`` is present and decodes to a trusted * ``redirect_uri``, the error params are propagated back to the client so @@ -37147,6 +37149,10 @@ export interface components { client_private_key_id?: string | null; /** Client Secret */ client_secret?: string | null; + /** Dcr Issuer */ + dcr_issuer?: string | null; + /** Dcr Server Url */ + dcr_server_url?: string | null; /** Id Jag Resource */ id_jag_resource?: string | null; /** Id Jag Resource Token Endpoint */ @@ -53354,6 +53360,7 @@ export interface operations { query?: { code?: string | null; state?: string | null; + iss?: string | null; error?: string | null; error_description?: string | null; error_uri?: string | null; From d94670674434cc9f9ac50b974ea2b1e3bde9be32 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Mon, 5 Oct 2026 12:33:56 -0500 Subject: [PATCH 002/191] feat(guardrails): add llm shield pii redaction and rehydration guardrail (#42645) * feat(guardrails): add llm shield pii redaction and rehydration guardrail LLM Shield is a self-hosted PII gateway. This adds it as a guardrail so a proxy operator can redact personal data out of outbound requests and have the original values restored in the model's reply. The substitution is reversible, which is the difference from a masking guardrail. Outbound text is replaced with placeholders held in a session vault inside the operator's own LLM Shield deployment, and the reply is restored before it reaches the caller, so the end user still sees real values while the provider never received them. Streaming responses are restored incrementally. LLM Shield holds back only the trailing characters that could still turn out to be part of a placeholder, so tokens are forwarded as they arrive rather than the whole response being collected first. A placeholder split across two chunks is never emitted in fragments. The integration talks to LLM Shield over HTTP and adds no dependency. Notes for reviewers: - The guardrail sets use_native_lifecycle_hooks, since redaction and restoration need the native pre-call, post-call and streaming hooks rather than the unified path. - Per-request state lives on the request dict, never on the guardrail instance, because the proxy registers a single instance process-wide. The streaming carry-over is a local of the generator for the same reason. - Every failure blocks the request. A redaction guardrail that fails open would send the exact data it exists to protect to the provider. * feat(ui): list llm shield in the guardrail garden Adds the card, preset and logo so operators can pick LLM Shield from the guardrails page the same way as the other partner guardrails. * docs(guardrails): add llm shield example config Shows both modes on one entry. Listing only pre_call redacts the request and then hands the placeholders back to the end user, so the test asserts both hooks are enabled. * feat(ui): use the llm shield brand mark for the guardrail logo * fix(guardrails): restore llm shield values in anthropic replies The /v1/messages reply is a plain dict with a content block list and no choices, so it fell through the restore path and went back to the caller still carrying placeholders. The request was redacted correctly, which is what made this easy to miss. Found by running all three endpoints against a live provider; the mocked tests all passed because they only built the OpenAI shape. Adds tests for the message shape and for leaving non-text blocks alone. * docs(guardrails): correct the llm shield start command * fix(guardrails): redact every request shape and restore every reply shape Three gaps, all of which let an enabled guardrail hand data to the provider or hand placeholders to the caller. Requests only walked `messages`. The Responses API `input` and tool call `arguments` went out untouched. Measured against a live provider: a request sent through `/v1/responses` reached the model with the real address in it while the guardrail reported as enabled. Request traversal now covers chat content (string and multimodal), tool call arguments, and `input` as a bare string or a list of items. Fixing that exposed the matching gap on the way back: the Responses API reply carries `output` items rather than `choices`, so it returned to the caller still holding placeholders. It now gets its own walk, handling text blocks as dicts or objects. The dashboard preset seeded only pre_call, so a guardrail created from the UI would redact the request and return the placeholders to the user. Presets can now seed both modes; the form already normalised either shape. Adds tests for each request shape, for both Responses API reply forms, and replaces a test that had asserted the `input` bypass as correct behaviour. * fix(guardrails): narrow the stream delta before writing to it basedpyright could not prove the delta was non-None on the write path, and reportOptionalMemberAccess has a zero budget. The guard is also clearer than relying on the text check to imply it. * fix(guardrails): mint the vault id instead of trusting the caller's The vault id was taken from caller-supplied session metadata, and every caller shares one LLM Shield key. Someone who knew or guessed another caller's session id could send a placeholder, have the model echo it back, and get that caller's plaintext restored into their own reply. Vault ids are now minted per request behind a per-process prefix, so a caller cannot name a vault this process uses. Redaction mints, restoration reads back, and a reply whose id does not match is left holding its placeholders rather than resolved against some other vault. Also covers two more request fields that were reaching the provider intact: the Responses API `instructions`, and the legacy `function_call.arguments` alongside `tool_calls`. The collectors move to module level, which drops the traversal back under the complexity limit and lets the code carry its own explanation instead of the comments that were restating it. * fix(guardrails): drop Final from a loop-assigned local basedpyright rejects a Final assigned inside a loop, and reportGeneralTypeIssues sits one over its budget ceiling. * fix(guardrails): redact completion prompts and responses tool items Two more provider-bound request shapes were reaching the model intact while the guardrail reported as enabled. /v1/completions carries its text in a top-level `prompt`, which the traversal never looked at. It is handled as a string and as the array form, where each entry is rewritten in place. Responses input items hold tool data outside `content`: a function_call item in `arguments`, a function_call_output item in `output`. Both are now collected alongside the item's content. Adds a test per shape. * fix(guardrails): redact the anthropic system prompt and string-array input Two more provider-bound shapes, found by walking the request types rather than waiting for them to be reported. /v1/messages carries its system prompt at the top level, as a string or a list of text blocks. It is one of the endpoints this guardrail claims to cover, and a system prompt is a natural place to put a customer's details. `input` as an array of bare strings, the embeddings and moderations shape, was skipped because the loop only handled item dicts. Verified against a live provider: a system prompt holding an address now reaches the model as a stand-in and is restored in the reply. * fix(guardrails): narrow prompt and input to a list before iterating Guarding with a conditional iterable left the value un-narrowed, so passing it on was an argument-type error and the element checks read as unreachable. An early return narrows it properly and reads better. * fix(guardrails): restore every streaming choice, not just the first Streaming rehydration read and rewrote choices[0] only, so with n>1 every later choice went back to the caller still holding its placeholders. Each choice is its own token stream, so the sliding window is now tracked per choice index rather than once per stream. A single shared window would have been worse than the bug: it would splice the characters held back for one choice onto the next one's delta. The final flush walks every choice the same way, and the two helpers that only ever looked at choices[0] are gone. Adds a test that both choices come back restored, and one that each choice gets its own window handed back rather than its neighbour's. * refactor(guardrails): name the guardrail llm_shield_proxy throughout The integration was called llm_shield in code, llm-shield in the example config, and LLM Shield in the dashboard, while the product and its PyPI package are both llm-shield-proxy. An operator who saw the guardrail in LiteLLM could not tell what to install. One identifier now: llm_shield_proxy for the enum value, module, directory, class, config model, logo and environment variables, with LLM Shield Proxy as the display name. That matches `pip install llm-shield-proxy`. Renames only; no behaviour change. * feat(guardrails): redact the participant name on a message `name` on a user or assistant turn identifies a person and was going to the provider intact. The proxy this integrates with already redacts it, so the integration was the weaker of the two. On a tool or function turn the same field carries the function's name, which has to arrive unchanged or the call stops routing. That case is skipped, and a test asserts the value is never even sent to the shield. * fix(guardrails): flush every held choice, and cover tool results and suffix Three review findings. The trailing flush walked the last chunk's choices, so a choice that finished earlier and stopped appearing lost whatever text was still held for it and its answer was truncated. It is now driven by the windows themselves and emits one chunk per choice, synthesising the choice when the terminal chunk omits it. That was data loss, not just under-redaction. An Anthropic tool_result carries its own content, as a string or as further blocks, and only each part's `text` was being collected. Handled recursively; image and audio parts still fall through untouched. The legacy completions `suffix` is forwarded to providers that support it and was never collected. Note the placement: it has to be gathered before the string-prompt early return, which is what the new test pins. * fix(guardrails): walk nested tool results iteratively, with a depth bound CI flagged _collect_content as recursive. It was, and worse, it was unbounded: a tool_result nests its own content, the nesting is caller controlled, and the descent had nothing to stop it. That is a JSON bomb, not a style issue. Now an explicit queue with a depth bound of 8. Real payloads nest one or two deep. The queue is walked in document order because the shield maps its replies back by position, so collection order is part of the contract. * fix(guardrails): redact Responses PromptObject variables A Responses request can send `prompt` as a PromptObject rather than a string. Its `variables` are substituted into the stored prompt on the provider side, so they are caller text, and the dict shape was falling through untouched. `id` and `version` pick which stored prompt to run and are left unchanged. * test(guardrails): assert the depth bound instead of only reaching the end The depth test asserted nothing, so it passed whether or not the bound held, and the test-quality gate counted it as a zero-assert test. It now sends a shallow value alongside a 200-deep chain and asserts the shallow one is collected while the value past the bound is not. * fix(guardrails): keep system-prompt values out of the restored reply Redaction put every span of a request into one vault, and the reply was restored against that same vault. System prompts are written by the application and the caller never sees them, so a caller who got the model to echo a placeholder back had its plaintext restored into their own reply -- a way to read a system prompt they were never shown. Server-authored spans now go into a vault of their own: system and developer turns, Anthropic's top-level `system`, and the Responses API `instructions`. Its id is deliberately never stored, so nothing restores against it. The reply is restored against the caller's vault alone, and an echoed placeholder from a system prompt comes back as the placeholder. Values the caller also wrote themselves are unaffected -- they are in the caller's vault too, and still restore. The extra round trip happens only when a request actually carries server-authored text. * style(guardrails): satisfy ruff format and annotate the new tests `ruff format` wanted the widened `_collect_responses_fields` signature on one line, and the three tests added with the split-vault fix needed return annotations to keep ANN201 level with the base. * fix(guardrails): restore tool calls in the LLM Shield guardrail The request walk redacted a tool call's `arguments` -- plus the legacy `function_call`, Anthropic `tool_use.input` leaves and the Responses API's `function_call` / `function_call_output` fields -- while the response walk restored only `message.content`. A placeholder therefore reached the caller inside a tool call, and nothing raised. This is the same change as the out-of-tree example adapter this file is copied from, kept body-identical on purpose: the response side now collects every restorable span in one positional rehydrate batch, streaming keeps a window per (choice index, tool-call index) and flushes each into the chunk carrying the finish_reason, and `apply_guardrail` restores `inputs["tool_calls"]` on the response side. The declared limit on restoring values inside a JSON string is documented in the module. * fix(guardrails): import copy, keep the vault id off the provider, drop recursion Three defects Greptile and veria-ai found on the reopened PR, all real: - `copy.deepcopy` was called in `apply_guardrail` with no `import copy`, a guaranteed NameError on every response carrying tool calls. It landed on 2026-09-13, ten days after the review that rated this branch safe, and no test reached it: every tool-call test covered the request side. Adds the import and a regression test on the response side. - The vault session id was stored in `metadata`, which is forwarded to the provider on /v1/responses. A provider holding the placeholders and the session id can call the shield's rehydrate endpoint and read back the plaintext this guardrail exists to withhold. Moves it to `litellm_metadata`, which is not forwarded, and reads it back from there only. - `_collect_json_leaves` recursed over model-controlled JSON; the repo's recursive_detector gate rejects that. Rewritten with an explicit stack, same depth bound. 52 tests pass. ruff format, ruff-strict and check_type_discipline all clean, with LIT counts identical to the merge base. * fix(guardrails): build llm_shield_proxy stream deltas without new mutable literals The lint job's LIT002 budget gate failed on this PR: the file added 11 mutable-collection constructions and the tree sits at its limit. Build the index-only tool-call continuation in one helper, keep read-only inputs as tuples, and annotate the lists the delta and texts fields require. Adds tests for the two tool-call flush paths the refactor touches, which had no coverage: held arguments landing in the finish_reason chunk next to that chunk's own fragment, and the trailing flush of a stream that ends without a finish_reason. * fix(guardrails): drop Final from loop-body locals in llm_shield_proxy basedpyright rejects Final on a name assigned inside a loop, and the eleven such locals put reportGeneralTypeIssues over its budget (112/101). The LIT010 Final rule already exempts loop-body assignments, so the annotations go. * feat(guardrails): restore llm_shield_proxy placeholders on native streams Anthropic /v1/messages and /v1/responses streams have no `choices`, so the streaming hook passed them through with placeholders still in them. Both are now restored incrementally, with the same per-stream windows as chat: - /v1/messages arrives as raw SSE. Frames are cut at event boundaries, text_delta and input_json_delta are restored per block index, and held text is emitted as one more delta ahead of content_block_stop. Signed thinking deltas, frames from other endpoints and non-SSE raw streams pass through unchanged. - /v1/responses events are restored per item and part. Held text goes out as a copy of the stream's last delta before its .done event, and the events that repeat the reply (.done, content_part.done, output_item.done, response.completed) are restored in full. The request side now also redacts Anthropic tool_use inputs and Responses reasoning summaries, and sends tool and function descriptions (including parameter schema descriptions) and the user / safety_identifier fields to the non-restorable vault, like system prompts. Tool results stay restorable: the model reads them to answer, so restoring them returns what the caller would have seen without the guardrail. * fix(guardrails): redact llm_shield_proxy predicted outputs and output schemas `prediction.content` is the caller's own draft of the reply, so it is redacted into the caller vault and restored with the reply. The descriptions in a structured-output schema (Chat response_format.json_schema, Responses text.format) are application authored like tool schemas, so they go to the non-restorable vault. * fix(guardrails): fail closed on deep llm_shield_proxy requests, widen coverage - Request walks no longer skip what lies past their depth bound. Content nested past it, and tool inputs or schemas past the new JSON bound, now block the request instead of reaching the provider unredacted. The old depth test asserted the skip; it now asserts the block. - Tool and output schemas are walked by their JSON Schema structure, and give up `title`, `examples` and `default` as well as `description`. `enum` and `const` still go out as sent. - Responses events are matched by shape: any `*.delta` with a string delta is a token stream, and any `*.done` restores every non-identifier text field plus the `part` or `item` it repeats. This covers reasoning_summary_part.done and MCP arguments, and future families. Audio deltas are left alone. - An SSE stream whose first chunk ends partway through a field name (`b"eve"`) is no longer taken for a non-SSE stream. * fix(guardrails): scan llm_shield_proxy schemas by default The schema walk collected an allowlist of keywords, so any keyword it did not list -- draft-07 `dependencies`, `$comment`, vendor `x-` extensions -- went to the provider in clear. Invert it: every string is collected except under keywords whose value must go out verbatim (types, formats, patterns, references, required lists, enum, const). Name -> subschema maps still treat their keys as property names, so a property called `type` is walked, not skipped. * fix(guardrails): redact llm_shield_proxy schema enum and const values `enum` and `const` were skipped by the schema walk, so a value holding PII went to the provider in clear. They now go to the caller's vault rather than the non-restorable one: the model emits the stand-in in its tool arguments or structured output, and restoring the reply turns it back into the value the schema allows, so the call still routes. * fix(guardrails): redact llm_shield_proxy web search user locations Web search forwards the user's approximate location, and its free-text `city` and `region` fields can hold an address. Collect them into the non-restorable vault, from Chat `web_search_options.user_location` and from the `user_location` of Responses and Anthropic web-search tools. * fix(guardrails): drop unused llm_shield_proxy suppressions Upstream added LIT013 (a *-ok marker that suppresses nothing) and LIT014 (at most one for and one if per comprehension). Remove the 34 markers that no longer suppress anything and flatten the finished streams with itertools.chain.from_iterable. * fix(guardrails): type the llm_shield_proxy request and reply walks Narrowing with isinstance(x, dict) leaves keys and values unknown, so every call that passed a narrowed value counted against the reportUnknownArgumentType budget. Parse into dict[str, object] and list[object] once, in _as_object and _as_array, type the carry keys and accumulators, and bind writers with functools.partial instead of lambdas. The shield's batch reply is now also checked to hold only strings. * fix(guardrails): keep restored llm_shield_proxy replies out of the cache, widen coverage Addresses the open veria-ai and Cursor Bugbot findings on #42645. - Restore a copy of the reply and of each stream chunk, never LiteLLM's own object. LiteLLM caches and logs that object, and placeholders are numbered per request, so two callers' redacted requests can share a cache key: restoring in place cached one caller's plaintext for the next. The deployment hook no longer restores either, since LiteLLM caches what it returns; the proxy's post-call hook restores model-level guardrails after the cache write. - Restore /v1/completions replies, streamed and not, which carry `choice.text`. - Redact Responses replay fields the reply side already restores: tool output sent as input_text parts, custom_tool_call `input`, code_interpreter_call `code`. - Redact typed Responses prompt variables (`{"type": "input_text", "text": ...}`). - Put Responses system and developer input items in the non-restorable vault, like their Chat counterparts. - Expose LLMShieldProxyGuardrailConfigModel through get_config_model, so the dashboard can collect the Shield URL and key. * fix(guardrails): redact llm_shield_proxy plain-text document blocks An Anthropic document block carries text inline, in a text source's `data` or a content source's `content`, and that text reached the provider unredacted. Collect both, plus the block's `title` and `context`; base64, URL and file sources pass untouched. * fix(guardrails): redact llm_shield_proxy extra_body overrides LiteLLM merges extra_body over the transformed request just before sending, so text placed there (input, messages, system, ...) replaced the redacted field on the wire. Walk extra_body with the same collectors as the request, keeping the caller / application split. * test(guardrails): import InMemoryCache directly in the llm_shield_proxy cache test litellm keeps a deprecated module-level `caching` bool, so `litellm.caching.caching` resolves to that bool once an earlier test in the same worker has set it, and the test failed with AttributeError depending on test order. * fix(guardrails): restore llm_shield_proxy replies for model-level use outside the proxy 71e68fd stopped the deployment post-call hook from restoring, so the response cache never holds restored plaintext. Inside the proxy that is right: the proxy's post-call hook restores after the cache write. But with model-level `guardrails` on the SDK, the deployment hooks are the only redact and restore steps, so callers got placeholders back. When the deployment pre-call hook is the one that redacts, it now records the request's vault id and marks the request no-cache / no-store; the deployment post-call hook restores only when that record matches. The cache key there is built from the redacted request and a cache hit skips the post-call hook, so a cached reply could neither be restored nor safely shared. Proxy requests carry no record and keep restoring in the proxy's post-call hook, after the cache write. * fix(guardrails): don't repeat usage in llm_shield_proxy end-of-stream flush chunks With n>=2 and stream_options.include_usage, the end-of-stream flush copies the last chunk the stream carried, which is the one holding usage, so each synthetic flush chunk repeated it and a consumer summing usage chunks counted the request twice. The copy now drops `usage`, matching a normal mid-stream chunk. Reported by @yucheng-berri. * fix(guardrails): keep restored llm_shield_proxy values out of telemetry, refuse SDK streams - The post-call restore hook no longer goes through log_guardrail_information, which recorded its whole return value, the restored reply, as guardrail_response. That field is exported to traces even with message logging turned off. - A model-level stream outside the proxy is refused once redacted. Nothing restores an SDK stream, and its cache writer reads the request from before the deployment hook, so it also got cached despite the no-store bypass. - Drop a narrating comment, and keep example_config.yaml to config only; the how-to lives in the docs PR. * refactor(guardrails): split llm_shield_proxy into payload, request walk and stream modules The module had grown past 1,700 lines. Shared payload types and helpers move to payload.py, the request walk to request_walk.py and the stream restorers to stream_restorers.py; llm_shield_proxy.py keeps the guardrail class. No behaviour change. * style(guardrails): drop routine comments from llm_shield_proxy AGENTS.md keeps source comments to tool directives and genuinely complex logic; the rationale stays in the docstrings. --- .../llm_shield_proxy/__init__.py | 33 + .../llm_shield_proxy/example_config.yaml | 14 + .../llm_shield_proxy/llm_shield_proxy.py | 744 ++++++ .../llm_shield_proxy/payload.py | 190 ++ .../llm_shield_proxy/request_walk.py | 363 +++ .../llm_shield_proxy/stream_restorers.py | 345 +++ litellm/types/guardrails.py | 1 + .../guardrail_hooks/llm_shield_proxy.py | 24 + ruff-strict.toml | 4 + .../guardrail_hooks/test_llm_shield_proxy.py | 2003 +++++++++++++++++ .../public/assets/logos/llm_shield_proxy.svg | 6 + .../_components/add_guardrail_form.tsx | 4 +- .../_components/guardrail_garden_configs.ts | 12 +- .../_components/guardrail_garden_data.test.ts | 1 + .../_components/guardrail_garden_data.ts | 10 + .../_components/guardrail_info_helpers.tsx | 3 + 16 files changed, 3755 insertions(+), 2 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/llm_shield_proxy.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py create mode 100644 ui/litellm-dashboard/public/assets/logos/llm_shield_proxy.svg diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/__init__.py new file mode 100644 index 00000000000..44b19f82218 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/__init__.py @@ -0,0 +1,33 @@ +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .llm_shield_proxy import LLMShieldProxyGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> LLMShieldProxyGuardrail: + import litellm + + _llm_shield_guardrail_callback: Final = LLMShieldProxyGuardrail( + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(_llm_shield_guardrail_callback) + return _llm_shield_guardrail_callback + + +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.LLM_SHIELD_PROXY.value: initialize_guardrail, +} + + +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.LLM_SHIELD_PROXY.value: LLMShieldProxyGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml new file mode 100644 index 00000000000..4f732773a75 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml @@ -0,0 +1,14 @@ +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "llm_shield_proxy" + litellm_params: + guardrail: llm_shield_proxy + mode: ["pre_call", "post_call"] + default_on: true + api_base: "http://localhost:8000" + api_key: os.environ/LLM_SHIELD_PROXY_API_KEY diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py new file mode 100644 index 00000000000..fa23a110a74 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -0,0 +1,744 @@ +import copy +import functools +import os +import uuid +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from typing import ( + TYPE_CHECKING, + Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__ + ClassVar, + Final, + Literal, + Optional, +) + +import httpx + +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.llm_shield_proxy import ( + LLMShieldProxyGuardrailConfigModel, + ) + from litellm.types.utils import CallTypes, LLMResponseTypes +from .payload import ( + JsonBody, + MutableRequest, + RequestTooDeep, + Slot, + SlotSink, + as_array, + as_object, + choice_index, + collect_json_leaves, + collect_response_item, + detached, + read_field, + read_list, + rehydrate_slots, + write_field, +) +from .request_walk import ( + locate_request_texts, +) +from .stream_restorers import ( + AnthropicSSERestorer, + CarryKey, + CarryWindows, + ResponsesStreamRestorer, + carry_sort_key, + continuation_delta, + responses_event_type, +) + +GUARDRAIL_NAME: Final = "llm_shield_proxy" + +_DEFAULT_API_BASE: Final = "http://localhost:8000" +_REDACT_PATH: Final = "/v1/guard/redact" +_REHYDRATE_PATH: Final = "/v1/guard/rehydrate" +_REHYDRATE_STREAM_PATH: Final = "/v1/guard/rehydrate/stream" + +_SESSION_METADATA_KEY: Final = "llm_shield_session_id" + +_DEPLOYMENT_RESTORE_KEY: Final = "llm_shield_restore_at_deployment" + +_VAULT_PREFIX: Final = f"litellm-{uuid.uuid4().hex}" + +_DEFAULT_TIMEOUT_SECONDS: Final = 10.0 + + +class LLMShieldProxyGuardrail(CustomGuardrail): + """Redacts PII before it leaves the proxy and restores it in the response. + + Unlike a masking guardrail, the substitution is reversible. Outbound text is + replaced with placeholders held in a session vault inside the user's own LLM + Shield deployment; the model's reply is then restored so the end user sees the + original values while the provider never received them. + + Streaming is restored incrementally rather than by buffering the response. LLM + Shield holds back only the trailing characters that could still turn out to be + part of a placeholder, so tokens are forwarded as they arrive and a placeholder + split across two chunks is never emitted in fragments. + """ + + use_native_lifecycle_hooks: ClassVar[bool] = True + + def __init__( + self, + guardrail_name: str = GUARDRAIL_NAME, + api_base: str | None = None, + api_key: str | None = None, + **kwargs: Any, # noqa: LIT008 # kwargs-ok: forwarded verbatim to CustomGuardrail.__init__ + ) -> None: + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + env_base: Final = os.environ.get("LLM_SHIELD_PROXY_API_BASE") + self.api_base: Final = (api_base or env_base or _DEFAULT_API_BASE).rstrip("/") + self.api_key: Final = api_key or os.environ.get("LLM_SHIELD_PROXY_API_KEY") + super().__init__(guardrail_name=guardrail_name, **kwargs) + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: parent's signature. + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + + @staticmethod + def get_config_model() -> type["LLMShieldProxyGuardrailConfigModel"]: + from litellm.types.proxy.guardrails.guardrail_hooks.llm_shield_proxy import ( + LLMShieldProxyGuardrailConfigModel, + ) + + return LLMShieldProxyGuardrailConfigModel + + async def async_pre_call_deployment_hook( + self, + kwargs: MutableRequest, + call_type: "CallTypes | None", + ) -> MutableRequest | None: + """Redacts a model-level guardrail's request, and keeps it out of the response cache. + + Outside the proxy this hook is the only redaction step, and the deployment post-call + hook the only restoration step. LiteLLM builds the cache key after this hook, from + the redacted request, and a cache hit returns before the post-call hook runs. So a + cached reply would either reach the caller unrestored or, stored after restoration, + hand this caller's values to the next caller whose redacted request matches. The + request is therefore neither read from nor written to the cache. Inside the proxy + this hook does not redact -- the proxy's pre-call hook already ran -- and caching + is left alone, because the proxy restores after the cache write. + + A streamed request is refused once redacted. No hook restores an SDK stream, and the + stream's cache writer reads the request from before this hook, so it would also be + cached despite the bypass. + """ + before: Final = self._minted_session_id(kwargs) + _ = await super().async_pre_call_deployment_hook(kwargs, call_type) + session_id: Final = self._minted_session_id(kwargs) + if session_id is None or session_id == before: + return kwargs + if kwargs.get("stream") is True: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=( + "LLM Shield Proxy cannot restore a streamed reply for a model-level guardrail " + "outside the LiteLLM proxy; send the request through the proxy or without stream=True." + ), + ) + metadata: Final = as_object(kwargs.get("litellm_metadata")) + if metadata is not None: + metadata[_DEPLOYMENT_RESTORE_KEY] = session_id + cache_controls: Final = as_object(kwargs.get("cache")) + kwargs["cache"] = {**(cache_controls or {}), "no-cache": True, "no-store": True} + return kwargs + + async def async_post_call_success_deployment_hook( + self, + request_data: MutableRequest, + response: "LLMResponseTypes", + call_type: "CallTypes | None", + ) -> "LLMResponseTypes | None": + """Restores the reply here only when the deployment pre-call hook redacted it. + + LiteLLM caches what this hook returns. Inside the proxy the request was redacted by + the proxy's pre-call hook and the proxy's post-call hook restores the reply after + the cache write, so restoring here as well would cache this caller's plaintext under + a key built from the redacted request. Outside the proxy nothing restores later, and + the pre-call deployment hook has already kept that request out of the cache. + """ + metadata: Final = as_object(request_data.get("litellm_metadata")) + marker: Final = metadata.get(_DEPLOYMENT_RESTORE_KEY) if metadata is not None else None + session_id: Final = self._minted_session_id(request_data) + if session_id is None or marker != session_id: + return None + return await super().async_post_call_success_deployment_hook(request_data, response, call_type) + + def _headers(self, session_id: str) -> JsonBody: + headers: Final[JsonBody] = { + "Content-Type": "application/json", + "X-Session-ID": session_id, + } + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + return headers + + async def _call_shield(self, path: str, session_id: str, payload: JsonBody) -> Mapping[str, object]: + """Posts to LLM Shield Proxy, failing closed on any transport or status error. + + A redaction guardrail that fails open sends the very data it exists to + protect to a third-party provider, so an unreachable or erroring shield + blocks the request instead of passing it through. + """ + try: + response: Final = await self.async_handler.post( + f"{self.api_base}{path}", + headers=self._headers(session_id), + json=payload, + timeout=_DEFAULT_TIMEOUT_SECONDS, + ) + response.raise_for_status() + return response.json() + except httpx.HTTPStatusError as exc: + verbose_proxy_logger.exception("LLM Shield Proxy returned %s for %s", exc.response.status_code, path) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"LLM Shield Proxy returned {exc.response.status_code}; blocking the request.", + ) from exc + except Exception as exc: + verbose_proxy_logger.exception("LLM Shield Proxy call to %s failed", path) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message="LLM Shield Proxy is unreachable; blocking the request.", + ) from exc + + async def _redact(self, texts: Sequence[str], session_id: str) -> Sequence[str]: + payload: Final[JsonBody] = {"texts": list(texts)} + body: Final = await self._call_shield(_REDACT_PATH, session_id, payload) + return self._same_length_or_raise(body.get("texts"), texts, "redact") + + async def _rehydrate(self, texts: Sequence[str], session_id: str) -> Sequence[str]: + payload: Final[JsonBody] = {"texts": list(texts)} + body: Final = await self._call_shield(_REHYDRATE_PATH, session_id, payload) + return self._same_length_or_raise(body.get("texts"), texts, "rehydrate") + + def _same_length_or_raise(self, returned: object, sent: Sequence[str], operation: str) -> Sequence[str]: + """Guards the positional mapping the callers rely on to write results back.""" + entries: Final = as_array(returned) + texts: Final = tuple(entry for entry in entries or () if isinstance(entry, str)) + if entries is None or len(entries) != len(sent) or len(texts) != len(entries): + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"LLM Shield Proxy {operation} returned an unexpected payload; blocking the request.", + ) + return texts + + @staticmethod + def _mint_session_id(data: MutableRequest) -> str: + """Mints a vault id for this request, overwriting anything already there. + + Redaction and restoration both happen inside one request/response pair, so + a fresh id per request is all that is needed, and it is what keeps one + caller from reaching another caller's vault. + """ + session_id: Final = f"{_VAULT_PREFIX}-{uuid.uuid4().hex}" + metadata: Final = data.setdefault("litellm_metadata", {}) + if isinstance(metadata, dict): + metadata[_SESSION_METADATA_KEY] = session_id + return session_id + + @staticmethod + def _minted_session_id(data: MutableRequest) -> str | None: + """The vault id this process minted for `data`, or None if it has none. + + Read only from `litellm_metadata`, the proxy-private store `_mint_session_id` writes + to. A caller can populate `metadata`; they cannot populate this. + """ + metadata: Final = as_object(data.get("litellm_metadata")) + existing: Final = metadata.get(_SESSION_METADATA_KEY) if metadata is not None else None + return existing if isinstance(existing, str) and existing.startswith(_VAULT_PREFIX) else None + + @staticmethod + def _session_id(data: MutableRequest) -> str: + """Reads back the vault id minted while redacting this request. + + Falls back to an unused id rather than to anything the caller supplied: a + reply that cannot be restored is a visible placeholder, while trusting a + caller-supplied id would hand them someone else's plaintext. + """ + existing: Final = LLMShieldProxyGuardrail._minted_session_id(data) + return existing if existing is not None else f"{_VAULT_PREFIX}-{uuid.uuid4().hex}" + + @staticmethod + def _locate_request_texts(data: MutableRequest) -> tuple[Sequence[Slot], Sequence[Slot]]: + return locate_request_texts(data) + + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: "DualCache", + data: MutableRequest, + call_type: str, + ) -> MutableRequest | None: + """Replaces PII anywhere in the outbound request with vault placeholders.""" + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is not True: + return data + + try: + slots, privileged = self._locate_request_texts(data) + except RequestTooDeep as exc: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Request {exc} nests deeper than LLM Shield Proxy inspects; blocking the request.", + ) from exc + if not slots and not privileged: + return data + + session_id: Final = self._mint_session_id(data) + if privileged: + await self._redact_into(privileged, f"{_VAULT_PREFIX}-{uuid.uuid4().hex}") + if slots: + await self._redact_into(slots, session_id) + return data + + async def _redact_into(self, slots: Sequence[Slot], session_id: str) -> None: + """Redacts every span in `slots` under one vault and writes the result back.""" + redacted: Final = await self._redact(tuple(text for text, _ in slots), session_id) + for (_, write), replacement in zip(slots, redacted): + write(replacement) + + async def async_post_call_success_hook( + self, + data: MutableRequest, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + """Restores the original values in a copy of a non-streaming response. + + The copy is what keeps plaintext out of the response cache. LiteLLM caches the + reply it received from the provider, and on some paths -- a native Anthropic dict, + an in-memory cache -- it stores the object itself rather than a serialised + snapshot. Restoring that object in place would cache this caller's values under a + key built from the redacted request, which another caller's identical-looking + request then hits. Left untouched, the cached reply holds placeholders, and a hit + is restored against the new caller's own vault. + """ + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: + return response + response = detached(response) # rebind-ok: everything below restores the copy. + + if self._is_anthropic_message_response(response): + return await self._restore_anthropic_response(response, data) + + response_slots: Final = self._responses_api_slots(response) + if response_slots: + return await self._restore_responses_api_response(response, response_slots, data) + + choices: Final = getattr(response, "choices", None) + if not choices: + return response + + pending: Final[SlotSink] = [] + for choice in choices: + message = getattr(choice, "message", None) + if message is None: + text = read_field(choice, "text") + if isinstance(text, str) and text: + pending.append((text, functools.partial(write_field, choice, "text"))) + continue + content = getattr(message, "content", None) + if isinstance(content, str) and content: + pending.append((content, functools.partial(setattr, message, "content"))) + for tool_call in getattr(message, "tool_calls", None) or (): + function = getattr(tool_call, "function", None) + arguments = getattr(function, "arguments", None) if function is not None else None + if isinstance(arguments, str) and arguments: + pending.append((arguments, functools.partial(setattr, function, "arguments"))) + legacy = getattr(message, "function_call", None) + legacy_arguments = getattr(legacy, "arguments", None) if legacy is not None else None + if isinstance(legacy_arguments, str) and legacy_arguments: + pending.append((legacy_arguments, functools.partial(setattr, legacy, "arguments"))) + if not pending: + return response + + restored: Final = await self._rehydrate(tuple(text for text, _ in pending), self._session_id(data)) + for (_, write), replacement in zip(pending, restored): + write(replacement) + return response + + @staticmethod + def _is_anthropic_message_response(response: object) -> bool: + """Anthropic's native /v1/messages reply arrives as a plain dict.""" + body: Final = as_object(response) + return body is not None and body.get("type") == "message" and isinstance(body.get("content"), list) + + async def _restore_anthropic_response(self, response: MutableRequest, data: MutableRequest) -> MutableRequest: + """Restores text blocks and tool inputs in an Anthropic native message reply. + + This shape has no `choices`, so without its own branch the reply would go + back to the caller still carrying placeholders. + + A `tool_use` block's payload is `input`, an arbitrary JSON object rather than a + string, and the request path redacts its string leaves -- so the reply's leaves + have to come back or the application invokes the tool with placeholders. + """ + slots: Final[SlotSink] = [] + for entry in read_list(response, "content"): + block = as_object(entry) + if block is None: + continue + kind = block.get("type") + text = block.get("text") + if kind == "text" and isinstance(text, str) and text: + slots.append((text, functools.partial(block.__setitem__, "text"))) + elif kind == "tool_use" and as_object(block.get("input")) is not None: + collect_json_leaves(block.get("input"), slots) + if not slots: + return response + + restored: Final = await self._rehydrate(tuple(text for text, _ in slots), self._session_id(data)) + for (_, write), replacement in zip(slots, restored): + write(replacement) + return response + + @staticmethod + def _responses_api_slots(response: object) -> Sequence[Slot]: + """Restorable spans in a Responses API reply. + + That shape carries `output` items rather than `choices`, so it needs its own + walk; without one the reply goes back to the caller still holding + placeholders even though the request was redacted correctly. Items and blocks + come through as dicts or as objects depending on how far the reply has been + deserialised, so both are handled. + + The item-level fields mirror `collect_responses_fields`, which walks the same + fields on the request side -- a function_call item holds `arguments`, a + function_call_output holds `output` -- so the two directions stay symmetric. + """ + slots: Final[SlotSink] = [] + for item in getattr(response, "output", None) or (): + collect_response_item(item, slots) + return tuple(slots) + + async def _restore_responses_api_response(self, response: Any, slots: Sequence[Slot], data: MutableRequest) -> Any: + """Puts the original values back into a Responses API reply.""" + await rehydrate_slots(slots, functools.partial(self._rehydrate, session_id=self._session_id(data))) + return response + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: MutableRequest, + ) -> AsyncGenerator[Any, None]: + """Restores original values incrementally, without buffering the stream. + + Each choice -- and each tool call within a choice -- is its own token stream, so + the sliding window is tracked per (choice index, tool call) pair. One shared + window would splice the characters held back for one stream onto another. The + windows are locals of this generator, so they are scoped to a single stream and + cannot leak between concurrent requests. + + The two native stream shapes have no `choices` and are restored by their own + walkers, with the same per-stream windows: Anthropic `/v1/messages` arrives as raw + SSE frames, and the Responses API as typed events. + """ + if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: + async for chunk in response: + yield chunk + return + + session_id: Final = self._session_id(request_data) + step: Final = functools.partial(self._stream_step, session_id=session_id) + rehydrate: Final = functools.partial(self._rehydrate, session_id=session_id) + sse: Final = AnthropicSSERestorer(step) + events: Final = ResponsesStreamRestorer(step, rehydrate) + carries: Final[CarryWindows] = {} + last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush. + + async for chunk in response: + if isinstance(chunk, (bytes, str)): + for frames in await sse.feed(chunk): + yield frames + continue + if responses_event_type(chunk) is not None: + for event in await events.restore(detached(chunk)): + yield event + continue + restored_chunk = detached(chunk) + last_chunk = restored_chunk + for choice in getattr(restored_chunk, "choices", None) or (): + await self._restore_choice(choice, carries, session_id) + yield restored_chunk + + for frames in await sse.finish(): + yield frames + for event in await events.finish(): + yield event + if last_chunk is not None and any(carries.values()): + async for trailing in self._flush_trailing(last_chunk, carries, session_id): + yield trailing + + async def _restore_choice(self, choice: object, carries: CarryWindows, session_id: str) -> None: + """Restores one choice's delta, advancing that choice's own windows. + + Content and each tool call are separate token streams, so each gets its own + window: `(choice_index, None)` for content, `(choice_index, tool_call_index)` for + one tool call's accumulating `arguments`. A shared window would splice the text + held back for one stream onto another. + """ + delta: Final = getattr(choice, "delta", None) + index: Final = choice_index(choice) + is_final: Final = bool(getattr(choice, "finish_reason", None)) + if isinstance(choice, TextChoices): + await self._restore_text_window(choice, (index, None), carries, session_id, is_final) + return + if delta is None: + return + + await self._restore_content_window(delta, (index, None), carries, session_id, is_final) + + for tool_call in getattr(delta, "tool_calls", None) or (): + await self._restore_tool_call_window(tool_call, index, carries, session_id) + + if is_final: + await self._flush_finished_choice(delta, index, carries, session_id) + + async def _restore_text_window( + self, + choice: object, + key: CarryKey, + carries: CarryWindows, + session_id: str, + is_final: bool, + ) -> None: + """Restores a Completions stream choice's `text` through its window.""" + carry: Final = carries.get(key, "") + text: Final = read_field(choice, "text") + if not isinstance(text, str) or not text: + if not (is_final and carry): + return + emitted, remaining = await self._stream_step(text if isinstance(text, str) else "", carry, is_final, session_id) + carries[key] = remaining # rebind-ok: this stream's window advances. + if emitted or text: + write_field(choice, "text", emitted) + + async def _restore_content_window( + self, + delta: Any, + key: CarryKey, + carries: CarryWindows, + session_id: str, + is_final: bool, + ) -> None: + """Restores one delta's content through its own window.""" + carry: Final = carries.get(key, "") + text: Final = getattr(delta, "content", None) + + if not isinstance(text, str) or not text: + if is_final and carry: + flushed, remaining = await self._stream_step("", carry, True, session_id) + carries[key] = remaining # rebind-ok: this stream's window advances. + if flushed: + delta.content = flushed + return + + emitted, remaining = await self._stream_step(text, carry, is_final, session_id) + carries[key] = remaining # rebind-ok: this stream's window advances. + delta.content = emitted + + async def _restore_tool_call_window( + self, + tool_call: Any, + choice_index: int, + carries: CarryWindows, + session_id: str, + ) -> None: + """Restores one streamed tool call's argument fragment. + + A tool call's `arguments` is a JSON document delivered as fragments that clients + concatenate per tool-call index, so each index gets a window of its own rather + than sharing the content stream's. + """ + tool_index: Final = read_field(tool_call, "index") + if not isinstance(tool_index, int): + return + function: Final = read_field(tool_call, "function") + if function is None: + return + arguments: Final = read_field(function, "arguments") + if not isinstance(arguments, str) or not arguments: + return + + key: Final = (choice_index, tool_index) + emitted, remaining = await self._stream_step(arguments, carries.get(key, ""), False, session_id) + carries[key] = remaining # rebind-ok: this tool call's window advances. + write_field(function, "arguments", emitted) + + async def _flush_finished_choice( + self, + delta: Any, + choice_index: int, + carries: CarryWindows, + session_id: str, + ) -> None: + """Emits everything this finishing choice still holds, into this chunk. + + A client parses a tool call's `arguments` when the chunk carrying the + finish_reason arrives, so a flush delivered afterwards is too late -- the client + has already tried to parse truncated JSON. Content lands back on `content`; held + tool-call text is appended as an index-only continuation entry, which is the shape + clients concatenate by index, so no id or name is needed. Appending is correct + even when this chunk already carried a fragment for that tool call. + """ + continuations: Final[list[dict[str, object]]] = [] # mutable-ok: built into this chunk's delta. + for key in sorted((held for held in carries if held[0] == choice_index), key=carry_sort_key): + carry = carries[key] + if not carry: + continue + _, tool_index = key + text, remaining = await self._stream_step("", carry, True, session_id) + carries[key] = remaining # rebind-ok: this stream's window advances. + if not text: + continue + if tool_index is None: + delta.content = text + else: + continuations.extend(continuation_delta(tool_index, text)) + if continuations: + existing: Final = tuple(getattr(delta, "tool_calls", None) or ()) + delta.tool_calls = [*existing, *continuations] + + async def _flush_trailing( + self, last_chunk: Any, carries: CarryWindows, session_id: str + ) -> AsyncGenerator[Any, None]: + """Empties every window still holding text, one chunk per window. + + This is the net for a stream that ended with no finish_reason at all; a stream + that ended with one is flushed into its own terminal chunk by + `_flush_finished_choice`, because that is the moment a client parses tool + arguments. + + Driven by the windows rather than by the last chunk's choices. A choice that + finished earlier is not present in the terminal chunk, and flushing only what + that chunk carries would drop its held text and truncate its answer. + """ + for key in sorted(carries, key=carry_sort_key): + carry = carries[key] + if not carry: + continue + choice_index, tool_index = key + text, remaining = await self._stream_step("", carry, True, session_id) + carries[key] = remaining # rebind-ok: this stream's window advances. + if not text: + continue + chunk = self._chunk_for_choice(last_chunk, choice_index) + if chunk is None: + continue + choice: object = chunk.choices[0] + delta = read_field(choice, "delta") + if isinstance(choice, TextChoices): + choice.text = text + elif tool_index is None: + write_field(delta, "content", text) + else: + write_field(delta, "content", None) + write_field(delta, "tool_calls", continuation_delta(tool_index, text)) + yield chunk + + @staticmethod + def _chunk_for_choice(last_chunk: Any, index: int) -> Any: + """A single-choice copy of the last chunk, carrying only `index`. + + Emitting one choice per chunk keeps a flush from reading as content on a + choice it does not belong to. + """ + chunk: Final = last_chunk.model_copy(deep=True) + raw_choices: Final = getattr(chunk, "choices", None) + if not raw_choices: + return None + choices: Final[tuple[object, ...]] = tuple(raw_choices) + position: Final = next((at for at, choice in enumerate(choices) if choice_index(choice) == index), 0) + kept: Final = raw_choices[position] + if getattr(kept, "delta", None) is None and not isinstance(kept, TextChoices): + return None + kept.index = index + kept.finish_reason = None + chunk.choices = [kept] + if hasattr(chunk, "usage"): + del chunk.usage + return chunk + + async def _stream_step(self, text: str, carry: str, final: bool, session_id: str) -> tuple[str, str]: + """Returns ``(text safe to emit now, window still being held)``.""" + body: Final = await self._call_shield( + _REHYDRATE_STREAM_PATH, + session_id, + {"text": text, "carry": carry, "final": final}, + ) + emitted: Final = body.get("text") + remaining: Final = body.get("carry") + if not isinstance(emitted, str) or not isinstance(remaining, str): + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message="LLM Shield Proxy stream rehydration returned an unexpected payload.", + ) + return emitted, remaining + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: MutableRequest, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """Unified entry point: what the UI's Test guardrail button and the translation + handlers call. + + `tool_calls` is handled on the response side only. LiteLLM populates the field here, + and on a reply it holds the model's tool arguments -- the same text the native hook + restores, and restoring one but not the other would leave the placeholder on + whichever path ran. The request side is left to the native pre-call hook, because + redacting it here as well would redact it twice. + """ + text_list: Final = tuple(inputs.get("texts") or ()) + tool_calls: Final = tuple(inputs.get("tool_calls") or ()) if input_type == "response" else () + if not text_list and not tool_calls: + return inputs + + restored_calls: Final[list[object]] = [copy.deepcopy(call) for call in tool_calls] # mutable-ok: a new list. + spans: Final[list[str]] = list(text_list) # mutable-ok: ordered batch, frozen before the call. + writers: Final[list[Callable[[str], None]]] = [] # mutable-ok: one per span appended below. + for call in restored_calls: + function = read_field(call, "function") + arguments = read_field(function, "arguments") if function is not None else None + if isinstance(arguments, str) and arguments: + spans.append(arguments) + writers.append(functools.partial(write_field, function, "arguments")) + + replaced: Final = ( + await self._redact(tuple(spans), self._mint_session_id(request_data)) + if input_type == "request" + else await self._rehydrate(tuple(spans), self._session_id(request_data)) + ) + restored_values: Final[list[str]] = list(replaced) # mutable-ok: sliced into the texts list. + + for write, replacement in zip(writers, restored_values[len(text_list) :]): + write(replacement) + merged: Final[JsonBody] = {**inputs} + if text_list: + merged["texts"] = restored_values[: len(text_list)] + if restored_calls: + merged["tool_calls"] = restored_calls + return merged diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py new file mode 100644 index 00000000000..644ac76efb8 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py @@ -0,0 +1,190 @@ +import copy +import functools +from collections.abc import Awaitable, Callable, Sequence +from typing import ( + Final, + TypeAlias, +) + +MutableRequest: TypeAlias = dict[str, object] + +JsonBody: TypeAlias = dict[str, object] + +MAX_JSON_DEPTH: Final = 64 + +Slot: TypeAlias = tuple[str, Callable[[str], None]] + +StreamStep: TypeAlias = Callable[[str, str, bool], Awaitable[tuple[str, str]]] + +Rehydrate: TypeAlias = Callable[[Sequence[str]], Awaitable[Sequence[str]]] + +SlotSink: TypeAlias = list[Slot] + +MutableSeq: TypeAlias = list[object] + + +def as_object(value: object) -> MutableRequest | None: + """`value` as a JSON object, or None. + + `isinstance(value, dict)` alone leaves the keys and values unknown to the type + checker. A JSON object's keys are strings, so the type is stated once, here. + """ + return value if isinstance(value, dict) else None + + +def as_array(value: object) -> MutableSeq | None: + """`value` as a JSON array, or None. See `as_object`.""" + return value if isinstance(value, list) else None + + +def detached(value: object) -> object: + """A deep copy of a reply or chunk, for restoring without touching LiteLLM's own object. + + LiteLLM keeps the object it handed the hooks to fill its response cache and its + logs, so writing restored plaintext into that object would put it there too. + """ + return copy.deepcopy(value) + + +def is_container(value: object) -> bool: + """Whether `value` is a JSON object or array, without narrowing it to unknown types.""" + return isinstance(value, (dict, list)) + + +def collect(container: MutableRequest, key: str, slots: SlotSink) -> None: + """Records the string at `key`, along with the write that replaces it.""" + value: Final = container.get(key) + if isinstance(value, str) and value: + slots.append((value, functools.partial(container.__setitem__, key))) + + +def collect_entry(entries: MutableSeq, index: int, slots: SlotSink) -> None: + """Records a string held directly in a list, rather than under a key.""" + value: Final = entries[index] + if isinstance(value, str) and value: + slots.append((value, functools.partial(entries.__setitem__, index))) + + +class RequestTooDeep(Exception): + """A request nests text past a walk's bound. + + Skipping the rest would forward it unredacted while the guardrail reports as + enabled, so the pre-call hook refuses the request instead. + """ + + +def collect_text_parts(container: MutableRequest, key: str, slots: SlotSink) -> None: + """Collects the `text` of every part in the list held at `key`.""" + for entry in as_array(container.get(key)) or (): + part = as_object(entry) + if part is not None: + collect(part, "text", slots) + + +def choice_index(choice: object) -> int: + """Streaming choices are matched across chunks by their index.""" + index: Final = getattr(choice, "index", 0) + return index if isinstance(index, int) else 0 + + +def read_field(holder: object, name: str) -> object: + """Reads one field from a dict or from an object. + + LiteLLM's replies arrive as Pydantic models on some paths and as plain dicts on + others, depending how far they have been deserialised, so every response walk here + has to handle both shapes. + """ + fields: Final = as_object(holder) + if fields is not None: + return fields.get(name) + return getattr(holder, name, None) + + +def read_list(holder: object, name: str) -> Sequence[object]: + """Reads a list field from a dict or an object; anything else reads as empty. + + The entries are the reply's own objects, so writing through them edits the reply. + """ + value: Final = read_field(holder, name) + if isinstance(value, tuple): + return value + return tuple(as_array(value) or ()) + + +def write_field(holder: object, name: str, value: object) -> None: + """Writes one string field back into a dict or an object. Pairs with read_field.""" + if isinstance(holder, dict): + holder[name] = value + else: + setattr(holder, name, value) + + +def collect_json_leaves(node: object, slots: SlotSink, *, strict: bool = False) -> None: + """Collects every string leaf of a JSON-ish structure, with a write-back per leaf. + + An Anthropic `tool_use` block carries `input`, an arbitrary JSON object rather than a + string, so a value worth restoring can sit at any depth. Bounded by `MAX_JSON_DEPTH`: + the shape is caller or model controlled, and the bound is what stops a crafted one from + becoming an unbounded descent. Walked with an explicit stack rather than recursively, + so a deeply nested value cannot spend stack frames proportional to attacker-chosen + depth. + + `strict` is for the request side, where a leaf left behind would reach the provider + unredacted: past the bound it raises `RequestTooDeep`. On the reply side a leaf past + the bound just keeps its placeholder, which leaks nothing, so it is skipped. + """ + pending: Final[list[tuple[object, int]]] = [(node, 0)] # mutable-ok: local walk stack. + while pending: + current, current_depth = pending.pop() + if current_depth > MAX_JSON_DEPTH: + if strict and is_container(current) and current: + raise RequestTooDeep("json") + continue + current_object = as_object(current) + if current_object is not None: + for key in tuple(current_object): + value = current_object[key] + if isinstance(value, str) and value: + slots.append((value, functools.partial(current_object.__setitem__, key))) + else: + pending.append((value, current_depth + 1)) + continue + entries = as_array(current) + if entries is not None: + for index, value in enumerate(entries): + if isinstance(value, str) and value: + slots.append((value, functools.partial(entries.__setitem__, index))) + else: + pending.append((value, current_depth + 1)) + + +def collect_response_item(item: object, slots: SlotSink) -> None: + """Restorable spans in one Responses API output item, dict or object. + + Mirrors `collect_responses_fields` on the request side -- a function_call or + mcp_call item holds `arguments`, their outputs `output`, a reasoning item `summary` + parts -- so the two directions stay symmetric. A custom tool call carries `input` and + a code interpreter call `code`, both model-written. + """ + for block in read_list(item, "content"): + for field in ("text", "refusal"): + text = read_field(block, field) + if isinstance(text, str) and text: + slots.append((text, lambda new, b=block, f=field: write_field(b, f, new))) + for part in read_list(item, "summary"): + text = read_field(part, "text") + if isinstance(text, str) and text: + slots.append((text, lambda new, p=part: write_field(p, "text", new))) + for field in ("arguments", "output", "input", "code"): + value = read_field(item, field) + if isinstance(value, str) and value: + slots.append((value, lambda new, i=item, f=field: write_field(i, f, new))) + + +async def rehydrate_slots(slots: Sequence[Slot], rehydrate: Rehydrate) -> None: + """Restores every span in `slots` in one batch and writes each result back.""" + if not slots: + return + restored: Final = await rehydrate(tuple(text for text, _ in slots)) + for (_, write), replacement in zip(slots, restored): + write(replacement) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py new file mode 100644 index 00000000000..c9c0d2ccc47 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py @@ -0,0 +1,363 @@ +from collections.abc import Sequence +from typing import ( + Final, +) + +from .payload import ( + MAX_JSON_DEPTH, + MutableRequest, + RequestTooDeep, + Slot, + SlotSink, + as_array, + as_object, + collect, + collect_entry, + collect_json_leaves, + collect_text_parts, + is_container, + read_list, +) + +PRIVILEGED_ROLES: Final = frozenset({"system", "developer"}) + +MAX_CONTENT_DEPTH: Final = 8 + +SCHEMA_STRUCTURAL_KEYWORDS: Final = frozenset( + ( + "type", + "format", + "pattern", + "required", + "dependentRequired", + "propertyOrdering", + "discriminator", + "contentEncoding", + "contentMediaType", + "$ref", + "$id", + "$schema", + "$anchor", + "$dynamicRef", + "$dynamicAnchor", + "$recursiveRef", + "$recursiveAnchor", + "$vocabulary", + ) +) + +SCHEMA_VALUE_KEYWORDS: Final = frozenset(("examples", "default")) + +SCHEMA_LITERAL_KEYWORDS: Final = frozenset(("enum", "const")) + +SCHEMA_MAP_KEYWORDS: Final = frozenset( + ("properties", "patternProperties", "$defs", "definitions", "dependentSchemas", "dependencies") +) + + +def collect_prompt(data: MutableRequest, slots: SlotSink) -> None: + """The Completions API sends its text in `prompt`, and its tail in `suffix`.""" + collect(data, "suffix", slots) + prompt: Final = data.get("prompt") + if isinstance(prompt, str): + collect(data, "prompt", slots) + return + prompt_object: Final = as_object(prompt) + if prompt_object is not None: + variables: Final = as_object(prompt_object.get("variables")) + if variables is not None: + for name in tuple(variables): + collect(variables, name, slots) + typed = as_object(variables[name]) + if typed is not None: + collect(typed, "text", slots) + return + entries: Final = as_array(prompt) + if entries is None: + return + for index in range(len(entries)): + collect_entry(entries, index, slots) + + +def collect_content(container: MutableRequest, slots: SlotSink) -> None: + """Collects `content`, a string or a list of typed parts. + + An Anthropic tool_result nests its own content, so this has to descend. It walks + with an explicit stack and a depth bound rather than by recursion: the nesting is + caller controlled, and an unbounded descent is a JSON bomb. Content nested past the + bound raises `RequestTooDeep` rather than being skipped. + """ + pending: Final[list[tuple[MutableRequest, int]]] = [(container, 0)] # mutable-ok: local queue, never escapes. + cursor = 0 # rebind-ok: advances through the queue. + while cursor < len(pending): + node, depth = pending[cursor] + cursor += 1 + content = node.get("content") + if isinstance(content, str): + collect(node, "content", slots) + continue + if depth >= MAX_CONTENT_DEPTH and content: + raise RequestTooDeep("content") + for item in as_array(content) or (): + part = as_object(item) + if part is None: + continue + collect(part, "text", slots) + if part.get("type") == "tool_use": + collect_json_leaves(part.get("input"), slots, strict=True) + source = as_object(part.get("source")) if part.get("type") == "document" else None + if source is not None: + collect(part, "title", slots) + collect(part, "context", slots) + if source.get("type") == "text": + collect(source, "data", slots) + elif source.get("type") == "content": + pending.append((source, depth + 1)) + if "content" in part: + pending.append((part, depth + 1)) + + +def collect_participant_name(message: MutableRequest, slots: SlotSink) -> None: + """Redacts `name` where it identifies a person, never where it names a function. + + On a user or assistant turn `name` is the participant, which is personal data. + On a tool or function turn the same field carries the function's name and has + to reach the provider unchanged, or the call no longer routes. + """ + if message.get("role") in ("tool", "function"): + return + collect(message, "name", slots) + + +def collect_tool_arguments(message: MutableRequest, slots: SlotSink) -> None: + """Tool arguments carry the values a user asked the model to act on.""" + for tool_call in read_list(message, "tool_calls"): + tool_call_object = as_object(tool_call) + function = as_object(tool_call_object.get("function")) if tool_call_object is not None else None + if function is not None: + collect(function, "arguments", slots) + legacy: Final = as_object(message.get("function_call")) + if legacy is not None: + collect(legacy, "arguments", slots) + + +def collect_system(data: MutableRequest, slots: SlotSink) -> None: + """Anthropic's /v1/messages carries its system prompt at the top level.""" + system: Final = data.get("system") + if isinstance(system, str): + collect(data, "system", slots) + return + collect_text_parts(data, "system", slots) + + +def collect_responses_fields(data: MutableRequest, slots: SlotSink, privileged: SlotSink) -> None: + """The Responses API sends text outside `messages`, in `instructions` and `input`. + + `instructions` is written by the application, not by the caller, so it is + collected into the privileged sink; `input` is the caller's own text, except for + system and developer items in it, which go to the privileged sink like their Chat + counterparts. + """ + collect(data, "instructions", privileged) + request_input: Final = data.get("input") + if isinstance(request_input, str): + collect(data, "input", slots) + return + entries: Final = as_array(request_input) + if entries is None: + return + for index, entry in enumerate(entries): + if isinstance(entry, str): + collect_entry(entries, index, slots) + continue + item = as_object(entry) + if item is None: + continue + collect_content(item, privileged if item.get("role") in PRIVILEGED_ROLES else slots) + collect(item, "arguments", slots) + collect(item, "output", slots) + collect_text_parts(item, "output", slots) + collect(item, "input", slots) + collect(item, "code", slots) + collect_text_parts(item, "summary", slots) + + +def collect_tool_definitions(data: MutableRequest, slots: SlotSink, privileged: SlotSink) -> None: + """Tool definitions are application-authored free text bound for the provider. + + A tool's description and the free text in its parameter schema are where callers put + examples and customer context, so they carry PII as often as a prompt does. They are + collected into the privileged sink, like a system prompt: redacted outbound, and never + restorable from the reply. `enum` and `const` values are the exception, and go to the + caller's vault -- see `SCHEMA_LITERAL_KEYWORDS`. Names and types are left as sent. + + Covers Chat `tools[].function`, the legacy `functions[]`, and the flat tool shape the + Responses API and Anthropic share, whose schema is `parameters` or `input_schema`. + """ + for key in ("tools", "functions"): + for entry in as_array(data.get(key)) or (): + tool = as_object(entry) + if tool is None: + continue + function = as_object(tool.get("function")) + for holder in (tool, function) if function is not None else (tool,): + collect(holder, "description", privileged) + collect_schema_text(holder.get("parameters"), slots, privileged) + collect_schema_text(holder.get("input_schema"), slots, privileged) + + +def collect_schema_text(schema: object, slots: SlotSink, privileged: SlotSink) -> None: + """Collects the text in a JSON Schema, at any depth. + + Scan by default: every string is collected except under the keywords in + `SCHEMA_STRUCTURAL_KEYWORDS`, whose values must go out verbatim. A list of keywords + *to* collect would leak every one it forgot -- draft-07 `dependencies`, a `$comment`, + a vendor `x-` extension -- which is how this walk started out. Free text goes to the + privileged sink; `enum` / `const` literals go to the caller's, so the model's use of + them is restored. + + Structure matters in two places. Under `properties` and the other name -> subschema + maps, keys are property names rather than keywords, so a property called `type` is a + subschema to walk, not a keyword to skip. And `examples` / `default` hold JSON values, + so all their strings are collected whatever the keys around them are called. Nested + past `MAX_JSON_DEPTH`, the request is refused. + """ + pending: Final[list[tuple[object, int]]] = [(schema, 0)] # mutable-ok: local walk stack. + while pending: + node, depth = pending.pop() + if depth > MAX_JSON_DEPTH: + if is_container(node) and node: + raise RequestTooDeep("schema") + continue + entries = as_array(node) + if entries is not None: + for index, item in enumerate(entries): + collect_entry(entries, index, privileged) + if is_container(item): + pending.append((item, depth + 1)) + continue + schema_object = as_object(node) + if schema_object is None: + continue + for keyword, value in tuple(schema_object.items()): + if keyword in SCHEMA_STRUCTURAL_KEYWORDS: + continue + subschemas = as_object(value) if keyword in SCHEMA_MAP_KEYWORDS else None + if keyword in SCHEMA_LITERAL_KEYWORDS: + collect(schema_object, keyword, slots) + collect_json_leaves(value, slots, strict=True) + elif keyword in SCHEMA_VALUE_KEYWORDS: + collect(schema_object, keyword, privileged) + collect_json_leaves(value, privileged, strict=True) + elif subschemas is not None: + pending.extend((child, depth + 1) for child in subschemas.values()) + elif isinstance(value, str): + collect(schema_object, keyword, privileged) + elif is_container(value): + pending.append((value, depth + 1)) + + +def collect_output_contracts(data: MutableRequest, slots: SlotSink, privileged: SlotSink) -> None: + """Text the caller sends to shape the reply rather than to prompt it. + + A predicted output (`prediction.content`) is the caller's own draft of the answer, so + it goes with their text: the model largely repeats it, and it has to come back. A + structured-output schema -- Chat `response_format.json_schema`, Responses + `text.format` -- is application-authored like a tool schema, so its free text goes + to the privileged sink, and its names and types stay as sent. + """ + prediction: Final = as_object(data.get("prediction")) + if prediction is not None: + collect(prediction, "content", slots) + collect_text_parts(prediction, "content", slots) + response_format: Final = as_object(data.get("response_format")) + text_options: Final = as_object(data.get("text")) + for declared in ( + response_format.get("json_schema") if response_format is not None else None, + text_options.get("format") if text_options is not None else None, + ): + wrapper = as_object(declared) + if wrapper is not None: + collect(wrapper, "description", privileged) + collect_schema_text(wrapper.get("schema"), slots, privileged) + + +def collect_user_locations(data: MutableRequest, privileged: SlotSink) -> None: + """Web search forwards the user's approximate location, whose `city` and `region` + are free text and can hold a street address. + + Chat carries it in `web_search_options.user_location.approximate`; the Responses + and Anthropic web-search tools carry it flat on the tool's `user_location`. Nothing + restores it from a reply, hence the privileged sink. + """ + options: Final = data.get("web_search_options") + tools: Final = as_array(data.get("tools")) or () + for declared in (options, *tools): + holder = as_object(declared) + location = as_object(holder.get("user_location")) if holder is not None else None + if location is None: + continue + approximate = as_object(location.get("approximate")) + for container in (location, approximate) if approximate is not None else (location,): + collect(container, "city", privileged) + collect(container, "region", privileged) + + +def collect_end_user_ids(data: MutableRequest, privileged: SlotSink) -> None: + """`user` and `safety_identifier` are forwarded to the provider and often hold an email. + + Only detected PII is replaced, so an opaque id reaches the provider unchanged. LiteLLM's + own end-user spend tracking reads the id resolved at authentication, before this hook + runs, so rewriting the field here does not move spend. Nothing restores these from a + reply, hence the privileged sink. + """ + collect(data, "user", privileged) + collect(data, "safety_identifier", privileged) + + +def locate_request_texts( + data: MutableRequest, +) -> tuple[Sequence[Slot], Sequence[Slot]]: + """Finds every redactable span, split by whether the caller can see it. + + Anything missed here reaches the provider in the clear while the guardrail + still reports as enabled, so the walk covers every request shape that + carries text. + + The split exists because the response is restored against one vault only. + Server-authored spans -- system and developer turns, Anthropic's top-level + `system`, the Responses API `instructions`, tool and output schemas -- go into a + vault nothing is ever restored against, so a caller who gets the model to + echo one of their placeholders back receives the placeholder, not the value + behind it. End-user identifiers go there too: nothing in a reply needs them. + + Tool *results* stay on the caller's side deliberately. The model reads them in + order to answer, so it can already repeat anything in them; restoring the + placeholder gives the caller the answer they would have had without this + guardrail, and an agent that reads a file and quotes an address from it needs + that address back. + + `extra_body` is walked the same way as the request itself. LiteLLM merges it over + the transformed request just before sending, so a field there -- `input`, + `messages`, `system` -- replaces the redacted one on the wire. + """ + slots: Final[SlotSink] = [] + privileged: Final[SlotSink] = [] + for payload in (data, as_object(data.get("extra_body"))): + if payload is None: + continue + for entry in read_list(payload, "messages"): + message = as_object(entry) + if message is not None: + sink = privileged if message.get("role") in PRIVILEGED_ROLES else slots + collect_content(message, sink) + collect_participant_name(message, sink) + collect_tool_arguments(message, sink) + collect_responses_fields(payload, slots, privileged) + collect_prompt(payload, slots) + collect_system(payload, privileged) + collect_tool_definitions(payload, slots, privileged) + collect_output_contracts(payload, slots, privileged) + collect_user_locations(payload, privileged) + collect_end_user_ids(payload, privileged) + return tuple(slots), tuple(privileged) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py new file mode 100644 index 00000000000..4ef8bbdbb3e --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py @@ -0,0 +1,345 @@ +import copy +import functools +import itertools +import json +import re +from enum import Enum +from types import MappingProxyType +from typing import ( + Final, + TypeAlias, +) + +from .payload import ( + JsonBody, + MutableRequest, + Rehydrate, + SlotSink, + StreamStep, + as_object, + collect_response_item, + read_field, + read_list, + rehydrate_slots, + write_field, +) + +ANTHROPIC_DELTA_FIELDS: Final = MappingProxyType({"text_delta": "text", "input_json_delta": "partial_json"}) + +SSE_EVENT_BOUNDARY: Final = re.compile(rb"(\r?\n\r?\n)") + +SSE_OPENINGS: Final = (b"event:", b"data:", b"id:", b"retry:", b":") + +RESPONSES_BINARY_DELTAS: Final = frozenset(("response.audio.delta",)) + +RESPONSES_STRUCTURAL_FIELDS: Final = frozenset( + ("type", "id", "item_id", "call_id", "name", "server_label", "status", "obfuscation") +) + +RESPONSES_TERMINAL_EVENTS: Final = frozenset(("response.completed", "response.incomplete")) + +CarryKey: TypeAlias = tuple[int, int | None] +CarryWindows: TypeAlias = dict[CarryKey, str] + +ResponsesStreamKey: TypeAlias = tuple[str, object, object, object] + + +def carry_sort_key(key: CarryKey) -> tuple[int, int]: + """Orders streaming windows without ever comparing None to an int. + + `sorted()` over the raw keys raises as soon as one choice holds both a content window + and a tool-call window, because `None < 0` is not orderable. Content sorts first, then + tool calls by their index. + """ + choice_index, tool_index = key + return (choice_index, -1 if tool_index is None else tool_index) + + +def continuation_delta(tool_index: int, text: str) -> list[dict[str, object]]: + """A `tool_calls` delta carrying `text` as an index-only continuation. + + Clients concatenate tool-call fragments by index, so no id or name is needed. + """ + return [{"index": tool_index, "function": {"arguments": text}}] + + +def opens_like_sse(head: bytes) -> bool | None: + """Whether a raw stream is SSE, judged by its opening bytes; None while undecidable. + + An SSE stream opens with a field name or a `:` comment. Anything else -- a JSON array + streamed in pieces, say -- has no event boundaries to wait for. A chunk that ends + partway through a field name decides nothing yet, so that case waits for more. + """ + opening: Final = head.lstrip() + if not opening: + return None + if opening.startswith(SSE_OPENINGS): + return True + if any(field.startswith(opening) for field in SSE_OPENINGS): + return None + return False + + +def responses_event_type(chunk: object) -> str | None: + """The event type of a Responses API stream event, or None for any other chunk. + + The type arrives as a plain string on dicts and as a str-valued Enum on LiteLLM's + event models. The Enum is unwrapped because it does not hash like its value, so it + would miss every lookup in the event tables above. + """ + if isinstance(chunk, (bytes, str)): + return None + kind: Final = read_field(chunk, "type") + value: Final = kind.value if isinstance(kind, Enum) else kind + return value if isinstance(value, str) and value.startswith("response.") else None + + +class AnthropicSSERestorer: + """Restores an Anthropic `/v1/messages` stream, which reaches the hook as raw SSE. + + Each content block is its own token stream with its own window, keyed by the block's + `index`: `text_delta` carries prose and `input_json_delta` a tool call's arguments. + When a block stops, whatever its window still holds is emitted as one more delta for + that block, just ahead of the `content_block_stop` frame, so the client has the whole + block before it is told the block is complete. + + Frames are processed whole. A network chunk can end in the middle of an event, so the + unfinished tail is kept until the rest arrives; that delays one partial event, never + a completed one. A frame that is not an Anthropic event -- another endpoint's SSE, or + anything that fails to parse -- is passed through byte for byte, and a raw stream that + does not open like SSE at all is passed through chunk by chunk, never buffered. + """ + + def __init__(self, step: StreamStep) -> None: + self._step: Final = step + self._carries: Final[dict[int, str]] = {} # mutable-ok: per-block windows advanced in place. + self._delta_types: Final[dict[int, str]] = {} # mutable-ok: each block's delta type, for its flush. + self._pending = b"" + self._as_text = False + self._is_sse: bool | None = None + + async def feed(self, chunk: bytes | str) -> tuple[bytes | str, ...]: + """Restores every event this chunk completes; holds back an unfinished tail.""" + if isinstance(chunk, str): + self._as_text = True + if self._is_sse is False: + return (chunk,) + raw: Final = chunk.encode("utf-8") if isinstance(chunk, str) else chunk + buffered: Final = self._pending + raw + if self._is_sse is None: + self._is_sse = opens_like_sse(buffered) + if self._is_sse is None: + self._pending = buffered + return () + if not self._is_sse: + self._pending = b"" + return self._emit(buffered) + boundaries: Final = tuple(SSE_EVENT_BOUNDARY.finditer(buffered)) + if not boundaries: + self._pending = buffered + return () + cut: Final = boundaries[-1].end() + self._pending = buffered[cut:] + parts: Final = SSE_EVENT_BOUNDARY.split(buffered[:cut]) + restored: Final = tuple( + [await self._restore_event(parts[index]) + parts[index + 1] for index in range(0, len(parts) - 1, 2)] + ) + return self._emit(b"".join(restored)) + + async def finish(self) -> tuple[bytes | str, ...]: + """Emits an unterminated final event and any window a block never closed.""" + held: Final = self._pending + self._pending = b"" + if not self._is_sse: + return self._emit(held) + tail: Final = await self._restore_event(held) if held.strip() else held + flushed: Final = await self._flush_all() + separator: Final = b"\n\n" if tail.strip() and flushed else b"" + return self._emit(tail + separator + flushed) + + def _emit(self, frames: bytes) -> tuple[bytes | str, ...]: + if not frames: + return () + return (frames.decode("utf-8") if self._as_text else frames,) + + async def _restore_event(self, block: bytes) -> bytes: + """Rewrites one SSE event, or returns it untouched if it carries nothing to restore.""" + try: + lines: Final = block.decode("utf-8").split("\n") + except UnicodeDecodeError: + return block + data_lines: Final = tuple(index for index, line in enumerate(lines) if line.startswith("data:")) + if len(data_lines) != 1: + return block + line: Final = lines[data_lines[0]] + try: + parsed: Final[object] = json.loads(line[len("data:") :]) + except ValueError: + return block + event: Final = as_object(parsed) + if event is None: + return block + kind: Final = event.get("type") + index: Final = event.get("index") + if kind == "content_block_stop" and isinstance(index, int): + return await self._flush(index) + block + if kind == "message_stop": + return await self._flush_all() + block + if kind != "content_block_delta" or not await self._restore_delta(event): + return block + ending: Final = "\r" if line.endswith("\r") else "" + rewritten: Final = ( + *lines[: data_lines[0]], + f"data: {json.dumps(event, ensure_ascii=False)}{ending}", + *lines[data_lines[0] + 1 :], + ) + return "\n".join(rewritten).encode("utf-8") + + async def _restore_delta(self, event: MutableRequest) -> bool: + """Advances one block's window through this delta. False if it holds no text.""" + index: Final = event.get("index") + delta: Final = event.get("delta") + if not isinstance(index, int) or not isinstance(delta, dict): + return False + delta_type: Final = delta.get("type") + if not isinstance(delta_type, str): + return False + field: Final = ANTHROPIC_DELTA_FIELDS.get(delta_type) + text: Final = delta.get(field) if field is not None else None + if field is None or not isinstance(text, str) or not text: + return False + emitted, remaining = await self._step(text, self._carries.get(index, ""), False) + self._carries[index] = remaining + self._delta_types[index] = delta_type + delta[field] = emitted + return True + + async def _flush(self, index: int) -> bytes: + """One synthetic delta frame carrying whatever `index`'s window still holds.""" + carry: Final = self._carries.pop(index, "") + delta_type: Final = self._delta_types.pop(index, None) + field: Final = ANTHROPIC_DELTA_FIELDS.get(delta_type) if isinstance(delta_type, str) else None + if not carry or field is None: + return b"" + text, _ = await self._step("", carry, True) + if not text: + return b"" + event: Final[JsonBody] = { + "type": "content_block_delta", + "index": index, + "delta": {"type": delta_type, field: text}, + } + return f"event: content_block_delta\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + + async def _flush_all(self) -> bytes: + flushed: Final = tuple([await self._flush(index) for index in tuple(self._carries)]) + return b"".join(flushed) + + +class ResponsesStreamRestorer: + """Restores a Responses API event stream. + + The event families are matched by shape rather than listed, so a text stream the + API adds later is restored by default instead of leaking a placeholder: + + - Any `*.delta` event whose `delta` is a string is a token stream (output_text, + refusal, function-call and MCP arguments, reasoning summaries, ...). Each gets its + own window, keyed by the family, the item id and the part index. + - Any `*.done` event closes the stream of the same family. Whatever its window still + holds goes out first, as a copy of that stream's last delta event -- so it carries + the stream's own ids, and repeats that event's `sequence_number`. Then every text + field on the done event is restored in full: its string fields other than + identifiers, plus any `part` or `item` it repeats. + - `response.completed` / `response.incomplete` repeat the whole reply, and are + restored the same way the non-streaming reply is. + """ + + def __init__(self, step: StreamStep, rehydrate: Rehydrate) -> None: + self._step: Final = step + self._rehydrate: Final = rehydrate + self._carries: Final[dict[ResponsesStreamKey, str]] = {} # mutable-ok: per-stream windows advanced in place. + self._last_deltas: Final[dict[ResponsesStreamKey, object]] = {} # mutable-ok: newest delta per stream. + + async def restore(self, event: object) -> tuple[object, ...]: + """The events to emit in place of `event`: any flush, then the event itself.""" + kind: Final = responses_event_type(event) + if kind is None: + return (event,) + if kind.endswith(".delta") and kind not in RESPONSES_BINARY_DELTAS: + await self._restore_delta(event, kind) + return (event,) + slots: Final[SlotSink] = [] + flushed: Final = await self._flush(responses_stream_key(event, kind)) if kind.endswith(".done") else () + if kind.endswith(".done"): + collect_event_text(event, slots) + part: Final = read_field(event, "part") + if part is not None: + collect_response_item({"content": [part]}, slots) + collect_response_item(read_field(event, "item"), slots) + elif kind in RESPONSES_TERMINAL_EVENTS: + for item in read_list(read_field(event, "response"), "output"): + collect_response_item(item, slots) + await rehydrate_slots(slots, self._rehydrate) + return (*flushed, event) + + async def finish(self) -> tuple[object, ...]: + """Flushes every stream the provider never closed, e.g. a truncated reply.""" + flushed: Final = tuple([await self._flush(key) for key in tuple(self._carries)]) + return tuple(itertools.chain.from_iterable(flushed)) + + async def _restore_delta(self, event: object, kind: str) -> None: + text: Final = read_field(event, "delta") + if not isinstance(text, str) or not text: + return + key: Final = responses_stream_key(event, kind) + emitted, remaining = await self._step(text, self._carries.get(key, ""), False) + self._carries[key] = remaining + self._last_deltas[key] = event + write_field(event, "delta", emitted) + + async def _flush(self, key: ResponsesStreamKey) -> tuple[object, ...]: + carry: Final = self._carries.pop(key, "") + template: Final = self._last_deltas.pop(key, None) + if not carry or template is None: + return () + text, _ = await self._step("", carry, True) + if not text: + return () + flush: Final = copy.deepcopy(template) + write_field(flush, "delta", text) + return (flush,) + + +def responses_stream_key(event: object, kind: str) -> ResponsesStreamKey: + """Identifies the delta stream an event belongs to, the same for its delta and done. + + The family is the event type without its `.delta` / `.done` suffix, so an output_text + stream and a refusal stream on the same part never share a window. + """ + family: Final = kind.rsplit(".", 1)[0] + part_index: Final = read_field(event, "content_index") + summary_index: Final = read_field(event, "summary_index") + return ( + family, + read_field(event, "item_id"), + read_field(event, "output_index"), + part_index if part_index is not None else summary_index, + ) + + +def collect_event_text(event: object, slots: SlotSink) -> None: + """Collects every top-level text field of a Responses API event, dict or model. + + Scan by default, with identifiers excluded, rather than a list of known fields: the + `.done` event of each stream family names its text differently (`text`, `refusal`, + `arguments`, ...), and a family added upstream would otherwise leak a placeholder. + """ + attributes: Final[object] = getattr(event, "__dict__", None) + fields: Final = as_object(event) or as_object(attributes) + if fields is None: + return + for name, value in tuple(fields.items()): + if name in RESPONSES_STRUCTURAL_FIELDS or name.endswith("_id"): + continue + if isinstance(value, str) and value: + slots.append((value, functools.partial(write_field, event, name))) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 46026c12d24..8f1140cef62 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -145,6 +145,7 @@ class SupportedGuardrailIntegrations(Enum): STRAIKER = "straiker" ALICE = "alice" AGENT_365 = "agent_365" + LLM_SHIELD_PROXY = "llm_shield_proxy" CONDUCT = "conduct" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield_proxy.py b/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield_proxy.py new file mode 100644 index 00000000000..967d7ee75c2 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield_proxy.py @@ -0,0 +1,24 @@ +from pydantic import Field + +from .base import GuardrailConfigModel + + +class LLMShieldProxyGuardrailConfigModel(GuardrailConfigModel): + api_key: str | None = Field( + default=None, + description=( + "The virtual key for the LLM Shield Proxy instance. If not provided, the " + "`LLM_SHIELD_PROXY_API_KEY` environment variable is checked." + ), + ) + api_base: str | None = Field( + default=None, + description=( + "The base URL of the LLM Shield Proxy instance. If not provided, the `LLM_SHIELD_PROXY_API_BASE` " + "environment variable is checked, then `http://localhost:8000`." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "LLM Shield Proxy" diff --git a/ruff-strict.toml b/ruff-strict.toml index 899a8ff3af5..9e0b11a1af4 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -30,6 +30,10 @@ external = [ # grows over time; typing it concretely (`object`) broke that forwarding call outright — # basedpyright turned every named param into a reportArgumentType error. Any is correct here. "litellm/proxy/guardrails/guardrail_hooks/alice/alice.py" = ["ANN401"] +# Same reason: `**kwargs` forwards verbatim to CustomGuardrail.__init__, and the lifecycle +# hook signatures inherit `Any` for `response` from CustomLogger, so narrowing them here +# would break the override rather than describe it. +"litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py" = ["ANN401"] [lint.mccabe] max-complexity = 15 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py new file mode 100644 index 00000000000..ed34c07022e --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -0,0 +1,2003 @@ +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from httpx import Request, Response + +import litellm +from litellm.caching.caching_handler import _PENDING_CACHE_WRITES +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy.guardrails.guardrail_hooks.llm_shield_proxy.llm_shield_proxy import ( + GUARDRAIL_NAME, + LLMShieldProxyGuardrail, +) +from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import ( + FunctionCallArgumentsDeltaEvent, + OutputTextDeltaEvent, + OutputTextDoneEvent, + ResponsesAPIStreamEvents, +) +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + TextChoices, + TextCompletionResponse, + Usage, +) + + +def _guardrail(**overrides: object) -> LLMShieldProxyGuardrail: + params: dict[str, object] = { + "api_key": "test-key", + "api_base": "http://shield.test", + "guardrail_name": GUARDRAIL_NAME, + "event_hook": "pre_call", + "default_on": True, + } + params.update(overrides) + return LLMShieldProxyGuardrail(**params) + + +def _response(payload: dict, status_code: int = 200) -> Response: + return Response( + status_code=status_code, + json=payload, + request=Request("POST", "http://shield.test/v1/guard/redact"), + ) + + +def _mock_post(guardrail: LLMShieldProxyGuardrail, *payloads: dict) -> AsyncMock: + """Queues one shield response per expected call.""" + mock = AsyncMock(side_effect=[_response(p) for p in payloads]) + guardrail.async_handler.post = mock # type: ignore[method-assign] + return mock + + +def _chunk(content: str | None, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)] + ) + + +async def _drain(generator) -> list: + return [chunk async for chunk in generator] + + +def _tool_chunk(arguments: str, finish_reason: str | None = None) -> ModelResponseStream: + """One streamed fragment of tool call 0's arguments.""" + tool_call = {"index": 0, "id": "call_1", "type": "function", "function": {"name": "send", "arguments": arguments}} + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[tool_call]), finish_reason=finish_reason)] + ) + + +def _field(holder: object, name: str) -> object: + """Reads a field from a dict or a model; the guardrail emits both shapes.""" + return holder.get(name) if isinstance(holder, dict) else getattr(holder, name) + + +class _FakeShield: + """The three guard endpoints over one fixed vault, placeholder -> original. + + The stream endpoint holds back a trailing `[` that has not closed yet, which is the + behaviour that makes a placeholder split across two chunks come out whole. + """ + + def __init__(self, vault: dict[str, str]) -> None: + self.vault = vault + self.urls: list[str] = [] + + def _restore(self, text: str) -> str: + for placeholder, original in self.vault.items(): + text = text.replace(placeholder, original) + return text + + async def post(self, url: str, headers: dict, json: dict, timeout: float) -> Response: + self.urls.append(url) + if url.endswith("/rehydrate/stream"): + text = self._restore(json["carry"] + json["text"]) + opening = text.rfind("[") + if json["final"] or opening == -1 or "]" in text[opening:]: + return _response({"text": text, "carry": ""}) + return _response({"text": text[:opening], "carry": text[opening:]}) + return _response({"texts": [self._restore(text) for text in json["texts"]]}) + + +def _shielded(vault: dict[str, str]) -> tuple[LLMShieldProxyGuardrail, _FakeShield]: + guardrail = _guardrail(event_hook="post_call") + shield = _FakeShield(vault) + guardrail.async_handler.post = shield.post # type: ignore[method-assign] + return guardrail, shield + + +def _sse(event: dict) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _sse_events(frames: list) -> list[dict]: + """Parses emitted SSE output, whatever its chunking, back into event payloads.""" + raw = b"".join(frame.encode() if isinstance(frame, str) else frame for frame in frames).decode() + return [ + json.loads(line[len("data:") :]) + for event in raw.split("\n\n") + for line in event.split("\n") + if line.startswith("data:") + ] + + +def _text_block_stream(*deltas: str) -> list[bytes]: + """An Anthropic /v1/messages stream with one text block made of `deltas`.""" + return [ + _sse({"type": "message_start", "message": {"id": "msg_1", "role": "assistant", "content": []}}), + _sse({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + *( + _sse({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": d}}) + for d in deltas + ), + _sse({"type": "content_block_stop", "index": 0}), + _sse({"type": "message_stop"}), + ] + + +async def _restore_stream(guardrail: LLMShieldProxyGuardrail, chunks: list) -> list: + async def stream(): + for chunk in chunks: + yield chunk + + return await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + +def test_llm_shield_guardrail_config(monkeypatch: pytest.MonkeyPatch): + """Should register through init_guardrails_v2 like any other provider.""" + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) + monkeypatch.setenv("LLM_SHIELD_PROXY_API_KEY", "test-key") + + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "llm_shield_proxy", + "litellm_params": {"guardrail": "llm_shield_proxy", "mode": "pre_call", "default_on": True}, + } + ], + config_file_path="", + ) + + registered = [cb for cb in litellm.callbacks if isinstance(cb, LLMShieldProxyGuardrail)] + assert len(registered) == 1 + assert registered[0].guardrail_name == "llm_shield_proxy" + + +class TestLLMShieldProxyInitialization: + def test_api_base_defaults_to_localhost(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LLM_SHIELD_PROXY_API_BASE", raising=False) + assert _guardrail(api_base=None).api_base == "http://localhost:8000" + + def test_api_base_reads_environment(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LLM_SHIELD_PROXY_API_BASE", "http://shield.internal:9000") + assert _guardrail(api_base=None).api_base == "http://shield.internal:9000" + + def test_trailing_slash_is_stripped(self): + assert _guardrail(api_base="http://shield.test/").api_base == "http://shield.test" + + def test_both_modes_can_be_enabled_on_one_entry(self): + """Redaction and restoration are two halves of one config entry. + + A deployment that lists only pre_call would redact the request and then hand + the placeholders straight back to the end user. + """ + guardrail = _guardrail(event_hook=["pre_call", "post_call"]) + data: dict = {"messages": []} + + assert guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is True + assert guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is True + assert guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.during_call) is False + + +class TestRedaction: + @pytest.mark.asyncio + async def test_string_content_is_redacted(self): + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["Email [EMAIL_1] about it"]}) + + data = {"messages": [{"role": "user", "content": "Email a@b.com about it"}]} + result = await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="completion" + ) + + assert result["messages"][0]["content"] == "Email [EMAIL_1] about it" + + @pytest.mark.asyncio + async def test_multimodal_text_parts_are_redacted(self): + """The list content shape is a historical bypass; text parts must be covered.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["call [PHONE_1]"]}) + + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "call 555-0100"}, + {"type": "image_url", "image_url": {"url": "http://x/y.png"}}, + ], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["content"][0]["text"] == "call [PHONE_1]" + assert data["messages"][0]["content"][1]["image_url"]["url"] == "http://x/y.png" + + @pytest.mark.asyncio + async def test_request_without_text_is_untouched(self): + """No text to redact means no call to LLM Shield Proxy. + + This deliberately uses a request with no caller text at all. An earlier + version used a Responses-API `input`, which asserted the very bypass that + let `input` reach the provider unredacted. + """ + guardrail = _guardrail() + mock = _mock_post(guardrail) + data = {"model": "gpt-4o", "temperature": 0.2} + + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + mock.assert_not_called() + + @pytest.mark.asyncio + async def test_session_id_is_reused_across_hooks(self): + """Rehydration can only resolve tokens minted under the same session.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["a@b.com"]}) + + data = {"messages": [{"role": "user", "content": "a@b.com"}]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + await guardrail._rehydrate(["[EMAIL_1]"], guardrail._session_id(data)) + + sessions = {call.kwargs["headers"]["X-Session-ID"] for call in mock.call_args_list} + assert len(sessions) == 1 + + +class TestRequestCoverage: + """Every request shape that carries caller text must be redacted. + + A shape missed here is not a cosmetic gap: the guardrail reports as enabled + while the raw value goes to the provider. + """ + + @pytest.mark.asyncio + async def test_responses_api_string_input_is_redacted(self): + """Measured against a live provider: `input` reached the model unredacted.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["Email [EMAIL_1] the invoice"]}) + + data = {"input": "Email jane.doe@example.com the invoice"} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["Email jane.doe@example.com the invoice"] + assert data["input"] == "Email [EMAIL_1] the invoice" + + @pytest.mark.asyncio + async def test_responses_api_list_input_is_redacted(self): + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[PHONE_1]"]}) + + data = { + "input": [ + {"role": "user", "content": "jane.doe@example.com"}, + {"role": "user", "content": [{"type": "input_text", "text": "555-0100"}]}, + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["input"][0]["content"] == "[EMAIL_1]" + assert data["input"][1]["content"][0]["text"] == "[PHONE_1]" + + @pytest.mark.asyncio + async def test_tool_call_arguments_are_redacted(self): + """Tool arguments carry the values the user asked the model to act on.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ['{"email": "[EMAIL_1]"}']}) + + data = { + "messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "send", "arguments": '{"email": "jane.doe@example.com"}'}, + } + ], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == '{"email": "[EMAIL_1]"}' + + @pytest.mark.asyncio + async def test_responses_api_instructions_are_redacted(self): + """`instructions` is provider-bound text that sits outside `messages`.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["contact [EMAIL_1]"]}) + + data = {"instructions": "contact jane.doe@example.com", "input": ""} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["contact jane.doe@example.com"] + assert data["instructions"] == "contact [EMAIL_1]" + + @pytest.mark.asyncio + async def test_legacy_function_call_arguments_are_redacted(self): + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ['{"email": "[EMAIL_1]"}']}) + + data = { + "messages": [ + { + "role": "assistant", + "function_call": {"name": "send", "arguments": '{"email": "jane.doe@example.com"}'}, + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["function_call"]["arguments"] == '{"email": "[EMAIL_1]"}' + + @pytest.mark.asyncio + async def test_completions_prompt_is_redacted(self): + """/v1/completions puts its text in a top-level `prompt`, not in messages.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["Email [EMAIL_1]"]}) + + data = {"prompt": "Email jane.doe@example.com"} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="atext_completion") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["Email jane.doe@example.com"] + assert data["prompt"] == "Email [EMAIL_1]" + + @pytest.mark.asyncio + async def test_completions_prompt_array_is_redacted(self): + """`prompt` also accepts an array, and each entry is provider-bound.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[PHONE_1]"]}) + + data = {"prompt": ["jane.doe@example.com", "555-0100"]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="atext_completion") + + assert data["prompt"] == ["[EMAIL_1]", "[PHONE_1]"] + + @pytest.mark.asyncio + async def test_responses_function_call_items_are_redacted(self): + """Responses input items hold tool data in `arguments` and `output`.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ['{"email": "[EMAIL_1]"}', "sent to [EMAIL_1]"]}) + + data = { + "input": [ + {"type": "function_call", "name": "send", "arguments": '{"email": "jane.doe@example.com"}'}, + {"type": "function_call_output", "call_id": "c1", "output": "sent to jane.doe@example.com"}, + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["input"][0]["arguments"] == '{"email": "[EMAIL_1]"}' + assert data["input"][1]["output"] == "sent to [EMAIL_1]" + + @pytest.mark.asyncio + async def test_responses_tool_output_parts_are_redacted(self): + """A function_call_output can carry its result as a list of input_text parts.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["sent to [EMAIL_1]"]}) + + data = { + "input": [ + { + "type": "function_call_output", + "call_id": "c1", + "output": [{"type": "input_text", "text": "sent to jane.doe@example.com"}], + }, + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["sent to jane.doe@example.com"] + assert data["input"][0]["output"][0]["text"] == "sent to [EMAIL_1]" + + @pytest.mark.asyncio + async def test_responses_custom_tool_call_input_is_redacted(self): + """A replayed custom_tool_call carries its payload in `input`, not `arguments`.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["email [EMAIL_1]"]}) + + data = { + "input": [ + {"type": "custom_tool_call", "call_id": "c1", "name": "mail", "input": "email jane.doe@example.com"}, + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["input"][0]["input"] == "email [EMAIL_1]" + assert data["input"][0]["name"] == "mail" + + @pytest.mark.asyncio + async def test_extra_body_overrides_are_redacted(self): + """LiteLLM merges `extra_body` over the request just before sending, so its fields win on the wire.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["safe", "Mail [EMAIL_1]", "[EMAIL_1]"]}) + + data = { + "model": "gpt-4o", + "input": "safe", + "extra_body": { + "input": "alice@example.com", + "messages": [{"role": "user", "content": "Mail alice@example.com"}], + "service_tier": "flex", + }, + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["safe", "Mail alice@example.com", "alice@example.com"] + assert data["extra_body"]["input"] == "[EMAIL_1]" + assert data["extra_body"]["messages"][0]["content"] == "Mail [EMAIL_1]" + assert data["extra_body"]["service_tier"] == "flex" + + def test_extra_body_system_text_is_privileged(self) -> None: + """An application-authored override is no more restorable than the field it replaces.""" + data = {"messages": [{"role": "user", "content": "U"}], "extra_body": {"system": "S", "instructions": "I"}} + + caller, privileged = LLMShieldProxyGuardrail._locate_request_texts(data) + + assert [text for text, _ in caller] == ["U"] + assert sorted(text for text, _ in privileged) == ["I", "S"] + + @pytest.mark.asyncio + async def test_anthropic_text_documents_are_redacted(self): + """A document block carries text inline, as `source.data` or as `source.content`.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[NAME_1] notes", "[EMAIL_1]", "Reach [EMAIL_2]"]}) + + data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "document", + "title": "Jane Doe notes", + "source": {"type": "text", "media_type": "text/plain", "data": "alice@example.com"}, + }, + { + "type": "document", + "source": {"type": "content", "content": [{"type": "text", "text": "Reach bob@example.com"}]}, + }, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0x"}}, + ], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="anthropic_messages") + + sent = mock.call_args_list[0].kwargs["json"]["texts"] + assert sent == ["Jane Doe notes", "alice@example.com", "Reach bob@example.com"] + blocks = data["messages"][0]["content"] + assert blocks[0]["source"]["data"] == "[EMAIL_1]" + assert blocks[1]["source"]["content"][0]["text"] == "Reach [EMAIL_2]" + assert blocks[2]["source"]["data"] == "JVBERi0x", "binary sources are not text" + + @pytest.mark.asyncio + async def test_responses_code_interpreter_code_is_redacted(self): + """A replayed code_interpreter_call carries the code the model wrote, which the reply side restores.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["send('[EMAIL_1]')"]}) + + data = {"input": [{"type": "code_interpreter_call", "id": "ci_1", "code": "send('jane.doe@example.com')"}]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["input"][0]["code"] == "send('[EMAIL_1]')" + + @pytest.mark.asyncio + async def test_anthropic_system_prompt_is_redacted(self): + """/v1/messages carries its system prompt at the top level, not in messages.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["the user is [EMAIL_1]"]}) + + data = {"system": "the user is jane.doe@example.com", "messages": []} + await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="anthropic_messages" + ) + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["the user is jane.doe@example.com"] + assert data["system"] == "the user is [EMAIL_1]" + + @pytest.mark.asyncio + async def test_anthropic_system_blocks_are_redacted(self): + """`system` also accepts a list of text blocks.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = {"system": [{"type": "text", "text": "jane.doe@example.com"}], "messages": []} + await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="anthropic_messages" + ) + + assert data["system"][0]["text"] == "[EMAIL_1]" + + @pytest.mark.asyncio + async def test_string_array_input_is_redacted(self): + """Embeddings and moderations send `input` as an array of bare strings.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[PHONE_1]"]}) + + data = {"input": ["jane.doe@example.com", "555-0100"]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aembedding") + + assert data["input"] == ["[EMAIL_1]", "[PHONE_1]"] + + @pytest.mark.asyncio + async def test_participant_name_is_redacted(self): + """`name` on a user turn identifies a person.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["hi", "[PERSON_1]"]}) + + data = {"messages": [{"role": "user", "name": "Jane Doe", "content": "hi"}]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["hi", "Jane Doe"] + assert data["messages"][0]["name"] == "[PERSON_1]" + + @pytest.mark.asyncio + async def test_tool_function_name_is_left_alone(self): + """On a tool turn the same field is the function name. + + Redacting it would stop the call routing, so this asserts it is never sent + to the shield at all. + """ + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["result"]}) + + data = {"messages": [{"role": "tool", "name": "get_weather", "content": "result"}]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["name"] == "get_weather" + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["result"] + + @pytest.mark.asyncio + async def test_anthropic_tool_result_content_is_redacted(self): + """A tool_result nests its own content, as a string or as more blocks.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[EMAIL_2]"]}) + + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "found jane.doe@example.com"}, + { + "type": "tool_result", + "tool_use_id": "t2", + "content": [{"type": "text", "text": "also bob@example.com"}], + }, + ], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["content"][0]["content"] == "[EMAIL_1]" + assert data["messages"][0]["content"][1]["content"][0]["text"] == "[EMAIL_2]" + + @pytest.mark.asyncio + async def test_nesting_past_the_bound_blocks_the_request(self): + """Nesting is caller controlled, so the descent has to stop somewhere -- and where + it stops, the request must not go out. + + This test used to assert the opposite: that text past the bound was skipped. That + sent `past-the-bound@example.com` to the provider unredacted while the guardrail + reported as enabled. + """ + guardrail = _guardrail() + mock = _mock_post(guardrail) + + deep: dict = {"type": "tool_result", "content": "past-the-bound@example.com"} + for _ in range(200): + deep = {"type": "tool_result", "content": [deep]} + data = {"messages": [{"role": "user", "content": [{"type": "text", "text": "shallow"}, deep]}]} + + with pytest.raises(GuardrailRaisedException): + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + mock.assert_not_called() + + @pytest.mark.asyncio + async def test_deep_tool_input_blocks_the_request(self): + """A tool_use input past the JSON bound must not be forwarded half-redacted.""" + guardrail = _guardrail() + mock = _mock_post(guardrail) + + deep: dict = {"email": "past-the-bound@example.com"} + for _ in range(100): + deep = {"next": deep} + block = {"type": "tool_use", "id": "t1", "name": "f", "input": deep} + data = {"messages": [{"role": "assistant", "content": [block]}]} + + with pytest.raises(GuardrailRaisedException): + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + mock.assert_not_called() + + @pytest.mark.asyncio + async def test_realistic_nesting_is_redacted_in_full(self): + """The bounds are far past real payloads: a tool input nested inside a tool result, + several JSON levels deep, is redacted whole rather than refused.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + tool_use = { + "type": "tool_use", + "id": "t1", + "name": "f", + "input": {"a": {"b": {"c": {"d": {"to": "x@example.com"}}}}}, + } + data = {"messages": [{"role": "user", "content": [{"type": "tool_result", "content": [tool_use]}]}]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert tool_use["input"]["a"]["b"]["c"]["d"]["to"] == "[EMAIL_1]" + + @pytest.mark.asyncio + async def test_responses_prompt_object_variables_are_redacted(self): + """A PromptObject's variables are substituted into the prompt provider side. + + The id and version pick which stored prompt to run and have to arrive + unchanged; the variables are caller text. + """ + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = {"prompt": {"id": "pmpt_123", "version": "2", "variables": {"customer": "jane.doe@example.com"}}} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["prompt"]["variables"]["customer"] == "[EMAIL_1]" + assert data["prompt"]["id"] == "pmpt_123" + assert data["prompt"]["version"] == "2" + + @pytest.mark.asyncio + async def test_responses_typed_prompt_variables_are_redacted(self): + """A variable can be a typed input rather than a string; its `text` is caller text.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = { + "prompt": { + "id": "pmpt_123", + "variables": { + "customer": {"type": "input_text", "text": "jane.doe@example.com"}, + "logo": {"type": "input_image", "image_url": "https://example.com/logo.png"}, + }, + } + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["prompt"]["variables"]["customer"] == {"type": "input_text", "text": "[EMAIL_1]"} + assert data["prompt"]["variables"]["logo"]["image_url"] == "https://example.com/logo.png" + + @pytest.mark.asyncio + async def test_completions_suffix_is_redacted(self): + """LiteLLM forwards the legacy `suffix` to providers that support it.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["signed [EMAIL_1]", "write to [EMAIL_1]"]}) + + data = {"prompt": "write to jane.doe@example.com", "suffix": "signed jane.doe@example.com"} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="atext_completion") + + assert data["suffix"] == "signed [EMAIL_1]" + + @pytest.mark.asyncio + async def test_every_shape_in_one_request_is_redacted(self): + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["a", "b", "c", "d"]}) + + data = { + "messages": [ + {"role": "user", "content": "one"}, + {"role": "user", "content": [{"type": "text", "text": "two"}]}, + { + "role": "assistant", + "tool_calls": [{"function": {"name": "f", "arguments": "three"}}], + }, + ], + "input": "four", + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["one", "two", "three", "four"] + assert data["messages"][0]["content"] == "a" + assert data["messages"][1]["content"][0]["text"] == "b" + assert data["messages"][2]["tool_calls"][0]["function"]["arguments"] == "c" + assert data["input"] == "d" + + @pytest.mark.asyncio + async def test_anthropic_tool_use_input_is_redacted(self): + """A replayed tool_use block carries its arguments as a JSON object, not a string.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[PHONE_1]"]}) + + data = { + "messages": [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "t1", + "name": "send", + "input": {"to": "jane.doe@example.com", "meta": {"phone": "555-0100"}}, + } + ], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["jane.doe@example.com", "555-0100"] + block = data["messages"][0]["content"][0] + assert block["input"] == {"to": "[EMAIL_1]", "meta": {"phone": "[PHONE_1]"}} + assert block["name"] == "send", "the tool name has to arrive unchanged for the call to route" + + @pytest.mark.asyncio + async def test_responses_reasoning_summary_is_redacted(self): + """A replayed reasoning item quotes the conversation in its summary parts.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["user asked about [EMAIL_1]"]}) + + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "summary": [{"type": "summary_text", "text": "user asked about jane.doe@example.com"}], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["input"][0]["summary"][0]["text"] == "user asked about [EMAIL_1]" + + def test_tool_schemas_give_up_their_free_text_and_nothing_else(self): + """Every string is collected except what must reach the model verbatim: names, + types, formats, patterns and required lists.""" + data = { + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "top", + "parameters": { + "type": "object", + "title": "title", + "properties": { + # A property that is itself named "description". + "description": {"type": "string", "description": "named"}, + "kind": {"type": "string", "enum": ["a", "b"], "const": "a", "description": "enum"}, + "deep": {"type": "array", "items": {"type": "object", "description": "nested"}}, + "to": { + "type": "string", + "format": "email", + "pattern": "^.+@.+$", + "examples": ["example"], + "default": "default", + }, + "choice": {"anyOf": [{"type": "object", "default": {"type": "object-default"}}]}, + # Property names that collide with keywords are subschemas all the same. + "type": {"type": "string", "description": "named-type"}, + }, + "required": ["to"], + "$defs": {"shared": {"description": "defined"}}, + # Keywords nobody listed: scanned by default. + "dependencies": {"mode": {"description": "dependent"}}, + "$comment": "comment", + "x-note": "vendor", + }, + }, + } + ] + } + caller, privileged = LLMShieldProxyGuardrail._locate_request_texts(data) + + assert sorted(text for text, _ in caller) == ["a", "a", "b"], "enum and const go to the caller vault" + assert sorted(text for text, _ in privileged) == [ + "comment", + "default", + "defined", + "dependent", + "enum", + "example", + "named", + "named-type", + "nested", + "object-default", + "title", + "top", + "vendor", + ] + + @pytest.mark.asyncio + async def test_enum_values_are_redacted_and_restored_in_the_tool_call(self): + """An enum value holding PII is redacted, and the model's use of the stand-in is + restored in its tool arguments, so the call still carries a value the schema allows.""" + guardrail = _guardrail(event_hook=["pre_call", "post_call"]) + shield = _FakeShield({"[EMAIL_1]": "ops@example.com"}) + redact_mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + data = { + "messages": [], + "tools": [ + { + "type": "function", + "function": { + "name": "notify", + "parameters": {"properties": {"to": {"type": "string", "enum": ["ops@example.com"]}}}, + }, + } + ], + } + + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + assert redact_mock.call_args_list[0].kwargs["json"]["texts"] == ["ops@example.com"] + assert data["tools"][0]["function"]["parameters"]["properties"]["to"]["enum"] == ["[EMAIL_1]"] + + guardrail.async_handler.post = shield.post # type: ignore[method-assign] + call = SimpleNamespace(function=SimpleNamespace(name="notify", arguments='{"to": "[EMAIL_1]"}')) + reply = ModelResponse(choices=[Choices(message=Message(content=None, tool_calls=None))]) + reply.choices[0].message.tool_calls = [call] + restored = await guardrail.async_post_call_success_hook(data=data, user_api_key_dict=None, response=reply) + + assert json.loads(restored.choices[0].message.tool_calls[0].function.arguments) == {"to": "ops@example.com"} + + def test_schema_nesting_past_the_bound_is_refused(self): + schema: dict = {"type": "object", "description": "past-the-bound@example.com"} + for _ in range(100): + schema = {"type": "object", "properties": {"next": schema}} + data = {"tools": [{"type": "function", "function": {"name": "f", "parameters": schema}}]} + + with pytest.raises(Exception, match="schema"): + LLMShieldProxyGuardrail._locate_request_texts(data) + + +class TestRestoration: + @pytest.mark.asyncio + async def test_openai_shape_is_restored(self): + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + response = ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content="[EMAIL_1]"))]) + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result.choices[0].message.content == "a@b.com" + + @pytest.mark.asyncio + async def test_responses_api_shape_is_restored(self): + """The Responses API reply carries output items, not choices. + + Measured against a live provider: once the request side was fixed the reply + came back still holding the placeholder, because this shape has no choices + to walk. + """ + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + response = SimpleNamespace(output=[SimpleNamespace(content=[{"type": "output_text", "text": "[EMAIL_1]"}])]) + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result.output[0].content[0]["text"] == "a@b.com" + + @pytest.mark.asyncio + async def test_responses_api_object_blocks_are_restored(self): + """Blocks arrive as objects too, depending on how far the reply is parsed.""" + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + block = SimpleNamespace(text="[EMAIL_1]") + response = SimpleNamespace(output=[SimpleNamespace(content=[block])]) + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result.output[0].content[0].text == "a@b.com" + + @pytest.mark.asyncio + async def test_anthropic_message_shape_is_restored(self): + """The /v1/messages reply is a plain dict with no choices. + + Measured against a live provider: without its own branch the reply went + back to the caller still carrying the placeholder, even though the + request had been redacted correctly. + """ + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + response = { + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "[EMAIL_1]"}], + } + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result["content"][0]["text"] == "a@b.com" + + @pytest.mark.asyncio + async def test_anthropic_non_text_blocks_are_left_alone(self): + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + response = { + "type": "message", + "role": "assistant", + "content": [ + {"type": "text", "text": "[EMAIL_1]"}, + {"type": "tool_use", "id": "t1", "name": "lookup", "input": {}}, + ], + } + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result["content"][0]["text"] == "a@b.com" + assert result["content"][1] == {"type": "tool_use", "id": "t1", "name": "lookup", "input": {}} + + +class TestVaultIsolation: + """The vault id must never be something a caller can choose. + + The vault holds the plaintext behind every placeholder. If a caller could name + the vault, they could send a placeholder, have the model echo it back, and get + another caller's value restored into their own reply. + """ + + @pytest.mark.asyncio + async def test_caller_supplied_session_id_is_not_used(self): + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = { + "messages": [{"role": "user", "content": "a@b.com"}], + "metadata": {"llm_shield_session_id": "victim-session"}, + "litellm_session_id": "victim-session", + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + used = mock.call_args_list[0].kwargs["headers"]["X-Session-ID"] + assert used != "victim-session" + assert data["litellm_metadata"]["llm_shield_session_id"] == used + + @pytest.mark.asyncio + async def test_session_id_is_not_forwarded_to_the_provider(self): + """The vault id is a capability, so it must stay out of provider-visible metadata. + + `metadata` is forwarded upstream on /v1/responses; `litellm_metadata` is not. A + provider holding both the placeholders and the session id could call the shield's + rehydrate endpoint and read back exactly what this guardrail withholds. + """ + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = {"messages": [{"role": "user", "content": "a@b.com"}], "metadata": {}} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + used = mock.call_args_list[0].kwargs["headers"]["X-Session-ID"] + assert "llm_shield_session_id" not in data["metadata"] + assert data["litellm_metadata"]["llm_shield_session_id"] == used + + @pytest.mark.asyncio + async def test_restore_ignores_a_foreign_session_id(self): + """A reply is left unrestored rather than resolved against another vault.""" + guardrail = _guardrail(event_hook="post_call") + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = {"metadata": {"llm_shield_session_id": "victim-session"}} + response = ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content="[EMAIL_1]"))]) + await guardrail.async_post_call_success_hook(data=data, user_api_key_dict=None, response=response) + + assert mock.call_args_list[0].kwargs["headers"]["X-Session-ID"] != "victim-session" + + @pytest.mark.asyncio + async def test_each_request_gets_its_own_vault(self): + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["[EMAIL_1]"]}) + + for _ in range(2): + await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"messages": [{"role": "user", "content": "a@b.com"}]}, + call_type="completion", + ) + + seen = {call.kwargs["headers"]["X-Session-ID"] for call in mock.call_args_list} + assert len(seen) == 2 + + + @pytest.mark.parametrize( + "data", + [ + pytest.param( + {"messages": [{"role": "system", "content": "S"}, {"role": "user", "content": "U"}]}, + id="system-turn", + ), + pytest.param( + {"messages": [{"role": "developer", "content": "S"}, {"role": "user", "content": "U"}]}, + id="developer-turn", + ), + pytest.param( + {"system": "S", "messages": [{"role": "user", "content": "U"}]}, + id="anthropic-top-level-system", + ), + pytest.param({"instructions": "S", "input": "U"}, id="responses-instructions"), + pytest.param( + { + "messages": [{"role": "user", "content": "U"}], + "tools": [{"type": "function", "function": {"name": "f", "description": "S"}}], + }, + id="chat-tool-description", + ), + pytest.param( + {"input": "U", "tools": [{"type": "function", "name": "f", "description": "S"}]}, + id="responses-tool-description", + ), + pytest.param( + {"messages": [{"role": "user", "content": "U"}], "functions": [{"name": "f", "description": "S"}]}, + id="legacy-function-description", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "U"}], + "tools": [{"name": "f", "input_schema": {"properties": {"to": {"description": "S"}}}}], + }, + id="anthropic-schema-description", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "U"}], + "response_format": { + "type": "json_schema", + "json_schema": {"name": "n", "description": "S", "schema": {"type": "object"}}, + }, + }, + id="chat-response-format", + ), + pytest.param( + { + "input": "U", + "text": { + "format": { + "type": "json_schema", + "name": "n", + "schema": {"properties": {"a": {"description": "S"}}}, + } + }, + }, + id="responses-text-format", + ), + pytest.param({"prediction": {"type": "content", "content": "U"}, "instructions": "S"}, id="prediction"), + pytest.param( + {"prediction": {"type": "content", "content": [{"type": "text", "text": "U"}]}, "instructions": "S"}, + id="prediction-parts", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "U"}], + "web_search_options": {"user_location": {"type": "approximate", "approximate": {"city": "S"}}}, + }, + id="chat-web-search-location", + ), + pytest.param( + { + "input": "U", + "tools": [{"type": "web_search", "user_location": {"type": "approximate", "region": "S"}}], + }, + id="responses-web-search-location", + ), + pytest.param({"messages": [{"role": "user", "content": "U"}], "user": "S"}, id="end-user-id"), + pytest.param({"input": "U", "safety_identifier": "S"}, id="safety-identifier"), + ], + ) + def test_server_authored_text_is_split_from_the_callers(self, data: dict) -> None: + """Every request shape must sort its server-authored spans out of the caller's.""" + caller, privileged = LLMShieldProxyGuardrail._locate_request_texts(data) + + assert [text for text, _ in caller] == ["U"] + assert [text for text, _ in privileged] == ["S"] + + @pytest.mark.asyncio + async def test_a_system_prompt_gets_a_vault_of_its_own(self) -> None: + """The reply is restored against the caller's vault, so the two cannot be one.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["[EMAIL_2]"]}) + + data = { + "messages": [ + {"role": "system", "content": "escalate to admin@corp.internal"}, + {"role": "user", "content": "email a@b.com"}, + ] + } + await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="completion" + ) + + privileged_id, caller_id = ( + call.kwargs["headers"]["X-Session-ID"] for call in mock.call_args_list + ) + assert privileged_id != caller_id + assert guardrail._session_id(data) == caller_id + + @pytest.mark.asyncio + async def test_the_system_prompt_vault_id_is_never_stored(self) -> None: + """Nothing can restore against the system vault later, because its id is not kept. + + This is what stops a caller from having the model echo a placeholder out of a + system prompt they cannot see and receiving the plaintext behind it. + """ + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["[EMAIL_2]"]}) + + data = { + "messages": [ + {"role": "system", "content": "escalate to admin@corp.internal"}, + {"role": "user", "content": "email a@b.com"}, + ] + } + await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="completion" + ) + + privileged_id = mock.call_args_list[0].kwargs["headers"]["X-Session-ID"] + assert privileged_id not in json.dumps(data, default=str) + + def test_responses_system_and_developer_items_are_privileged(self) -> None: + """Responses `input` carries system and developer turns as items, like Chat messages. + + In the caller's vault, a caller could have the model echo a placeholder out of a + system message they cannot see and receive the plaintext behind it. + """ + data = { + "input": [ + {"role": "system", "content": "S"}, + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "D"}]}, + {"role": "user", "content": "U"}, + ] + } + caller, privileged = LLMShieldProxyGuardrail._locate_request_texts(data) + + assert [text for text, _ in caller] == ["U"] + assert [text for text, _ in privileged] == ["S", "D"] + + +class TestFailClosed: + @pytest.mark.asyncio + async def test_unreachable_shield_blocks_the_request(self): + """Failing open would send the PII upstream, defeating the guardrail.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(side_effect=ConnectionError("refused")) + + with pytest.raises(GuardrailRaisedException): + await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"messages": [{"role": "user", "content": "a@b.com"}]}, + call_type="completion", + ) + + @pytest.mark.asyncio + async def test_error_status_blocks_the_request(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_response({"error": "nope"}, status_code=500)) + + with pytest.raises(GuardrailRaisedException): + await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"messages": [{"role": "user", "content": "a@b.com"}]}, + call_type="completion", + ) + + @pytest.mark.asyncio + async def test_short_payload_blocks_the_request(self): + """A response that loses an entry would silently misalign the write-back.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": []}) + + with pytest.raises(GuardrailRaisedException): + await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"messages": [{"role": "user", "content": "a@b.com"}]}, + call_type="completion", + ) + + +class TestStreamingRehydration: + @pytest.mark.asyncio + async def test_split_placeholder_is_not_emitted_in_fragments(self): + """The window holds back a partial placeholder and releases it once complete.""" + guardrail = _guardrail(event_hook="post_call") + # Shield holds "[EMAIL" back, then releases the restored value. + _mock_post( + guardrail, + {"text": "Email ", "carry": "[EMAIL"}, + {"text": "a@b.com about it", "carry": ""}, + ) + + async def stream(): + yield _chunk("Email [EMAIL") + yield _chunk("_1] about it", finish_reason="stop") + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + emitted = [c.choices[0].delta.content for c in chunks] + assert emitted == ["Email ", "a@b.com about it"] + # No fragment of the placeholder ever reached the client. + assert not any("[EMAIL" in (text or "") for text in emitted) + + @pytest.mark.asyncio + async def test_carry_is_returned_to_the_next_call(self): + guardrail = _guardrail(event_hook="post_call") + mock = _mock_post( + guardrail, + {"text": "", "carry": "hold"}, + {"text": "held-and-more", "carry": ""}, + ) + + async def stream(): + yield _chunk("hold") + yield _chunk("-and-more", finish_reason="stop") + + await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + assert mock.call_args_list[0].kwargs["json"]["carry"] == "" + assert mock.call_args_list[1].kwargs["json"]["carry"] == "hold" + assert mock.call_args_list[1].kwargs["json"]["final"] is True + + @pytest.mark.asyncio + async def test_every_choice_is_restored(self): + """With n>1 a later choice must not be handed back still holding a placeholder.""" + guardrail = _guardrail(event_hook="post_call") + _mock_post( + guardrail, + {"text": "first@example.com", "carry": ""}, + {"text": "second@example.com", "carry": ""}, + ) + + async def stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="[EMAIL_1]"), finish_reason="stop"), + StreamingChoices(index=1, delta=Delta(content="[EMAIL_2]"), finish_reason="stop"), + ] + ) + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + restored = [choice.delta.content for choice in chunks[0].choices] + assert restored == ["first@example.com", "second@example.com"] + + @pytest.mark.asyncio + async def test_choice_windows_do_not_cross_contaminate(self): + """Each choice is its own token stream, so each carries its own window. + + One shared window would send the characters held back for choice 0 up + against choice 1's next delta and splice the two streams together. + """ + guardrail = _guardrail(event_hook="post_call") + mock = _mock_post( + guardrail, + {"text": "", "carry": "A-held"}, + {"text": "", "carry": "B-held"}, + {"text": "a-done", "carry": ""}, + {"text": "b-done", "carry": ""}, + ) + + async def stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="a1")), + StreamingChoices(index=1, delta=Delta(content="b1")), + ] + ) + yield ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="a2"), finish_reason="stop"), + StreamingChoices(index=1, delta=Delta(content="b2"), finish_reason="stop"), + ] + ) + + await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + sent = [call.kwargs["json"] for call in mock.call_args_list] + assert sent[2]["carry"] == "A-held", "choice 0 must get its own window back" + assert sent[3]["carry"] == "B-held", "choice 1 must get its own window back" + + @pytest.mark.asyncio + async def test_a_choice_missing_from_the_last_chunk_still_flushes(self): + """Held text must not be dropped because its choice ended earlier. + + Choice 1 finishes and stops appearing, then the stream ends without a + finish_reason for choice 0. Flushing only the terminal chunk's choices would + discard whatever choice 1 was still holding and truncate its answer. + """ + guardrail = _guardrail(event_hook="post_call") + _mock_post( + guardrail, + {"text": "", "carry": "held-0"}, + {"text": "", "carry": "held-1"}, + {"text": "zero-done", "carry": ""}, + {"text": "one-done", "carry": ""}, + ) + + async def stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="a")), + StreamingChoices(index=1, delta=Delta(content="b")), + ] + ) + yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=None))]) + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + flushed = { + choice.index: choice.delta.content for chunk in chunks for choice in chunk.choices if choice.delta.content + } + assert flushed.get(1) == "one-done", "choice 1's held text was dropped" + assert flushed.get(0) == "zero-done" + + @pytest.mark.asyncio + async def test_held_tool_arguments_land_in_the_finishing_chunk(self): + """A client parses tool arguments on finish_reason, so the flush must ride that chunk. + + The finishing chunk also carries its own fragment for the same tool call. That + entry has to survive, with the held text appended after it as a continuation. + """ + guardrail = _guardrail(event_hook="post_call") + _mock_post( + guardrail, + {"text": '{"to": "', "carry": "[EMA"}, + {"text": "", "carry": '[EMAIL_1]"}'}, + {"text": 'a@example.com"}', "carry": ""}, + ) + + async def stream(): + yield _tool_chunk('{"to": "[EMA') + yield _tool_chunk('IL_1]"}', finish_reason="tool_calls") + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + assert len(chunks) == 2, "the flush must not arrive after the finish_reason chunk" + final_calls = chunks[1].choices[0].delta.tool_calls + assert len(final_calls) == 2, "the finishing chunk's own fragment was dropped" + assert _field(final_calls[1], "index") == 0 + arguments = "".join( + _field(_field(call, "function"), "arguments") or "" + for chunk in chunks + for call in chunk.choices[0].delta.tool_calls + ) + assert json.loads(arguments) == {"to": "a@example.com"} + + @pytest.mark.asyncio + async def test_held_tool_arguments_flush_when_the_stream_ends_unfinished(self): + """No finish_reason at all: a trailing chunk carries the held arguments alone.""" + guardrail = _guardrail(event_hook="post_call") + _mock_post( + guardrail, + {"text": '{"to": "', "carry": "[EMA"}, + {"text": 'a@example.com"}', "carry": ""}, + ) + + async def stream(): + yield _tool_chunk('{"to": "[EMA') + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + assert len(chunks) == 2 + trailing = chunks[1].choices[0].delta.tool_calls + assert trailing == [{"index": 0, "function": {"arguments": 'a@example.com"}'}}], ( + "the copied chunk's own fragment was already delivered and must not repeat" + ) + + @pytest.mark.asyncio + async def test_chunks_are_forwarded_as_they_arrive(self): + """Restoration must not buffer the stream into a single terminal chunk.""" + guardrail = _guardrail(event_hook="post_call") + _mock_post( + guardrail, + {"text": "one ", "carry": ""}, + {"text": "two ", "carry": ""}, + {"text": "three", "carry": ""}, + ) + + async def stream(): + yield _chunk("one ") + yield _chunk("two ") + yield _chunk("three", finish_reason="stop") + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + assert len(chunks) == 3 + assert [c.choices[0].delta.content for c in chunks] == ["one ", "two ", "three"] + + +class TestApplyGuardrailToolCalls: + """The unified entry point the UI's Test button and the translation handlers use.""" + + @pytest.mark.asyncio + async def test_response_tool_call_arguments_are_rehydrated(self): + """Regression: this path deep-copied tool calls with `copy` never imported. + + 47 tests passed with a guaranteed NameError here, because every tool-call test + covered the request side and this is the only path that reaches the copy. + """ + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["hi", '{"email": "a@b.com"}']}) + + data = {"litellm_metadata": {"llm_shield_session_id": "shield-abc"}} + inputs = { + "texts": ["hi"], + "tool_calls": [{"function": {"name": "send", "arguments": '{"email": "[EMAIL_1]"}'}}], + } + + merged = await guardrail.apply_guardrail(inputs=inputs, request_data=data, input_type="response") + + assert merged["tool_calls"][0]["function"]["arguments"] == '{"email": "a@b.com"}' + assert inputs["tool_calls"][0]["function"]["arguments"] == '{"email": "[EMAIL_1]"}' + + +class TestAnthropicStreamRestoration: + """/v1/messages streams reach the hook as raw SSE frames, with no `choices` to walk.""" + + VAULT = {"[EMAIL_1]": "a@example.com"} + + @pytest.mark.asyncio + async def test_split_placeholder_is_restored_and_never_fragmented(self): + guardrail, _ = _shielded(self.VAULT) + + out = await _restore_stream(guardrail, _text_block_stream("Mail [EMA", "IL_1] now")) + + deltas = [e["delta"]["text"] for e in _sse_events(out) if e["type"] == "content_block_delta"] + assert "".join(deltas) == "Mail a@example.com now" + assert not any("[EMA" in d for d in deltas), "a placeholder fragment reached the client" + + @pytest.mark.asyncio + async def test_held_text_lands_before_its_block_stops(self): + """A trailing `[` that never became a placeholder is still part of the answer.""" + guardrail, _ = _shielded(self.VAULT) + + out = await _restore_stream(guardrail, _text_block_stream("Mail [EMAIL_1], x = a[")) + + types = [e["type"] for e in _sse_events(out)] + deltas = [e["delta"]["text"] for e in _sse_events(out) if e["type"] == "content_block_delta"] + assert "".join(deltas) == "Mail a@example.com, x = a[" + assert types.index("content_block_stop") > max(i for i, t in enumerate(types) if t == "content_block_delta") + + @pytest.mark.asyncio + async def test_events_split_across_network_chunks_are_restored(self): + """A chunk can end mid-event; the frame is parsed once it is whole.""" + guardrail, _ = _shielded(self.VAULT) + raw = b"".join(_text_block_stream("Mail [EMA", "IL_1] now")) + + out = await _restore_stream(guardrail, [raw[i : i + 7] for i in range(0, len(raw), 7)]) + + deltas = [e["delta"]["text"] for e in _sse_events(out) if e["type"] == "content_block_delta"] + assert "".join(deltas) == "Mail a@example.com now" + + @pytest.mark.asyncio + async def test_str_frames_stay_str(self): + guardrail, _ = _shielded(self.VAULT) + + out = await _restore_stream(guardrail, [frame.decode() for frame in _text_block_stream("[EMAIL_1]")]) + + assert all(isinstance(frame, str) for frame in out) + assert "a@example.com" in "".join(out) + + @pytest.mark.asyncio + async def test_tool_input_json_is_restored(self): + guardrail, _ = _shielded(self.VAULT) + frames = [ + _sse({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "input": {}}}), + *( + _sse( + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": p}, + } + ) + for p in ('{"to": "[EMAI', 'L_1]"}') + ), + _sse({"type": "content_block_stop", "index": 1}), + ] + + out = await _restore_stream(guardrail, frames) + + partial = "".join(e["delta"]["partial_json"] for e in _sse_events(out) if e["type"] == "content_block_delta") + assert json.loads(partial) == {"to": "a@example.com"} + + @pytest.mark.asyncio + async def test_signed_thinking_and_foreign_frames_pass_through_byte_for_byte(self): + """Rewriting a signed thinking block breaks it; other frames are not ours to touch.""" + guardrail, shield = _shielded(self.VAULT) + thinking = {"type": "thinking_delta", "thinking": "[EMAIL_1]"} + frames = [ + _sse({"type": "content_block_delta", "index": 0, "delta": thinking}), + b'data: {"candidates": [{"content": {"parts": [{"text": "[EMAIL_1]"}]}}]}\n\n', + b"data: not json\n\n", + ] + + out = await _restore_stream(guardrail, frames) + + assert b"".join(out) == b"".join(frames) + assert shield.urls == [] + + @pytest.mark.asyncio + async def test_a_raw_stream_that_is_not_sse_is_never_buffered(self): + """Without event boundaries to wait for, buffering would hold the whole reply.""" + guardrail, _ = _shielded(self.VAULT) + chunks = [b'[{"candidates": []}', b', {"candidates": []}]'] + + out = await _restore_stream(guardrail, chunks) + + assert out == chunks + + @pytest.mark.asyncio + @pytest.mark.parametrize("cut", [1, 3, 5, 6]) + async def test_a_field_name_split_by_the_first_chunk_still_reads_as_sse(self, cut: int): + """`b"eve"` then `b"nt: ..."` is still SSE; deciding on the first chunk alone + would pass the whole stream through with its placeholders.""" + guardrail, _ = _shielded(self.VAULT) + raw = b"".join(_text_block_stream("Mail [EMAIL_1]")) + + out = await _restore_stream(guardrail, [raw[:cut], raw[cut:]]) + + deltas = [e["delta"]["text"] for e in _sse_events(out) if e["type"] == "content_block_delta"] + assert "".join(deltas) == "Mail a@example.com" + + +class TestResponsesStreamRestoration: + """/v1/responses streams are typed events, with no `choices` to walk.""" + + VAULT = {"[EMAIL_1]": "a@example.com"} + + @staticmethod + def _text_delta(delta: str, sequence_number: int, content_index: int = 0) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=content_index, + delta=delta, + sequence_number=sequence_number, + ) + + @pytest.mark.asyncio + async def test_deltas_and_done_text_are_restored(self): + guardrail, _ = _shielded(self.VAULT) + done = OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id="msg_1", + output_index=0, + content_index=0, + text="Mail [EMAIL_1] x[", + ) + + out = await _restore_stream( + guardrail, [self._text_delta("Mail [EMA", 1), self._text_delta("IL_1] x[", 2), done] + ) + + deltas = [e.delta for e in out if isinstance(e, OutputTextDeltaEvent)] + assert not any("[EMA" in d for d in deltas), "a placeholder fragment reached the client" + assert "".join(deltas) == "Mail a@example.com x[" + assert out[-1].text == "Mail a@example.com x[", "the done event repeats the full, restored text" + assert isinstance(out[-2], OutputTextDeltaEvent), "held text must land before the done event" + + @pytest.mark.asyncio + async def test_function_call_arguments_are_restored(self): + guardrail, _ = _shielded(self.VAULT) + events = [ + FunctionCallArgumentsDeltaEvent( + type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + item_id="fc_1", + output_index=1, + delta=part, + ) + for part in ('{"to": "[EMAI', 'L_1]"}') + ] + + out = await _restore_stream(guardrail, events) + + assert json.loads("".join(e.delta for e in out)) == {"to": "a@example.com"} + + @pytest.mark.asyncio + async def test_a_truncated_stream_still_flushes(self): + """No done event at all: whatever the window holds goes out at the end.""" + guardrail, _ = _shielded(self.VAULT) + + out = await _restore_stream(guardrail, [self._text_delta("see [EMAIL_1] a[", 1)]) + + assert "".join(e.delta for e in out) == "see a@example.com a[" + + @pytest.mark.asyncio + async def test_completed_response_is_restored(self): + """The terminal event repeats the whole reply, and clients read it as the answer.""" + guardrail, _ = _shielded(self.VAULT) + block = {"type": "output_text", "text": "Mail [EMAIL_1]"} + call = SimpleNamespace(type="function_call", arguments='{"to": "[EMAIL_1]"}') + completed = SimpleNamespace( + type="response.completed", + response=SimpleNamespace(output=[SimpleNamespace(content=[block]), call]), + ) + + (out,) = await _restore_stream(guardrail, [completed]) + + restored_block, restored_call = out.response.output[0].content[0], out.response.output[1] + assert restored_block["text"] == "Mail a@example.com" + assert restored_call.arguments == '{"to": "a@example.com"}' + + @pytest.mark.asyncio + async def test_streams_on_different_parts_do_not_share_a_window(self): + guardrail, _ = _shielded(self.VAULT) + events = [ + self._text_delta("one [EMA", 1), + self._text_delta("two", 2, content_index=1), + self._text_delta("IL_1]", 3), + ] + + out = await _restore_stream(guardrail, events) + + by_part: dict[int, str] = {} + for event in out: + by_part[event.content_index] = by_part.get(event.content_index, "") + event.delta + assert by_part == {0: "one a@example.com", 1: "two"} + + @pytest.mark.asyncio + async def test_reasoning_summary_part_done_is_restored(self): + """The summary part repeats the whole summary text after its deltas.""" + guardrail, _ = _shielded(self.VAULT) + part = SimpleNamespace(type="summary_text", text="asked about [EMAIL_1]") + event = SimpleNamespace( + type="response.reasoning_summary_part.done", item_id="rs_1", output_index=0, summary_index=0, part=part + ) + + (out,) = await _restore_stream(guardrail, [event]) + + assert out.part.text == "asked about a@example.com" + + @pytest.mark.asyncio + async def test_mcp_call_arguments_are_restored(self): + """A stream family outside the chat-era set: matched by shape, not by name.""" + guardrail, _ = _shielded(self.VAULT) + deltas = [ + {"type": "response.mcp_call_arguments.delta", "item_id": "mcp_1", "output_index": 0, "delta": d} + for d in ('{"to": "[EMAI', 'L_1]"}') + ] + done = { + "type": "response.mcp_call_arguments.done", + "item_id": "mcp_1", + "output_index": 0, + "arguments": '{"to": "[EMAIL_1]"}', + } + + out = await _restore_stream(guardrail, [*deltas, done]) + + assert json.loads("".join(e["delta"] for e in out[:-1])) == {"to": "a@example.com"} + assert json.loads(out[-1]["arguments"]) == {"to": "a@example.com"} + assert out[-1]["item_id"] == "mcp_1", "identifiers are not text and stay as sent" + + @pytest.mark.asyncio + async def test_audio_deltas_are_not_sent_to_the_shield(self): + """Audio arrives base64-encoded; restoring it would cost a round trip for nothing.""" + guardrail, shield = _shielded(self.VAULT) + audio = {"type": "response.audio.delta", "item_id": "a_1", "output_index": 0, "delta": "UklGRiQAAABXQVZF"} + + out = await _restore_stream(guardrail, [audio]) + + assert out == [audio] + assert shield.urls == [] + + +class TestResponseCacheIsolation: + """The reply LiteLLM caches must keep its placeholders. + + Placeholders are numbered per request, so two callers' redacted requests can be + identical and share a cache key. LiteLLM keeps the provider's reply object -- for a + native Anthropic dict or an in-memory cache, the object itself -- so restoring it in + place would hand one caller's values to the next caller who hits that key. + """ + + @pytest.mark.asyncio + async def test_a_cache_hit_is_restored_against_the_new_callers_vault(self): + cache = InMemoryCache() + provider_reply = {"type": "message", "content": [{"type": "text", "text": "Repeat [EMAIL_1]"}]} + cache.set_cache("redacted-request", provider_reply) + + alice, _ = _shielded({"[EMAIL_1]": "alice@example.com"}) + to_alice = await alice.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=cache.get_cache("redacted-request") + ) + bob, _ = _shielded({"[EMAIL_1]": "bob@example.com"}) + to_bob = await bob.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=cache.get_cache("redacted-request") + ) + + assert to_alice["content"][0]["text"] == "Repeat alice@example.com" + assert to_bob["content"][0]["text"] == "Repeat bob@example.com" + assert cache.get_cache("redacted-request")["content"][0]["text"] == "Repeat [EMAIL_1]" + + @pytest.mark.asyncio + async def test_a_model_response_is_restored_as_a_copy(self): + """A reply as `acompletion` returns it, hidden params and all.""" + reply = await litellm.acompletion( + model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], mock_response="Mail [EMAIL_1]" + ) + guardrail, _ = _shielded({"[EMAIL_1]": "a@example.com"}) + + restored = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=reply + ) + + assert restored.choices[0].message.content == "Mail a@example.com" + assert reply.choices[0].message.content == "Mail [EMAIL_1]" + + @pytest.mark.asyncio + async def test_responses_reply_is_restored_as_a_copy(self): + guardrail, _ = _shielded({"[EMAIL_1]": "a@example.com"}) + block = {"type": "output_text", "text": "Mail [EMAIL_1]"} + reply = SimpleNamespace(output=[SimpleNamespace(content=[block])]) + + restored = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=reply + ) + + assert restored.output[0].content[0]["text"] == "Mail a@example.com" + assert block["text"] == "Mail [EMAIL_1]" + + @pytest.mark.asyncio + async def test_stream_chunks_are_restored_as_copies(self): + """LiteLLM assembles the reply it caches from the chunks it yielded.""" + guardrail, _ = _shielded({"[EMAIL_1]": "a@example.com"}) + chunk = _chunk("Mail [EMAIL_1]", finish_reason="stop") + event = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="Mail [EMAIL_1]", + ) + + out = await _restore_stream(guardrail, [chunk, event]) + + assert out[0].choices[0].delta.content == "Mail a@example.com" + assert chunk.choices[0].delta.content == "Mail [EMAIL_1]" + assert "Mail [EMAIL_1]" == event.delta + + +class TestCompletionsRestoration: + """`/v1/completions` replies carry their text on the choice, with no message or delta.""" + + VAULT = {"[EMAIL_1]": "a@example.com"} + + @pytest.mark.asyncio + async def test_completions_reply_is_restored(self): + guardrail, _ = _shielded(self.VAULT) + reply = TextCompletionResponse(choices=[TextChoices(index=0, text="Mail [EMAIL_1]", finish_reason="stop")]) + + restored = await guardrail.async_post_call_success_hook( + data={"prompt": "x"}, user_api_key_dict=None, response=reply + ) + + assert restored.choices[0].text == "Mail a@example.com" + + @pytest.mark.asyncio + async def test_completions_stream_is_restored_across_chunks(self): + guardrail, _ = _shielded(self.VAULT) + chunks = [ + TextCompletionResponse(choices=[TextChoices(index=0, text="Mail [EMAI")]), + TextCompletionResponse(choices=[TextChoices(index=0, text="L_1] now", finish_reason="stop")]), + ] + + out = await _restore_stream(guardrail, chunks) + + assert "".join(chunk.choices[0].text for chunk in out) == "Mail a@example.com now" + + @pytest.mark.asyncio + async def test_completions_stream_without_finish_reason_is_flushed(self): + guardrail, _ = _shielded(self.VAULT) + chunks = [TextCompletionResponse(choices=[TextChoices(index=0, text="Mail [EMAIL_1")])] + + out = await _restore_stream(guardrail, chunks) + + assert "".join(chunk.choices[0].text or "" for chunk in out) == "Mail [EMAIL_1" + + +class TestProxyWiring: + def test_dashboard_config_model_is_exposed(self): + """The guardrail garden reads the provider's fields from `get_config_model`.""" + model = LLMShieldProxyGuardrail.get_config_model() + + assert model is not None + assert {"api_key", "api_base"} <= set(model.model_fields) + + @pytest.mark.asyncio + async def test_the_deployment_hook_leaves_a_proxy_reply_for_the_proxy_hook(self): + """Inside the proxy the deployment hook must not restore: LiteLLM caches what it returns. + + A proxy request was redacted by the proxy's pre-call hook, so it carries no + deployment-restore marker, and the proxy's post-call hook restores it after the cache + write. A caller-sent marker that does not match the minted vault id is ignored. + """ + guardrail, shield = _shielded({"[EMAIL_1]": "a@example.com"}) + reply = ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content="[EMAIL_1]"))]) + data = {"messages": [], "guardrails": [GUARDRAIL_NAME], "litellm_metadata": {}} + LLMShieldProxyGuardrail._mint_session_id(data) + data["litellm_metadata"]["llm_shield_restore_at_deployment"] = "caller-chosen" + + result = await guardrail.async_post_call_success_deployment_hook( + request_data=data, response=reply, call_type=None + ) + + assert result is None + assert reply.choices[0].message.content == "[EMAIL_1]" + assert shield.urls == [] + + @pytest.mark.asyncio + async def test_model_level_use_outside_the_proxy_is_restored_and_never_cached(self, monkeypatch): + """SDK use with model-level `guardrails`: the deployment hooks are the only redact and + restore steps, so the reply is restored there, and the request bypasses the cache -- + its key is built from the redacted request, and a cache hit would skip restoration. + """ + vault = {"[EMAIL_1]": "alice@example.com"} + + async def shield(url: str, headers: dict, json: dict, timeout: float) -> Response: + texts = json["texts"] + if url.endswith("/redact"): + return _response({"texts": [t.replace("alice@example.com", "[EMAIL_1]") for t in texts]}) + restored = [] + for text in texts: + for placeholder, original in vault.items(): + text = text.replace(placeholder, original) + restored.append(text) + return _response({"texts": restored}) + + guardrail = _guardrail(event_hook=["pre_call", "post_call"], default_on=False) + guardrail.async_handler.post = shield # type: ignore[method-assign] + cache = InMemoryCache() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + monkeypatch.setattr(litellm.cache, "cache", cache) + + reply = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Repeat alice@example.com"}], + mock_response="Repeat [EMAIL_1]", + guardrails=[GUARDRAIL_NAME], + ) + + # LiteLLM writes the cache from background tasks; let them land before looking. + await asyncio.gather(*_PENDING_CACHE_WRITES) + + assert reply.choices[0].message.content == "Repeat alice@example.com" + assert cache.cache_dict == {}, "the redacted request's reply must not be cached" + + @pytest.mark.asyncio + async def test_model_level_streaming_outside_the_proxy_is_refused(self, monkeypatch): + """No hook restores an SDK stream, and its cache writer misses the bypass, so it fails closed.""" + guardrail = _guardrail(event_hook=["pre_call", "post_call"], default_on=False) + _mock_post(guardrail, {"texts": ["Repeat [EMAIL_1]"]}) + cache = InMemoryCache() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + monkeypatch.setattr(litellm.cache, "cache", cache) + + with pytest.raises(GuardrailRaisedException, match="cannot restore a streamed reply"): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Repeat alice@example.com"}], + mock_response="Repeat [EMAIL_1]", + stream=True, + guardrails=[GUARDRAIL_NAME], + ) + await asyncio.gather(*_PENDING_CACHE_WRITES) + + assert cache.cache_dict == {} + + @pytest.mark.asyncio + async def test_restored_values_are_not_recorded_as_guardrail_telemetry(self): + """Guardrail logging is exported to traces even with message logging off, so the + restored reply must not land in it.""" + guardrail, _ = _shielded({"[EMAIL_1]": "alice@example.com"}) + reply = ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content="[EMAIL_1]"))]) + data = {"messages": [], "metadata": {}} + + restored = await guardrail.async_post_call_success_hook(data=data, user_api_key_dict=None, response=reply) + + assert restored.choices[0].message.content == "alice@example.com" + assert "alice@example.com" not in json.dumps(data, default=str) + + +class TestStreamUsage: + @pytest.mark.asyncio + async def test_trailing_flush_does_not_repeat_usage(self): + """With n>=2 and include_usage, the last chunk carries usage; a flush copied from it must not. + + Any consumer that sums usage chunks would otherwise count the request twice. + """ + guardrail, _ = _shielded({"[EMAIL_1]": "a@example.com"}) + both = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="Mail [EMAI")), + StreamingChoices(index=1, delta=Delta(content="Call [EMAI")), + ] + ) + usage_chunk = ModelResponseStream( + choices=[StreamingChoices(index=1, delta=Delta(content=None))], + usage=Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12), + ) + + out = await _restore_stream(guardrail, [both, usage_chunk]) + + with_usage = [chunk for chunk in out if getattr(chunk, "usage", None) is not None] + assert with_usage == [out[1]], "only the provider's own usage chunk carries usage" + flushed = out[2:] + assert flushed, "the held-back text is flushed at end of stream" + assert sorted(chunk.choices[0].index for chunk in flushed) == [0, 1] + assert [chunk.choices[0].delta.content for chunk in flushed] == ["[EMAI", "[EMAI"] diff --git a/ui/litellm-dashboard/public/assets/logos/llm_shield_proxy.svg b/ui/litellm-dashboard/public/assets/logos/llm_shield_proxy.svg new file mode 100644 index 00000000000..0dd78b078c9 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/llm_shield_proxy.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx index 29df7c8bf3d..a02cae097a7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx @@ -73,7 +73,9 @@ interface GuardrailPreset { provider: string; categoryName?: string; guardrailNameSuggestion: string; - mode: string; + // A guardrail that both rewrites the request and repairs the response needs two + // modes seeded, not one; the form already normalises either shape. + mode: string | string[]; defaultOn: boolean; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index 10ca58294b9..684049c54a9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -2,7 +2,9 @@ export interface GuardrailPreset { provider: string; categoryName?: string; guardrailNameSuggestion: string; - mode: string; + // A guardrail that both rewrites the request and repairs the response needs two + // modes seeded, not one; the form already normalises either shape. + mode: string | string[]; defaultOn: boolean; } @@ -325,6 +327,14 @@ export const GUARDRAIL_PRESETS: Record = { // MCP-only: default_on is the only activation path on the MCP hook defaultOn: true, }, + llm_shield_proxy: { + provider: "LLM Shield Proxy", + guardrailNameSuggestion: "LLM Shield Proxy", + // Both halves are required. With only pre_call the request is redacted and the + // placeholders are handed straight back to the caller. + mode: ["pre_call", "post_call"], + defaultOn: false, + }, conduct: { provider: "Conduct", guardrailNameSuggestion: "Conduct Guard", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts index a27dd344c95..e756c8c3e58 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts @@ -29,6 +29,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record = { straiker: "straiker.svg", alice: "alice.svg", agent_365: "microsoft_azure.svg", + llm_shield_proxy: "llm_shield_proxy.svg", conduct: "conduct.png", }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index d88a333d6f1..d1ead5f589a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -484,6 +484,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Agentic", "MCP", "Tool Misuse", "Observability"], providerKey: "Agent365", }, + { + id: "llm_shield_proxy", + name: "LLM Shield Proxy", + description: + "Self-hosted PII redaction that puts the original values back into the model's response, so the provider never receives personal data while the end user still sees it.", + category: "partner", + logo: guardrailLogoMap["LLM Shield Proxy"], + tags: ["PII", "Data Privacy", "Compliance", "Streaming"], + providerKey: "LLM Shield Proxy", + }, { id: "conduct", name: "Conduct Guard", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index 8df7dfb1403..09e0f21c670 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -1,6 +1,7 @@ import aimSecurityLogo from "../../../../../public/assets/logos/aim_security.jpeg"; import aktoLogo from "../../../../../public/assets/logos/akto.svg"; import aliceLogo from "../../../../../public/assets/logos/alice.svg"; +import llmShieldProxyLogo from "../../../../../public/assets/logos/llm_shield_proxy.svg"; import conductLogo from "../../../../../public/assets/logos/conduct.png"; import aporiaLogo from "../../../../../public/assets/logos/aporia.png"; import bedrockLogo from "../../../../../public/assets/logos/bedrock.svg"; @@ -86,6 +87,7 @@ export const guardrail_provider_map: Record = { QostodianNexus: "qostodian_nexus", Repelloai: "repelloai", Alice: "alice", + "LLM Shield Proxy": "llm_shield_proxy", Conduct: "conduct", }; @@ -211,6 +213,7 @@ export const guardrailLogoMap = { Straiker: straikerLogo.src, Alice: aliceLogo.src, "Microsoft Agent 365": microsoftAzureLogo.src, + "LLM Shield Proxy": llmShieldProxyLogo.src, "Conduct Guard": conductLogo.src, } satisfies Record; From 9b6a6a0b71633a89466f2e8490bd5c49dab5a68d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 5 Oct 2026 17:37:52 +0000 Subject: [PATCH 003/191] refactor(tracing): generate existing HTTP request models from Rust schemas (#44591) * test(tracing): pin HTTP request compatibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(tracing): generate existing HTTP request models from Rust schemas Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(tracing): bind trace query params to generated request models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci(rust): raise the native wheel size gate to 48 MB Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * perf(tracing): read the trace list clock without a thread-pool dispatch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(traces): share the trace page-size bounds between schema and reader Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(ui): type trace request queries against the generated OpenAPI schema Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs: encode the trace contract boundary in AGENTS.md Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ui): format trace request aliases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/verify_linux_native_wheel.py | 2 +- .../crates/traces-cache/src/reader.rs | 3 +- litellm-rust/crates/traces/AGENTS.md | 5 + .../crates/traces/src/bin/export_schema.rs | 10 +- litellm-rust/crates/traces/src/lib.rs | 1 + litellm-rust/crates/traces/src/request.rs | 56 +++++ litellm-rust/crates/traces/src/schema.rs | 32 +++ .../crates/traces/tests/request_schema.rs | 93 ++++++++ litellm/proxy/tracing_endpoints.py | 43 ++-- litellm/rust_bridge/trace/AGENTS.md | 7 + .../rust_bridge/trace/generated/requests.py | 59 +++++ litellm/tracing/AGENTS.md | 1 + scripts/generate_trace_types.py | 11 +- scripts/trace_codegen/README.md | 8 +- .../schemas/traces/TraceDetailRequest.json | 29 +++ .../schemas/traces/TraceErrorPageRequest.json | 19 ++ .../schemas/traces/TraceListRequest.json | 33 +++ .../schemas/traces/TraceQueryRequest.json | 14 ++ .../schemas/traces/TraceSpanRequest.json | 11 + tests/unit/proxy/test_tracing_endpoints.py | 211 +++++++++++++++++- .../src/components/lens/traces/AGENTS.md | 3 + .../src/components/lens/traces/api.ts | 7 +- .../lens/traces/list/runSearch/runSql.ts | 3 +- .../lens/traces/requestTypes.test-d.ts | 21 ++ .../src/components/lens/traces/types.ts | 7 + .../src/components/networking.tsx | 19 +- 26 files changed, 666 insertions(+), 42 deletions(-) create mode 100644 litellm-rust/crates/traces/src/request.rs create mode 100644 litellm-rust/crates/traces/tests/request_schema.rs create mode 100644 litellm/rust_bridge/trace/AGENTS.md create mode 100644 litellm/rust_bridge/trace/generated/requests.py create mode 100644 scripts/trace_codegen/schemas/traces/TraceDetailRequest.json create mode 100644 scripts/trace_codegen/schemas/traces/TraceErrorPageRequest.json create mode 100644 scripts/trace_codegen/schemas/traces/TraceListRequest.json create mode 100644 scripts/trace_codegen/schemas/traces/TraceQueryRequest.json create mode 100644 scripts/trace_codegen/schemas/traces/TraceSpanRequest.json create mode 100644 ui/litellm-dashboard/src/components/lens/traces/AGENTS.md create mode 100644 ui/litellm-dashboard/src/components/lens/traces/requestTypes.test-d.ts diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 465918f5a81..4c4cb399dce 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -214,7 +214,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 45_000_000 + native_size_limit: Final = 48_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), diff --git a/litellm-rust/crates/traces-cache/src/reader.rs b/litellm-rust/crates/traces-cache/src/reader.rs index 5c68d6e58bd..b0f83a89118 100644 --- a/litellm-rust/crates/traces-cache/src/reader.rs +++ b/litellm-rust/crates/traces-cache/src/reader.rs @@ -15,6 +15,7 @@ use litellm_traces::{ ListTracesParams, ReadAccessParams, SpanDetailParams, SpanErrorParams, TraceIdentityParams, TraceSpansParams, }, + request::{TRACE_PAGE_SIZE_MAX, TRACE_PAGE_SIZE_MIN}, resolve_trace, to_ui_content, }; @@ -133,7 +134,7 @@ impl TraceReader { cursor: Option<&str>, page_size: u32, ) -> Result, ReadError> { - if !(1..=500).contains(&page_size) { + if !(u32::from(TRACE_PAGE_SIZE_MIN)..=u32::from(TRACE_PAGE_SIZE_MAX)).contains(&page_size) { return Err(ReadError::InvalidParameters); } let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else { diff --git a/litellm-rust/crates/traces/AGENTS.md b/litellm-rust/crates/traces/AGENTS.md index a9249f31aa8..f31750f8712 100644 --- a/litellm-rust/crates/traces/AGENTS.md +++ b/litellm-rust/crates/traces/AGENTS.md @@ -4,3 +4,8 @@ - Keep ClickHouse schema, row encoding and queries in `litellm-traces-clickhouse`; keep PyO3 conversion in `python-bridge` - Test decoding and normalization through the public API - Expose one top-level `Error` enum in `src/error.rs` for decoding and normalization failures +- Own the tracing HTTP contracts: request types in `src/request.rs` and response views exported from `src/schema.rs` +- Python and the dashboard consume them only through generated code: `uv run scripts/generate_trace_types.py` writes `litellm/rust_bridge/trace/generated/`, and `npm run gen:api` in `ui/litellm-dashboard` regenerates `schema.d.ts` from the proxy's OpenAPI +- Declare each bound once as a constant and read it from both the schema attribute and the runtime check +- GET request types accept unknown fields because the routes ignore unknown query parameters; body request types use `deny_unknown_fields` +- Changing a request or response shape changes the public API; ship it in its own behavior-change PR diff --git a/litellm-rust/crates/traces/src/bin/export_schema.rs b/litellm-rust/crates/traces/src/bin/export_schema.rs index 25d1250ef12..edf89fa514f 100644 --- a/litellm-rust/crates/traces/src/bin/export_schema.rs +++ b/litellm-rust/crates/traces/src/bin/export_schema.rs @@ -1,6 +1,8 @@ fn main() { - println!( - "{}", - serde_json::to_string_pretty(&litellm_traces::schema::schemas()).unwrap() - ); + let schemas = if std::env::args().nth(1).as_deref() == Some("--requests") { + litellm_traces::schema::request_schemas() + } else { + litellm_traces::schema::schemas() + }; + println!("{}", serde_json::to_string_pretty(&schemas).unwrap()); } diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index a44679a59f3..7a72183de2c 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -15,6 +15,7 @@ mod normalize; mod otlp; pub mod query; mod query_access; +pub mod request; mod resolve; #[cfg(feature = "schema")] pub mod schema; diff --git a/litellm-rust/crates/traces/src/request.rs b/litellm-rust/crates/traces/src/request.rs new file mode 100644 index 00000000000..0ade6112f02 --- /dev/null +++ b/litellm-rust/crates/traces/src/request.rs @@ -0,0 +1,56 @@ +pub const TRACE_PAGE_SIZE_MIN: u16 = 1; +pub const TRACE_PAGE_SIZE_MAX: u16 = 500; + +#[macro_rules_attribute::apply(request_type)] +#[derive(Clone, Debug)] +pub struct TraceListRequest { + /// Window start, unix ms. Default: 24h ago + #[serde(default)] + pub start_ms: Option, + /// Window end, unix ms. Default: now + #[serde(default)] + pub end_ms: Option, + #[serde(default)] + #[cfg_attr(feature = "schema", schemars(length(max = 512)))] + pub cursor: Option, +} + +#[macro_rules_attribute::apply(request_type)] +#[derive(Clone, Debug)] +pub struct TraceDetailRequest { + #[serde(default)] + pub trace_ref: String, + #[serde(default)] + #[cfg_attr(feature = "schema", schemars(length(max = 512)))] + pub cursor: Option, + #[serde(default)] + #[cfg_attr( + feature = "schema", + schemars(range(min = TRACE_PAGE_SIZE_MIN, max = TRACE_PAGE_SIZE_MAX)) + )] + pub page_size: Option, +} + +#[macro_rules_attribute::apply(request_type)] +#[derive(Clone, Debug)] +pub struct TraceSpanRequest { + #[serde(default)] + pub trace_ref: String, +} + +#[macro_rules_attribute::apply(request_type)] +#[derive(Clone, Debug)] +pub struct TraceErrorPageRequest { + #[serde(default)] + pub trace_ref: String, + #[serde(default)] + #[cfg_attr(feature = "schema", schemars(length(max = 512)))] + pub cursor: Option, +} + +#[macro_rules_attribute::apply(request_type)] +#[derive(Clone, Debug)] +#[serde(deny_unknown_fields)] +pub struct TraceQueryRequest { + pub sql: String, +} diff --git a/litellm-rust/crates/traces/src/schema.rs b/litellm-rust/crates/traces/src/schema.rs index cfa1d8e201e..56e166019d7 100644 --- a/litellm-rust/crates/traces/src/schema.rs +++ b/litellm-rust/crates/traces/src/schema.rs @@ -36,6 +36,13 @@ fn received() -> Schema { .into_root_schema_for::() } +fn requested() -> Schema { + SchemaSettings::draft2020_12() + .for_deserialize() + .into_generator() + .into_root_schema_for::() +} + fn emitted() -> Schema { SchemaSettings::draft2020_12() .for_serialize() @@ -58,3 +65,28 @@ pub fn schemas() -> BTreeMap<&'static str, Schema> { ("SpanErrorPage", emitted::()), ]) } + +pub fn request_schemas() -> BTreeMap<&'static str, Schema> { + BTreeMap::from([ + ( + "TraceListRequest", + requested::(), + ), + ( + "TraceDetailRequest", + requested::(), + ), + ( + "TraceSpanRequest", + requested::(), + ), + ( + "TraceErrorPageRequest", + requested::(), + ), + ( + "TraceQueryRequest", + requested::(), + ), + ]) +} diff --git a/litellm-rust/crates/traces/tests/request_schema.rs b/litellm-rust/crates/traces/tests/request_schema.rs new file mode 100644 index 00000000000..598f73021fb --- /dev/null +++ b/litellm-rust/crates/traces/tests/request_schema.rs @@ -0,0 +1,93 @@ +#![cfg(feature = "schema")] + +use litellm_traces::request::{ + TraceDetailRequest, TraceErrorPageRequest, TraceListRequest, TraceQueryRequest, + TraceSpanRequest, +}; +use litellm_traces::schema::request_schemas; +use rstest::rstest; +use serde_json::json; + +#[rstest] +#[case::list("TraceListRequest")] +#[case::detail("TraceDetailRequest")] +#[case::span("TraceSpanRequest")] +#[case::error_page("TraceErrorPageRequest")] +fn get_request_schemas_ignore_unknown_fields(#[case] name: &str) { + let schemas = request_schemas(); + let schema = serde_json::to_value(&schemas[name]).unwrap(); + assert_ne!(schema["additionalProperties"], false); +} + +#[rstest] +fn query_request_schema_rejects_unknown_fields() { + let schemas = request_schemas(); + let schema = serde_json::to_value(&schemas["TraceQueryRequest"]).unwrap(); + assert_eq!(schema["additionalProperties"], false); +} + +#[rstest] +fn request_schemas_preserve_explicit_constraints() { + let schemas = request_schemas(); + let detail = serde_json::to_value(&schemas["TraceDetailRequest"]).unwrap(); + let list = serde_json::to_value(&schemas["TraceListRequest"]).unwrap(); + let span = serde_json::to_value(&schemas["TraceSpanRequest"]).unwrap(); + let error_page = serde_json::to_value(&schemas["TraceErrorPageRequest"]).unwrap(); + let query = serde_json::to_value(&schemas["TraceQueryRequest"]).unwrap(); + + assert_eq!(detail["properties"]["page_size"]["minimum"], 1); + assert_eq!(detail["properties"]["page_size"]["maximum"], 500); + assert_eq!(list["properties"]["cursor"]["maxLength"], 512); + assert_eq!(detail["properties"]["cursor"]["maxLength"], 512); + assert_eq!(error_page["properties"]["cursor"]["maxLength"], 512); + assert!(list["properties"]["start_ms"].get("minimum").is_none()); + assert!(list["properties"]["start_ms"].get("maximum").is_none()); + assert!(list["properties"]["end_ms"].get("minimum").is_none()); + assert!(list["properties"]["end_ms"].get("maximum").is_none()); + assert_eq!(detail["properties"]["trace_ref"]["default"], ""); + assert_eq!(span["properties"]["trace_ref"]["default"], ""); + assert_eq!(error_page["properties"]["trace_ref"]["default"], ""); + assert_eq!(query["required"], json!(["sql"])); + assert_eq!( + list["properties"]["start_ms"]["description"], + "Window start, unix ms. Default: 24h ago" + ); + assert_eq!( + list["properties"]["end_ms"]["description"], + "Window end, unix ms. Default: now" + ); +} + +#[rstest] +fn request_models_deserialize_defaults_and_null_cursors() { + let list: TraceListRequest = serde_json::from_value(json!({})).unwrap(); + let detail: TraceDetailRequest = serde_json::from_value(json!({})).unwrap(); + let span: TraceSpanRequest = serde_json::from_value(json!({})).unwrap(); + let error_page: TraceErrorPageRequest = serde_json::from_value(json!({})).unwrap(); + + assert!(list.start_ms.is_none()); + assert!(list.end_ms.is_none()); + assert!(list.cursor.is_none()); + assert_eq!(detail.trace_ref, ""); + assert!(detail.cursor.is_none()); + assert!(detail.page_size.is_none()); + assert_eq!(span.trace_ref, ""); + assert_eq!(error_page.trace_ref, ""); + assert!(error_page.cursor.is_none()); + + let null_cursor: TraceListRequest = serde_json::from_value(json!({"cursor": null})).unwrap(); + assert!(null_cursor.cursor.is_none()); +} + +#[rstest] +fn get_request_models_ignore_unknown_fields_and_query_model_rejects_them() { + assert!(serde_json::from_value::(json!({"unknown": true})).is_ok()); + assert!(serde_json::from_value::(json!({"unknown": true})).is_ok()); + assert!(serde_json::from_value::(json!({"unknown": true})).is_ok()); + assert!(serde_json::from_value::(json!({"unknown": true})).is_ok()); + assert!( + serde_json::from_value::(json!({"sql": "SELECT 1", "unknown": true})) + .is_err() + ); + assert!(serde_json::from_value::(json!({})).is_err()); +} diff --git a/litellm/proxy/tracing_endpoints.py b/litellm/proxy/tracing_endpoints.py index 0d3723c45b8..6ec0c6da3e7 100644 --- a/litellm/proxy/tracing_endpoints.py +++ b/litellm/proxy/tracing_endpoints.py @@ -29,6 +29,13 @@ from litellm.proxy.common_utils.http_parsing_utils import is_otlp_trace_request from litellm.proxy.tracing_runtime import provide_receiver, require_receiver from litellm.rust_bridge.trace.errors import TraceChanged from litellm.rust_bridge.trace.generated.models import TraceQueryHelp +from litellm.rust_bridge.trace.generated.requests import ( + TraceDetailRequest, + TraceErrorPageRequest, + TraceListRequest, + TraceQueryRequest, + TraceSpanRequest, +) from litellm.rust_bridge.trace.generated.types import ( AllQueryScope, OwnedQueryScope, @@ -49,6 +56,10 @@ router = APIRouter(tags=["agent tracing"]) MS_PER_DAY: Final = 24 * 60 * 60 * 1000 +async def current_time_ms() -> int: + return int(time.time() * 1000) + + @dataclass(frozen=True, slots=True) class TraceAccessContext: receiver: TraceReceiver | None @@ -183,28 +194,21 @@ def read_failure(error: TraceChanged | ValueError | OverflowError | RuntimeError @router.get("/v1/traces", response_model=TracePage) async def list_agent_traces( context: Annotated[TraceAccessContext, Depends(provide_trace_access)], - start_ms: Annotated[int | None, Query(description="Window start, unix ms. Default: 24h ago")] = None, - end_ms: Annotated[int | None, Query(description="Window end, unix ms. Default: now")] = None, - cursor: Annotated[str | None, Query(max_length=512)] = None, + now_ms: Annotated[int, Depends(current_time_ms)], + request: Annotated[TraceListRequest, Query()], ) -> TracePage: - now_ms: Final = int(time.time() * 1000) try: tracing, scope = context.reader() return await tracing.list_traces( scope=scope, - start_ms=start_ms if start_ms is not None else now_ms - MS_PER_DAY, - end_ms=end_ms if end_ms is not None else now_ms, - cursor=cursor, + start_ms=request.start_ms if request.start_ms is not None else now_ms - MS_PER_DAY, + end_ms=request.end_ms if request.end_ms is not None else now_ms, + cursor=request.cursor, ) except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: raise read_failure(error) from error -class TraceQueryRequest(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - sql: str - - @dataclass(frozen=True, slots=True) class TraceQueryAccess: storage: ClickHouseStorage @@ -272,13 +276,11 @@ async def help_agent_trace_queries( async def get_agent_trace( trace_id: str, context: Annotated[TraceAccessContext, Depends(provide_trace_access)], - trace_ref: Annotated[str, Query()] = "", - cursor: Annotated[str | None, Query(max_length=512)] = None, - page_size: Annotated[int | None, Query(ge=1, le=500)] = None, + request: Annotated[TraceDetailRequest, Query()], ) -> Trace: tracing, scope = context.reader() try: - trace: Final = await tracing.get_trace(trace_id, scope, trace_ref, cursor, page_size) + trace: Final = await tracing.get_trace(trace_id, scope, request.trace_ref, request.cursor, request.page_size) except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: raise read_failure(error) from error if trace is None: @@ -291,11 +293,11 @@ async def get_agent_trace_span( trace_id: str, span_id: str, context: Annotated[TraceAccessContext, Depends(provide_trace_access)], - trace_ref: Annotated[str, Query()] = "", + request: Annotated[TraceSpanRequest, Query()], ) -> SpanDetail: tracing, scope = context.reader() try: - span: Final = await tracing.get_span(trace_id, span_id, scope, trace_ref) + span: Final = await tracing.get_span(trace_id, span_id, scope, request.trace_ref) except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: raise read_failure(error) from error if span is None: @@ -308,12 +310,11 @@ async def get_agent_trace_span_error( trace_id: str, span_id: str, context: Annotated[TraceAccessContext, Depends(provide_trace_access)], - trace_ref: Annotated[str, Query()] = "", - cursor: Annotated[str | None, Query(max_length=512)] = None, + request: Annotated[TraceErrorPageRequest, Query()], ) -> SpanErrorPage: try: tracing, scope = context.reader() - page: Final = await tracing.get_span_error(trace_id, span_id, scope, trace_ref, cursor) + page: Final = await tracing.get_span_error(trace_id, span_id, scope, request.trace_ref, request.cursor) except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: raise read_failure(error) from error if page is None: diff --git a/litellm/rust_bridge/trace/AGENTS.md b/litellm/rust_bridge/trace/AGENTS.md new file mode 100644 index 00000000000..5132b6ce74c --- /dev/null +++ b/litellm/rust_bridge/trace/AGENTS.md @@ -0,0 +1,7 @@ +# Trace contract boundary + +- `generated/` is output of `scripts/generate_trace_types.py` from the `litellm-traces` schemas; change the Rust type and regenerate, never edit these files +- CI runs the generator with `--check` and fails on drift +- `storage.py` validates every native result against the generated response models before returning it +- Keep the trace methods in `_native.pyi` matching `litellm-rust/crates/python-bridge/src/routes/traces.rs` +- Native trace methods take scalar arguments; moving them to the generated request types changes overflow errors, so it needs its own behavior-change PR diff --git a/litellm/rust_bridge/trace/generated/requests.py b/litellm/rust_bridge/trace/generated/requests.py new file mode 100644 index 00000000000..dd5304ea09b --- /dev/null +++ b/litellm/rust_bridge/trace/generated/requests.py @@ -0,0 +1,59 @@ +# @generated by scripts/generate_trace_types.py, do not edit + +from __future__ import annotations + +from typing import Annotated, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field + + +class TraceDetailRequest(BaseModel): + model_config = ConfigDict( + frozen=True, + ) + + trace_ref: str = "" + cursor: str | None = Field(None, max_length=512) + page_size: int | None = Field(None, ge=1, le=500) + + +class TraceErrorPageRequest(BaseModel): + model_config = ConfigDict( + frozen=True, + ) + + trace_ref: str = "" + cursor: str | None = Field(None, max_length=512) + + +class TraceListRequest(BaseModel): + model_config = ConfigDict( + frozen=True, + ) + + start_ms: int | None = Field(None, description="Window start, unix ms. Default: 24h ago") + end_ms: int | None = Field(None, description="Window end, unix ms. Default: now") + cursor: str | None = Field(None, max_length=512) + + +class TraceQueryRequest(BaseModel): + model_config = ConfigDict( + extra="forbid", + frozen=True, + ) + + sql: str + + +class TraceSpanRequest(BaseModel): + model_config = ConfigDict( + frozen=True, + ) + + trace_ref: str = "" + + +TraceWireRequests: TypeAlias = Annotated[ + TraceDetailRequest | TraceErrorPageRequest | TraceListRequest | TraceQueryRequest | TraceSpanRequest, + Field(..., title="TraceWireRequests"), +] diff --git a/litellm/tracing/AGENTS.md b/litellm/tracing/AGENTS.md index d71664c38eb..82c4d4ac8a1 100644 --- a/litellm/tracing/AGENTS.md +++ b/litellm/tracing/AGENTS.md @@ -4,3 +4,4 @@ - Use `litellm.rust_bridge.trace.storage.ClickHouseStorage` for ClickHouse; keep trace schema, SQL and encoding in `litellm-traces`, and generic transport in `litellm-storage-clickhouse` - Derive tenant fields from authentication and overwrite matching fields supplied by the exporter - Test confirmed writes, failures, tenant isolation and read behavior through public functions +- Trace routes in `litellm/proxy/tracing_endpoints.py` bind query parameters with `Annotated[, Query()]` and bodies with the generated model; never redeclare field constraints in `Query(...)` or a local model diff --git a/scripts/generate_trace_types.py b/scripts/generate_trace_types.py index 39c93dabf49..52a61452711 100644 --- a/scripts/generate_trace_types.py +++ b/scripts/generate_trace_types.py @@ -34,7 +34,7 @@ class GeneratorConfig(BaseModel): options: tuple[str, ...] -def export(crate: str) -> Mapping[str, Mapping[str, JsonValue]]: +def export(crate: str, extra_args: tuple[str, ...] = ()) -> Mapping[str, Mapping[str, JsonValue]]: result: Final = subprocess.run( ( "cargo", @@ -48,6 +48,7 @@ def export(crate: str) -> Mapping[str, Mapping[str, JsonValue]]: f"export-{crate}-schema", "--features", "schema", + *(("--", *extra_args) if extra_args else ()), ), check=True, stdout=subprocess.PIPE, @@ -163,8 +164,9 @@ def main() -> int: sys.stderr.write(f"requires datamodel-code-generator=={config.version}\n") return 1 domain: Final = export("traces") + requests: Final = export("traces", ("--requests",)) clickhouse: Final = export("traces-clickhouse") - exported: Final = tuple(schema_files(domain, clickhouse)) + exported: Final = tuple(schema_files(domain, clickhouse, requests)) schema_results: Final = tuple(publish(path, content, args.check) for path, content in exported) schema_set_matches: Final = reconcile_schemas(frozenset(path for path, _ in exported), args.check) with TemporaryDirectory(prefix="trace-codegen-") as temporary: @@ -176,9 +178,11 @@ def main() -> int: directory, config, ) + request_models: Final = generate(requests, "requests", directory, config) python_results: Final = ( publish(GENERATED / "types.py", types.read_text(), args.check), publish(GENERATED / "models.py", models.read_text(), args.check), + publish(GENERATED / "requests.py", request_models.read_text(), args.check), ) return 0 if all((schema_set_matches, *schema_results, *python_results)) else 1 @@ -186,8 +190,9 @@ def main() -> int: def schema_files( domain: Mapping[str, Mapping[str, JsonValue]], clickhouse: Mapping[str, Mapping[str, JsonValue]], + requests: Mapping[str, Mapping[str, JsonValue]], ) -> Iterator[tuple[Path, str]]: - for crate, schemas in (("traces", domain), ("traces-clickhouse", clickhouse)): + for crate, schemas in (("traces", domain), ("traces-clickhouse", clickhouse), ("traces", requests)): for name, schema in schemas.items(): yield TOOLING / "schemas" / crate / f"{name}.json", json.dumps(schema, indent=2, sort_keys=True) + "\n" diff --git a/scripts/trace_codegen/README.md b/scripts/trace_codegen/README.md index f9071001c78..bb4f511a569 100644 --- a/scripts/trace_codegen/README.md +++ b/scripts/trace_codegen/README.md @@ -1,9 +1,13 @@ -Run `uv run scripts/generate_trace_types.py` from the repository root to export Rust schemas and regenerate the Python trace contracts. Run the same command with `--check` to compare fresh output with the committed schemas and Python files +Run `uv run scripts/generate_trace_types.py` from the repository root to export Rust schemas and regenerate the Python trace contracts. Run `uv run scripts/generate_trace_types.py --check` to compare fresh output with the committed schemas and Python files The script pins datamodel-code-generator in its inline dependency metadata. Rust uses the workspace's locked Schemars version through each owning crate's optional `schema` feature. Neither tool is a Python runtime dependency -Each crate exports its own roots using JSON Schema 2020-12. Request parameters use Schemars' deserialization contract. Trace views and query help use its serialization contract. Lens rows use their ClickHouse deserialization schemas, including quoted numbers and numeric boolean flags +The `litellm-traces` Rust request types own the generated request models in `litellm/rust_bridge/trace/generated/requests.py`. The GET routes bind their query parameters directly to the generated models. GET request types allow unknown fields because existing clients' unknown query parameters are ignored. The SQL body model forbids extra fields + +Each crate exports its own roots using JSON Schema 2020-12. Request parameters use Schemars' deserialization contract and carry only explicitly declared constraints. Their schemas skip the integer-bounds transform because it would add i64 bounds to `start_ms` and `end_ms`, narrowing what Python accepts, and replace `page_size`'s explicit 1..500 range with 0..65535. Trace views and query help use the serialization contract. Lens rows use their ClickHouse deserialization schemas, including quoted numbers and numeric boolean flags The templates preserve tuple conversion, immutable tuple defaults, and bounded `ReadOnly` TypedDict fields. Pydantic models use the generator's frozen-model option and each schema's extra-field policy. ClickHouse numeric schemas select bounded, normalized Python scalar types through schema metadata consumed by the model template +Changing the public request contract requires a separate behavior-change PR + Edit the owning Rust contract, schema annotation, or generation configuration, then regenerate. Never edit `litellm/rust_bridge/trace/generated/` manually. The SQL response envelope remains handwritten in `queries.py` diff --git a/scripts/trace_codegen/schemas/traces/TraceDetailRequest.json b/scripts/trace_codegen/schemas/traces/TraceDetailRequest.json new file mode 100644 index 00000000000..60cb76dfedc --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/TraceDetailRequest.json @@ -0,0 +1,29 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "properties": { + "cursor": { + "default": null, + "maxLength": 512, + "type": [ + "string", + "null" + ] + }, + "page_size": { + "default": null, + "format": "uint16", + "maximum": 500, + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "trace_ref": { + "default": "", + "type": "string" + } + }, + "title": "TraceDetailRequest", + "type": "object" +} diff --git a/scripts/trace_codegen/schemas/traces/TraceErrorPageRequest.json b/scripts/trace_codegen/schemas/traces/TraceErrorPageRequest.json new file mode 100644 index 00000000000..950dab8e83a --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/TraceErrorPageRequest.json @@ -0,0 +1,19 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "properties": { + "cursor": { + "default": null, + "maxLength": 512, + "type": [ + "string", + "null" + ] + }, + "trace_ref": { + "default": "", + "type": "string" + } + }, + "title": "TraceErrorPageRequest", + "type": "object" +} diff --git a/scripts/trace_codegen/schemas/traces/TraceListRequest.json b/scripts/trace_codegen/schemas/traces/TraceListRequest.json new file mode 100644 index 00000000000..a19c3965c0b --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/TraceListRequest.json @@ -0,0 +1,33 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "properties": { + "cursor": { + "default": null, + "maxLength": 512, + "type": [ + "string", + "null" + ] + }, + "end_ms": { + "default": null, + "description": "Window end, unix ms. Default: now", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "start_ms": { + "default": null, + "description": "Window start, unix ms. Default: 24h ago", + "format": "int64", + "type": [ + "integer", + "null" + ] + } + }, + "title": "TraceListRequest", + "type": "object" +} diff --git a/scripts/trace_codegen/schemas/traces/TraceQueryRequest.json b/scripts/trace_codegen/schemas/traces/TraceQueryRequest.json new file mode 100644 index 00000000000..1361b89ad90 --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/TraceQueryRequest.json @@ -0,0 +1,14 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "additionalProperties": false, + "properties": { + "sql": { + "type": "string" + } + }, + "required": [ + "sql" + ], + "title": "TraceQueryRequest", + "type": "object" +} diff --git a/scripts/trace_codegen/schemas/traces/TraceSpanRequest.json b/scripts/trace_codegen/schemas/traces/TraceSpanRequest.json new file mode 100644 index 00000000000..4f5bb8fff69 --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/TraceSpanRequest.json @@ -0,0 +1,11 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "properties": { + "trace_ref": { + "default": "", + "type": "string" + } + }, + "title": "TraceSpanRequest", + "type": "object" +} diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index 516a0415545..11d63ad545c 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -5,12 +5,15 @@ Tests for the agent tracing endpoints (litellm/proxy/tracing_endpoints.py). from collections.abc import AsyncGenerator, Mapping from contextlib import asynccontextmanager from types import ModuleType -from typing import Final, Literal -from unittest.mock import AsyncMock, MagicMock +from typing import Final, Literal, TypedDict +from unittest.mock import AsyncMock, MagicMock, call import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from httpx import Response +from pydantic import TypeAdapter +from typing_extensions import ReadOnly from litellm.constants import TRACE_READ_RETRY_AFTER_SECONDS from litellm.proxy import tracing_endpoints @@ -97,6 +100,29 @@ SPAN_DETAIL_RESPONSE: Final = { "output_ui": {"kind": "text", "text": ""}, "attributes": {}, } +SPAN_ERROR_RESPONSE: Final = { + "span_id": "s1", + "message": "span error", + "total_chars": 10, + "next_cursor": None, +} +NOW_MS: Final = 1_800_000_000_000 + + +class RequestValidationError(TypedDict): + type: ReadOnly[str] + loc: ReadOnly[list[str | int]] + + +def _validation_errors(response: Response) -> tuple[RequestValidationError, ...]: + return tuple(TypeAdapter(list[RequestValidationError]).validate_python(response.json()["detail"])) + + +def _assert_validation_error(response: Response, error_type: str, location: tuple[str | int, ...]) -> None: + assert response.status_code == 422, response.text + assert any( + error["type"] == error_type and tuple(error["loc"]) == location for error in _validation_errors(response) + ) @pytest.mark.parametrize( @@ -262,6 +288,43 @@ def test_list_traces_defaults_to_last_24h(client, receiver): assert kwargs["cursor"] is None +@pytest.mark.parametrize( + ("params", "expected_start_ms", "expected_end_ms"), + ( + ({}, NOW_MS - tracing_endpoints.MS_PER_DAY, NOW_MS), + ({"start_ms": 123}, 123, NOW_MS), + ({"end_ms": -7}, NOW_MS - tracing_endpoints.MS_PER_DAY, -7), + ), +) +def test_list_traces_resolves_default_bounds_from_injected_clock( + client: TestClient, + receiver: MagicMock, + params: Mapping[str, int], + expected_start_ms: int, + expected_end_ms: int, +) -> None: + client.app.dependency_overrides[tracing_endpoints.current_time_ms] = lambda: NOW_MS + response: Final = client.get("/v1/traces", params=params) + assert response.status_code == 200, response.text + receiver.list_traces.assert_awaited_once_with( + scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, + start_ms=expected_start_ms, + end_ms=expected_end_ms, + cursor=None, + ) + + +def test_list_traces_forwards_large_and_negative_bounds_unchanged(client: TestClient, receiver: MagicMock) -> None: + response: Final = client.get("/v1/traces", params={"start_ms": 2**63, "end_ms": -1, "cursor": "next"}) + assert response.status_code == 200, response.text + receiver.list_traces.assert_awaited_once_with( + scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, + start_ms=2**63, + end_ms=-1, + cursor="next", + ) + + def test_get_trace_404_and_200(client, receiver): assert client.get("/v1/traces/missing").status_code == 404 receiver.get_trace.return_value = TRACE_RESPONSE @@ -289,6 +352,125 @@ def test_trace_detail_passes_scoped_reference(client, receiver, suffix, cursor, ) +@pytest.mark.parametrize("page_size", (1, 500)) +def test_trace_detail_accepts_page_size_bounds(client: TestClient, receiver: MagicMock, page_size: int) -> None: + receiver.get_trace.return_value = TRACE_RESPONSE + response: Final = client.get("/v1/traces/t1", params={"trace_ref": "run-one", "page_size": page_size}) + assert response.status_code == 200, response.text + receiver.get_trace.assert_awaited_once_with( + "t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", None, page_size + ) + + +@pytest.mark.parametrize( + ("value", "error_type"), + (("0", "greater_than_equal"), ("501", "less_than_equal"), ("abc", "int_parsing")), +) +def test_trace_detail_reports_page_size_validation( + client: TestClient, receiver: MagicMock, value: str, error_type: str +) -> None: + response: Final = client.get("/v1/traces/t1", params={"page_size": value}) + _assert_validation_error(response, error_type, ("query", "page_size")) + receiver.get_trace.assert_not_awaited() + + +def test_trace_read_routes_accept_and_forward_512_character_cursors( + client: TestClient, receiver: MagicMock +) -> None: + cursor: Final = "x" * 512 + receiver.get_trace.return_value = TRACE_RESPONSE + receiver.get_span_error = AsyncMock(return_value=SPAN_ERROR_RESPONSE) + client.app.dependency_overrides[tracing_endpoints.current_time_ms] = lambda: NOW_MS + + list_response: Final = client.get("/v1/traces", params={"cursor": cursor}) + detail_response: Final = client.get("/v1/traces/t1", params={"trace_ref": "run-one", "cursor": cursor}) + error_response: Final = client.get( + "/v1/traces/t1/spans/s1/error", params={"trace_ref": "run-one", "cursor": cursor} + ) + + assert list_response.status_code == 200, list_response.text + assert detail_response.status_code == 200, detail_response.text + assert error_response.status_code == 200, error_response.text + receiver.list_traces.assert_awaited_once_with( + scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, + start_ms=NOW_MS - tracing_endpoints.MS_PER_DAY, + end_ms=NOW_MS, + cursor=cursor, + ) + receiver.get_trace.assert_awaited_once_with( + "t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", cursor, None + ) + receiver.get_span_error.assert_awaited_once_with( + "t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", cursor + ) + + +@pytest.mark.parametrize( + "path", + ("/v1/traces", "/v1/traces/t1", "/v1/traces/t1/spans/s1/error"), +) +def test_trace_read_routes_reject_513_character_cursors( + client: TestClient, receiver: MagicMock, path: str +) -> None: + response: Final = client.get(path, params={"cursor": "x" * 513}) + _assert_validation_error(response, "string_too_long", ("query", "cursor")) + + +def test_trace_read_routes_ignore_unknown_query_parameters(client: TestClient, receiver: MagicMock) -> None: + receiver.get_trace.return_value = TRACE_RESPONSE + receiver.get_span.return_value = SPAN_DETAIL_RESPONSE + receiver.get_span_error = AsyncMock(return_value=SPAN_ERROR_RESPONSE) + + list_params: Final = {"start_ms": 1, "end_ms": 2, "cursor": "list-cursor"} + list_response: Final = client.get("/v1/traces", params=list_params) + list_unknown_response: Final = client.get("/v1/traces", params={**list_params, "foo": "bar"}) + detail_params: Final = {"trace_ref": "run-one", "cursor": "detail-cursor", "page_size": 10} + detail_response: Final = client.get("/v1/traces/t1", params=detail_params) + detail_unknown_response: Final = client.get("/v1/traces/t1", params={**detail_params, "foo": "bar"}) + span_response: Final = client.get("/v1/traces/t1/spans/s1", params={"trace_ref": "run-one"}) + span_unknown_response: Final = client.get( + "/v1/traces/t1/spans/s1", params={"trace_ref": "run-one", "foo": "bar"} + ) + error_params: Final = {"trace_ref": "run-one", "cursor": "error-cursor"} + error_response: Final = client.get("/v1/traces/t1/spans/s1/error", params=error_params) + error_unknown_response: Final = client.get( + "/v1/traces/t1/spans/s1/error", params={**error_params, "foo": "bar"} + ) + + assert list_response.status_code == 200, list_response.text + assert list_unknown_response.status_code == 200, list_unknown_response.text + assert detail_response.status_code == 200, detail_response.text + assert detail_unknown_response.status_code == 200, detail_unknown_response.text + assert span_response.status_code == 200, span_response.text + assert span_unknown_response.status_code == 200, span_unknown_response.text + assert error_response.status_code == 200, error_response.text + assert error_unknown_response.status_code == 200, error_unknown_response.text + receiver.list_traces.assert_has_awaits( + ( + call(scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, start_ms=1, end_ms=2, cursor="list-cursor"), + call(scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, start_ms=1, end_ms=2, cursor="list-cursor"), + ) + ) + receiver.get_trace.assert_has_awaits( + ( + call("t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", "detail-cursor", 10), + call("t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", "detail-cursor", 10), + ) + ) + receiver.get_span.assert_has_awaits( + ( + call("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one"), + call("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one"), + ) + ) + receiver.get_span_error.assert_has_awaits( + ( + call("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", "error-cursor"), + call("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", "error-cursor"), + ) + ) + + @pytest.mark.parametrize( "path,method", ( @@ -606,6 +788,8 @@ def test_sql_and_help_use_authenticated_scope( result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) assert result.status_code == 200, result.text assert result.json() == SQL_ENVELOPE + assert type(result.json()["data"][0]["value"]) is str + assert type(result.json()["statistics"]["elapsed"]) is float receiver.storage.query_sql.assert_awaited_once_with("SELECT * FROM otel_traces", expected_scope, "test-secret") help_result: Final = client.get("/v1/traces/query/help") assert help_result.status_code == 200, help_result.text @@ -616,6 +800,29 @@ def test_sql_and_help_use_authenticated_scope( assert receiver.storage.query_sql.await_count == 1 +@pytest.mark.parametrize( + ("body", "error_type", "location"), + ( + (b"{}", "missing", ("body", "sql")), + (b'{"sql": null}', "string_type", ("body", "sql")), + (b'{"sql": 1}', "string_type", ("body", "sql")), + (b'{"sql": "SELECT 1", "extra": true}', "extra_forbidden", ("body", "extra")), + (b"{", "json_invalid", ("body", 1)), + (b"[]", "model_attributes_type", ("body",)), + ), +) +def test_sql_query_rejects_invalid_request_bodies( + client: TestClient, receiver: MagicMock, body: bytes, error_type: str, location: tuple[str | int, ...] +) -> None: + client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret" + receiver.storage.query_sql = AsyncMock() + response: Final = client.post( + "/v1/traces/query", content=body, headers={"content-type": "application/json"} + ) + _assert_validation_error(response, error_type, location) + receiver.storage.query_sql.assert_not_awaited() + + @pytest.mark.parametrize("auth", (UserAPIKeyAuth(), UserAPIKeyAuth(team_id="a", project_id="p"))) def test_sql_rejects_missing_identity_without_querying( client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth diff --git a/ui/litellm-dashboard/src/components/lens/traces/AGENTS.md b/ui/litellm-dashboard/src/components/lens/traces/AGENTS.md new file mode 100644 index 00000000000..bc378fd192c --- /dev/null +++ b/ui/litellm-dashboard/src/components/lens/traces/AGENTS.md @@ -0,0 +1,3 @@ +- Derive trace request and response types from `paths` and `components` in `@/lib/http/schema` (see `types.ts`); never hand-write a trace request or response shape +- Check every trace request `query` or body literal with `satisfies` against those types +- `requestTypes.test-d.ts` pins the generated request shapes; a failure there means the backend contract changed diff --git a/ui/litellm-dashboard/src/components/lens/traces/api.ts b/ui/litellm-dashboard/src/components/lens/traces/api.ts index 03ea4d068bb..6b9039eae7c 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/api.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/api.ts @@ -9,7 +9,7 @@ import { apiClient, getProxyBaseUrl, } from "../../networking"; -import type { SpanDetail, SpanErrorPage, Trace, TracePage } from "./types"; +import type { SpanDetail, SpanErrorPage, Trace, TraceListQuery, TracePage } from "./types"; export interface TraceWindow { readonly startMs: number; @@ -53,7 +53,10 @@ export function liveTracesApi(accessToken: string): TracesApi { }), list: (window) => agentTraceListCall({ accessToken, ...window }), anyRecorded: async () => { - const page = await apiClient.get("/v1/traces", { accessToken, query: { start_ms: 0 } }); + const page = await apiClient.get("/v1/traces", { + accessToken, + query: { start_ms: 0 } satisfies TraceListQuery, + }); return page.data.length > 0; }, trace: (traceId, traceRef, cursor) => agentTraceCall(accessToken, traceId, traceRef, cursor), diff --git a/ui/litellm-dashboard/src/components/lens/traces/list/runSearch/runSql.ts b/ui/litellm-dashboard/src/components/lens/traces/list/runSearch/runSql.ts index 96b5ae6455c..73b65f05eae 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/list/runSearch/runSql.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/list/runSearch/runSql.ts @@ -1,5 +1,6 @@ import { getProxyBaseUrl } from "@/components/networking"; import type { TimeWindow } from "@/components/shared/timeline/Timeline"; +import type { TraceQueryBody } from "../../types"; import { valueMatcher } from "@/components/shared/search/language"; import type { SearchFilter, SearchQuery } from "@/components/shared/search/searchQuery"; @@ -71,7 +72,7 @@ export const traceQueryCommand = (sql: string): string => [ `curl -s "${getProxyBaseUrl().replace(/\/$/, "")}/v1/traces/query" \\`, ` -H "Authorization: Bearer $LITELLM_API_KEY" -H "Content-Type: application/json" -d @- <<'EOF'`, - JSON.stringify({ sql }), + JSON.stringify({ sql } satisfies TraceQueryBody), "EOF", ].join("\n"); diff --git a/ui/litellm-dashboard/src/components/lens/traces/requestTypes.test-d.ts b/ui/litellm-dashboard/src/components/lens/traces/requestTypes.test-d.ts new file mode 100644 index 00000000000..3ef107cbcad --- /dev/null +++ b/ui/litellm-dashboard/src/components/lens/traces/requestTypes.test-d.ts @@ -0,0 +1,21 @@ +import { expectTypeOf, test } from "vitest"; +import type { SpanErrorQuery, SpanQuery, TraceDetailQuery, TraceListQuery, TraceQueryBody } from "./types"; + +test("trace request aliases match the generated OpenAPI shapes", () => { + expectTypeOf().toEqualTypeOf<{ + start_ms?: number | null; + end_ms?: number | null; + cursor?: string | null; + }>(); + expectTypeOf().toEqualTypeOf<{ + trace_ref?: string; + cursor?: string | null; + page_size?: number | null; + }>(); + expectTypeOf().toEqualTypeOf<{ trace_ref?: string }>(); + expectTypeOf().toEqualTypeOf<{ + trace_ref?: string; + cursor?: string | null; + }>(); + expectTypeOf().toEqualTypeOf<{ sql: string }>(); +}); diff --git a/ui/litellm-dashboard/src/components/lens/traces/types.ts b/ui/litellm-dashboard/src/components/lens/traces/types.ts index a424e010670..06a2b93b4ae 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/types.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/types.ts @@ -4,6 +4,13 @@ export type Trace = paths["/v1/traces/{trace_id}"]["get"]["responses"][200]["con export type TracePage = paths["/v1/traces"]["get"]["responses"][200]["content"]["application/json"]; export type SpanErrorPage = paths["/v1/traces/{trace_id}/spans/{span_id}/error"]["get"]["responses"][200]["content"]["application/json"]; +export type TraceListQuery = NonNullable; +export type TraceDetailQuery = NonNullable; +export type SpanQuery = NonNullable; +export type SpanErrorQuery = NonNullable< + paths["/v1/traces/{trace_id}/spans/{span_id}/error"]["get"]["parameters"]["query"] +>; +export type TraceQueryBody = components["schemas"]["TraceQueryRequest"]; type ApiSpanDetail = paths["/v1/traces/{trace_id}/spans/{span_id}"]["get"]["responses"][200]["content"]["application/json"]; export type Span = Trace["spans"][number]; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 98f3a120a49..a7e9387cbe6 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -117,7 +117,16 @@ import type { ComplexityRouterConfigPayload } from "./add_model/build_complexity import type { AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; import type { VectorStoreIndex } from "@/app/(dashboard)/vector-stores/_components/IndexesTab"; import type { RoutingDecision } from "./logs/detail/RoutingDecisionCard"; -import type { SpanDetail, SpanErrorPage, Trace, TracePage } from "./lens/traces/types"; +import type { + SpanDetail, + SpanErrorPage, + SpanErrorQuery, + SpanQuery, + Trace, + TraceDetailQuery, + TraceListQuery, + TracePage, +} from "./lens/traces/types"; import { createApiClient, deriveErrorMessage, @@ -1968,7 +1977,7 @@ export const agentTraceListCall = async ({ endMs: number; cursor?: string | null; }): Promise => { - const query = { start_ms: startMs, end_ms: endMs, cursor: cursor ?? undefined }; + const query = { start_ms: startMs, end_ms: endMs, cursor: cursor ?? undefined } satisfies TraceListQuery; return apiClient.get(`/v1/traces`, { accessToken, query }); }; @@ -1983,7 +1992,7 @@ export const agentTraceCall = async ( ): Promise => apiClient.get(`/v1/traces/${encodeURIComponent(traceId)}`, { accessToken, - query: { trace_ref: traceRef || undefined, cursor: cursor ?? undefined, page_size: 200 }, + query: { trace_ref: traceRef || undefined, cursor: cursor ?? undefined, page_size: 200 } satisfies TraceDetailQuery, }); export const agentTraceSpanCall = async ( @@ -1994,7 +2003,7 @@ export const agentTraceSpanCall = async ( ): Promise => apiClient.get(`/v1/traces/${encodeURIComponent(traceId)}/spans/${encodeURIComponent(spanId)}`, { accessToken, - query: { trace_ref: traceRef || undefined }, + query: { trace_ref: traceRef || undefined } satisfies SpanQuery, }); export const agentTraceSpanErrorCall = async ( @@ -2005,7 +2014,7 @@ export const agentTraceSpanErrorCall = async ( ): Promise => apiClient.get(`/v1/traces/${encodeURIComponent(traceId)}/spans/${encodeURIComponent(spanId)}/error`, { accessToken, - query: { trace_ref: options.traceRef || undefined, cursor: options.cursor || undefined }, + query: { trace_ref: options.traceRef || undefined, cursor: options.cursor || undefined } satisfies SpanErrorQuery, }); export const adminSpendLogsCall = async (accessToken: string) => { From 8364f88cbb6f51be00422e42700038708d1866cc Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 5 Oct 2026 10:49:16 -0700 Subject: [PATCH 004/191] ci: run unit selections from GHA test-path and drop the CircleCI unit jobs (#44461) * ci: run unit selections from GHA test-path and drop the CircleCI unit jobs * ci: keep existing shard token order and header note * ci: throwaway, drop tests/unit/repositories from the unit shard to show assert-ci-coverage fails * ci: revert throwaway assert-ci-coverage check * ci: stop crediting --ignore paths as invoked in assert_ci_coverage * ci: install the caching, extra_proxy and proxy-runtime extras in the GHA unit sync --------- Co-authored-by: yuneng --- .circleci/config.yml | 30 -- .circleci/scripts/unit_selection.sh | 187 --------- .circleci/tests.yml | 416 -------------------- .github/scripts/assert_ci_coverage.py | 80 ++-- .github/workflows/_test-unit-base.yml | 16 +- .github/workflows/test-unit-proxy-db.yml | 100 +++-- .github/workflows/test-unit.yml | 114 ++++-- tests/unit/test_assert_ci_coverage.py | 48 +-- tests/unit/test_unit_shard_missing_paths.py | 1 - 9 files changed, 226 insertions(+), 766 deletions(-) delete mode 100755 .circleci/scripts/unit_selection.sh delete mode 100644 .circleci/tests.yml diff --git a/.circleci/config.yml b/.circleci/config.yml index f2d51ba7a0f..2653135da41 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3292,35 +3292,6 @@ jobs: - store_artifacts: path: test-results - unit: - machine: - image: ubuntu-2204:2024.04.1 - resource_class: large - working_directory: ~/project - parallelism: 4 - steps: - - setup_litellm_test_deps - - run: - name: Generate Prisma client - command: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - run: - name: Run unit tests - command: | - mkdir -p test-results/unit - shard="$(find tests/unit -name 'test_*.py' | sort | circleci tests split --split-by=timings --timings-type=filename)" - if [ -z "${shard}" ]; then echo "shard ${CIRCLE_NODE_INDEX} received no tests/unit files; nothing to run"; exit 0; fi - mapfile -t files < <(printf '%s\n' "${shard}") - set +e - LITELLM_LOCAL_MODEL_COST_MAP=True uv run --no-sync pytest "${files[@]}" -p no:rerunfailures -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short -o junit_family=xunit1 --junitxml=test-results/unit/junit.xml - status=$? - set -e - if [ "$status" -eq 5 ]; then echo "pytest collected no tests from tests/unit; passing"; exit 0; fi - exit "$status" - - store_test_results: - path: test-results - - store_artifacts: - path: test-results - postgres_suite: parameters: test_path: @@ -3594,7 +3565,6 @@ workflows: jobs: - using_litellm_on_windows - windows_release_wheel - - unit - provider_replay_harness - base_sdk_install - local_testing_part1 diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh deleted file mode 100755 index 02df32d5eab..00000000000 --- a/.circleci/scripts/unit_selection.sh +++ /dev/null @@ -1,187 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -flag="${1:?usage: unit_selection.sh }" - -legacy_flags=( - caching-local - core-utils - enterprise-package - enterprise-routing - integrations - llm-other-providers - llm-vertex-ai - mcp-integration - misc - proxy-db-auth-checks - proxy-db-budgets - proxy-db-custom-logging - proxy-db-db-and-spend - proxy-db-endpoints-and-responses - proxy-db-guardrails-hooks - proxy-db-jwt-and-keys - proxy-db-key-generation - proxy-db-logging-misc - proxy-db-proxy-runtime - proxy-db-proxy-server-core - proxy-db-proxy-utils - proxy-extras - proxy-infra - responses-caching-types -) - -legacy_paths() { - case "$1" in - caching-local) echo tests/unit/caching ;; - core-utils) echo tests/unit/litellm_core_utils ;; - enterprise-package) - echo tests/unit/enterprise/integrations - echo tests/unit/enterprise/proxy/auth - echo tests/unit/enterprise/proxy/guardrails - echo tests/unit/enterprise/proxy/hooks - echo tests/unit/enterprise/proxy/management_endpoints - echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py - echo tests/unit/enterprise/proxy/test_liteadmin.py - echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;; - enterprise-routing) - echo tests/unit/google_genai - echo tests/unit/router_strategy - echo tests/unit/router_utils - echo tests/unit/proxy/common_utils/test_cache_aware_routing.py - echo tests/unit/enterprise/enterprise_callbacks/send_emails - echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py - echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py - echo tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py - echo tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py - echo tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py - echo tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py - echo tests/unit/enterprise/proxy/test_enterprise_routes.py - echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py - echo tests/unit/enterprise/proxy/test_managed_files_access_check.py - echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; - integrations) echo tests/unit/integrations ;; - llm-other-providers) find tests/unit/llms -name 'test_*.py' -not -path 'tests/unit/llms/vertex_ai/*' ;; - llm-vertex-ai) echo tests/unit/llms/vertex_ai ;; - mcp-integration) - echo tests/unit/experimental_mcp_client - echo tests/unit/proxy/_experimental/mcp_server - echo tests/unit/responses/mcp - echo tests/mcp_tests/test_proxy_mcp_e2e.py ;; - misc) - find tests/unit -maxdepth 1 -name 'test_*.py' - echo tests/unit/test_router - echo tests/unit/a2a_protocol - echo tests/unit/batches - echo tests/unit/chat_completions - echo tests/unit/completion_extras - echo tests/unit/containers - echo tests/unit/embeddings - echo tests/unit/endpoints - echo tests/unit/files - echo tests/unit/harness - echo tests/unit/images - echo tests/unit/interactions - echo tests/unit/messages - echo tests/unit/rag - echo tests/unit/rerank_api - echo tests/unit/rust_bridge - echo tests/unit/secret_managers - echo tests/unit/vector_stores - echo tests/unit/videos ;; - proxy-db-auth-checks) - echo tests/unit/proxy/auth/test_auth_checks.py - echo tests/unit/proxy/auth/test_user_api_key_auth.py - echo tests/unit/proxy/test_credential_slot_registry.py - echo tests/unit/proxy/test_deprecated_key_grace_period.py ;; - proxy-db-budgets) - echo tests/unit/proxy/auth/test_default_end_user_budget_simple.py - echo tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py - echo tests/unit/proxy/test_zero_cost_model_budget_bypass.py ;; - proxy-db-custom-logging) - echo tests/unit/proxy/test_custom_callback_input.py - echo tests/unit/proxy/test_custom_logger_s3_gcs.py ;; - proxy-db-db-and-spend) - echo tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py - echo tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py - echo tests/unit/proxy/db/test_update_daily_tag_spend.py - echo tests/unit/proxy/test_db_schema_changes.py - echo tests/unit/proxy/test_prisma_client_backoff_retry.py - echo tests/unit/proxy/test_update_spend.py - echo tests/unit/skills/test_skills_db.py ;; - proxy-db-endpoints-and-responses) - echo tests/unit/proxy/lens - echo tests/unit/proxy/auth/test_models_fallback_endpoint.py - echo tests/unit/proxy/common_utils/test_check_batch_cost.py - echo tests/unit/proxy/common_utils/test_check_responses_cost.py - echo tests/unit/proxy/common_utils/test_realtime_cache.py - echo tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py - echo tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py - echo tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py - echo tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py - echo tests/unit/proxy/response_polling - echo tests/unit/proxy/test_custom_tokenizer_bug.py - echo tests/unit/proxy/test_get_favicon.py - echo tests/unit/proxy/test_get_image.py - echo tests/unit/proxy/test_prompt_test_endpoint.py - echo tests/unit/proxy/test_reducto_ocr_route.py - echo tests/unit/proxy/test_response_polling_pre_call_checks.py - echo tests/unit/proxy/test_ui_path_detection.py ;; - proxy-db-guardrails-hooks) - echo tests/unit/proxy/hooks/test_banned_keyword_list.py - echo tests/unit/proxy/test_proxy_setting_guardrails.py - echo tests/unit/proxy/test_unit_test_proxy_hooks.py ;; - proxy-db-jwt-and-keys) - echo tests/unit/proxy/auth/test_jwt.py - echo tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py - echo tests/unit/proxy/test_proxy_custom_auth.py ;; - proxy-db-key-generation) echo tests/unit/proxy/management_endpoints/test_key_generate_prisma.py ;; - proxy-db-logging-misc) - echo tests/unit/proxy/management_helpers/test_audit_logs_proxy.py - echo tests/unit/proxy/spend_tracking/test_search_api_logging.py - echo tests/unit/proxy/test_proxy_reject_logging.py ;; - proxy-db-proxy-runtime) - echo tests/unit/proxy/auth/test_multipart_bypass_repro.py - echo tests/unit/proxy/auth/test_proxy_routes.py - echo tests/unit/proxy/middleware/test_request_size_limit_middleware.py - echo tests/unit/proxy/test_proxy_config_unit_test.py - echo tests/unit/proxy/test_proxy_token_counter.py - echo tests/unit/proxy/test_server_root_path.py ;; - proxy-db-proxy-server-core) - echo tests/unit/proxy/test__lazy_features.py - echo tests/unit/proxy/test_aproxy_startup.py - echo tests/unit/proxy/test_proxy_server.py ;; - proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;; - proxy-extras) echo tests/unit/litellm_proxy_extras ;; - proxy-infra) - echo tests/unit/gateway - echo tests/unit/proxy/management - echo tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py - echo tests/unit/proxy/roi_calculator ;; - responses-caching-types) - find tests/unit/responses -name 'test_*.py' -not -path 'tests/unit/responses/mcp/*' - echo tests/unit/types ;; - *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;; - esac -} - -expand() { - while read -r path; do - if [ -d "$path" ]; then - find "$path" -name 'test_*.py' - elif [ -f "$path" ]; then - echo "$path" - else - echo "unit_selection.sh: $path does not exist" >&2 - exit 1 - fi - done -} - -if [ "$flag" = unit ]; then - comm -23 \ - <(find tests/unit -name 'test_*.py' | sort) \ - <(for legacy in "${legacy_flags[@]}"; do legacy_paths "$legacy"; done | expand | sort) - exit 0 -fi - -legacy_paths "$flag" | expand | sort diff --git a/.circleci/tests.yml b/.circleci/tests.yml deleted file mode 100644 index a9cd21bad5e..00000000000 --- a/.circleci/tests.yml +++ /dev/null @@ -1,416 +0,0 @@ -version: 2.1 - -commands: - wait_for_service: - parameters: - url: - type: string - timeout: - type: string - default: "60" - steps: - - run: - name: "Wait for << parameters.url >>" - command: | - TIMEOUT=<< parameters.timeout >> - URL="<< parameters.url >>" - ELAPSED=0 - echo "Waiting up to ${TIMEOUT}s for ${URL} ..." - if echo "$URL" | grep -q '^tcp://'; then - HOST=$(echo "$URL" | sed 's|tcp://||' | cut -d: -f1) - PORT=$(echo "$URL" | sed 's|tcp://||' | cut -d: -f2) - while ! bash -c "echo > /dev/tcp/$HOST/$PORT" 2>/dev/null; do - sleep 2; ELAPSED=$((ELAPSED+2)) - if [ "$ELAPSED" -ge "$TIMEOUT" ]; then echo "Timed out"; exit 1; fi - done - else - while ! curl -sf --max-time 5 "$URL" > /dev/null 2>&1; do - sleep 2; ELAPSED=$((ELAPSED+2)) - if [ "$ELAPSED" -ge "$TIMEOUT" ]; then echo "Timed out"; exit 1; fi - done - fi - echo "Service ready after ${ELAPSED}s" - install_uv: - steps: - - run: - name: Install uv (pinned 0.10.9) - command: | - curl -LsSf -o /tmp/uv-install.sh https://astral.sh/uv/0.10.9/install.sh - echo "7fc46e39cb97290b57169c0c813a17970585ac519139f19006453c99b5f2f45f /tmp/uv-install.sh" | sha256sum -c - - env UV_NO_MODIFY_PATH=1 sh /tmp/uv-install.sh - rm -f /tmp/uv-install.sh - echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$BASH_ENV" - export PATH="$HOME/.local/bin:$PATH" - install_rust: - steps: - - run: - name: Install Rust (rustup 1.28.2, toolchain 1.98.0) - command: | - case "$(uname -m)" in - x86_64) - RUSTUP_TRIPLE=x86_64-unknown-linux-gnu - RUSTUP_SHA256=20a06e644b0d9bd2fbdbfd52d42540bdde820ea7df86e92e533c073da0cdd43c - ;; - aarch64) - RUSTUP_TRIPLE=aarch64-unknown-linux-gnu - RUSTUP_SHA256=e3853c5a252fca15252d07cb23a1bdd9377a8c6f3efa01531109281ae47f841c - ;; - *) - echo "install_rust: unsupported architecture $(uname -m)" >&2 - exit 1 - ;; - esac - curl -sSLf -o /tmp/rustup-init \ - "https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init" - echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c - - chmod +x /tmp/rustup-init - /tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.98.0 - rm -f /tmp/rustup-init - echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV" - export PATH="$HOME/.cargo/bin:$PATH" - rustc --version - cargo --version - install_codecov_cli: - steps: - - run: - name: Install Codecov CLI (pinned v11.3.1) - when: always - command: | - curl -sSLf -o /tmp/codecov https://cli.codecov.io/v11.3.1/linux/codecov - curl -sSLf -o /tmp/codecov.SHA256SUM https://cli.codecov.io/v11.3.1/linux/codecov.SHA256SUM - [ "$(cat /tmp/codecov.SHA256SUM)" = "ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d codecov" ] - (cd /tmp && sha256sum -c codecov.SHA256SUM) - chmod +x /tmp/codecov - mkdir -p "$HOME/.local/bin" - mv /tmp/codecov "$HOME/.local/bin/codecov" - setup_litellm_enterprise_pip: - steps: - - run: - name: "Install local version of litellm-enterprise" - command: | - uv run --no-sync python -c "import litellm_enterprise; print('litellm-enterprise OK:', litellm_enterprise.__file__)" - setup_test_deps: - steps: - - install_uv - - install_rust - - restore_cache: - keys: - - v3-integration-uv-cache-{{ checksum "uv.lock" }} - - run: - name: Install Dependencies - command: | - uv sync --frozen --all-groups --all-extras --python 3.12 - - setup_litellm_enterprise_pip - - save_cache: - paths: - - ~/.cache/uv - key: v3-integration-uv-cache-{{ checksum "uv.lock" }} - - run: - name: Generate Prisma client - command: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - skip_unless_relevant: - parameters: - category: - type: string - default: backend - base_ref: - type: string - default: "" - pull_request_url: - type: string - default: "" - steps: - - run: - name: "Skip job when no << parameters.category >>-relevant files changed" - command: | - export CIRCLE_PULL_REQUEST="${CIRCLE_PULL_REQUEST:-<< parameters.pull_request_url >>}" - export PATH_FILTER_BASE_BRANCH="<< parameters.base_ref >>" - [ -n "$PATH_FILTER_BASE_BRANCH" ] || unset PATH_FILTER_BASE_BRANCH - bash .circleci/scripts/path_filter.sh << parameters.category >> - start_postgres: - parameters: - db_name: - type: string - default: circle_test - image: - type: string - default: postgres:14@sha256:6a70deda415ec296f977890e11aba04a0db9f632a362e3fce45e845e3db74f26 - steps: - - run: - name: Start PostgreSQL - command: | - docker run -d \ - --name postgres-db \ - -e POSTGRES_USER=postgres \ - -e POSTGRES_PASSWORD=postgres \ - -e POSTGRES_DB=<< parameters.db_name >> \ - -p 5432:5432 \ - << parameters.image >> - - wait_for_service: - url: tcp://localhost:5432 - timeout: "60" - start_redis: - steps: - - run: - name: Start Redis - command: | - docker run -d \ - --name redis-cache \ - -p 6379:6379 \ - redis:7-alpine@sha256:7aec734b2bb298a1d769fd8729f13b8514a41bf90fcdd1f38ec52267fbaa8ee6 - - wait_for_service: - url: tcp://localhost:6379 - timeout: "60" - -jobs: - unit: - parameters: - flag: - type: string - default: unit - shards: - type: integer - default: 6 - workers: - type: integer - default: 4 - dist: - type: string - default: loadscope - base_ref: - type: string - default: "" - pull_request_url: - type: string - default: "" - legacy_mcp_peer: - type: boolean - default: false - reruns: - type: integer - default: 0 - machine: - image: ubuntu-2204:2024.04.1 - resource_class: large - working_directory: ~/project - parallelism: << parameters.shards >> - environment: - COVERAGE_CORE: sysmon - LITELLM_LOCAL_MODEL_COST_MAP: "True" - steps: - - checkout - - skip_unless_relevant: - base_ref: << parameters.base_ref >> - pull_request_url: << parameters.pull_request_url >> - - setup_test_deps - - when: - condition: << parameters.legacy_mcp_peer >> - steps: - - run: - name: Install MCP SDK1 peer - command: | - uv venv --python 3.12 .venv-mcp-peer - uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' - echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - - run: - name: "Run << parameters.flag >> shard" - no_output_timeout: 20m - command: | - mkdir -p test-results/<< parameters.flag >> - selection="$(bash .circleci/scripts/unit_selection.sh << parameters.flag >>)" || { echo "unit_selection.sh failed for << parameters.flag >>"; exit 1; } - [ -n "${selection}" ] || { echo "unit_selection.sh produced no files for << parameters.flag >>"; exit 1; } - shard="$(printf '%s\n' "${selection}" | circleci tests split --split-by=timings --timings-type=filename)" || { echo "circleci tests split failed for << parameters.flag >>"; exit 1; } - [ -n "${shard}" ] || { echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.flag >> files; nothing to run"; exit 0; } - mapfile -t files < <(printf '%s\n' "${shard}") - xdist_args=() - if [ "<< parameters.workers >>" -gt 0 ]; then xdist_args=(-n << parameters.workers >> --dist=<< parameters.dist >>); fi - rerun_args=(-p no:rerunfailures) - if [ "<< parameters.reruns >>" -gt 0 ]; then rerun_args=(--reruns << parameters.reruns >> --reruns-delay 1 --rerun-except "from pytest-timeout"); fi - test_env=(PATH="$PATH" HOME="$HOME" CI=true COVERAGE_CORE="$COVERAGE_CORE" LITELLM_LOCAL_MODEL_COST_MAP="$LITELLM_LOCAL_MODEL_COST_MAP") - if [ -n "${MCP_TEST_PEER_PYTHON:-}" ]; then test_env+=(MCP_TEST_PEER_PYTHON="$MCP_TEST_PEER_PYTHON"); fi - set +e - env -i "${test_env[@]}" \ - uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 "${xdist_args[@]}" --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml - status=$? - set -e - if [ "$status" -eq 5 ]; then echo "pytest collected no tests from the shard; passing"; exit 0; fi - exit "$status" - - install_codecov_cli - - run: - name: Upload coverage - when: always - command: | - [ -f coverage.xml ] || { echo "no coverage.xml produced; skipping upload"; exit 0; } - codecov upload-process --disable-search -f coverage.xml -F << parameters.flag >> -C "$CIRCLE_SHA1" -n "<< parameters.flag >>-${CIRCLE_NODE_INDEX}-${CIRCLE_BUILD_NUM}" --git-service github - - store_test_results: - path: test-results - - store_artifacts: - path: test-results - - store_artifacts: - path: coverage.xml - documentation: - machine: - image: ubuntu-2204:2024.04.1 - resource_class: large - working_directory: ~/project - steps: - - checkout - - setup_test_deps - - run: - name: Checkout litellm-docs - command: rm -rf docs/my-website && git clone --depth 1 https://github.com/BerriAI/litellm-docs.git docs/my-website - - run: - name: Run documentation validation - command: | - uv run --no-sync python ./tests/documentation_tests/test_env_keys.py - uv run --no-sync python ./tests/documentation_tests/test_router_settings.py - uv run --no-sync python ./tests/documentation_tests/test_api_docs.py - uv run --no-sync python ./tests/documentation_tests/test_circular_imports.py - integration: - parameters: - suite: - type: string - base_ref: - type: string - default: "" - pull_request_url: - type: string - default: "" - machine: - image: ubuntu-2204:2024.04.1 - resource_class: large - working_directory: ~/project - steps: - - checkout - - skip_unless_relevant: - base_ref: << parameters.base_ref >> - pull_request_url: << parameters.pull_request_url >> - - setup_test_deps - - start_postgres: - image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 - - start_redis - - run: - name: Run owned integration contracts - command: env -i PATH="$PATH" HOME="$HOME" CIRCLE_SHA1="$CIRCLE_SHA1" CIRCLE_WORKFLOW_ID="$CIRCLE_WORKFLOW_ID" bash .circleci/scripts/run_integration.sh << parameters.suite >> - no_output_timeout: 15m - - run: - name: Stop owned database and Redis - when: always - command: | - mkdir -p test-results/integration-<< parameters.suite >> - docker logs postgres-db > test-results/integration-<< parameters.suite >>/postgres.log 2>&1 || true - docker logs redis-cache > test-results/integration-<< parameters.suite >>/redis.log 2>&1 || true - docker rm -f postgres-db redis-cache - test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)" - - store_test_results: - path: test-results - - store_artifacts: - path: test-results - -workflows: - tests: - when: (pipeline.event.name == "push" and pipeline.git.branch == "main") or pipeline.event.name == "api" or (pipeline.event.name == "pull_request" and (pipeline.event.github.pull_request.base.ref == "main" or pipeline.event.github.pull_request.base.ref starts-with "litellm_")) - jobs: - - unit: - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-<< matrix.flag >> - shards: 1 - workers: 2 - reruns: 2 - matrix: - parameters: - flag: [caching-local, proxy-extras, enterprise-routing] - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-mcp-integration - flag: mcp-integration - shards: 1 - workers: 2 - legacy_mcp_peer: true - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-<< matrix.flag >> - shards: 1 - reruns: 2 - matrix: - parameters: - flag: - - enterprise-package - - proxy-infra - - responses-caching-types - - proxy-db-auth-checks - - proxy-db-jwt-and-keys - - proxy-db-proxy-server-core - - proxy-db-proxy-runtime - - proxy-db-custom-logging - - proxy-db-logging-misc - - proxy-db-db-and-spend - - proxy-db-guardrails-hooks - - proxy-db-budgets - - proxy-db-endpoints-and-responses - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-llm-vertex-ai - flag: llm-vertex-ai - shards: 2 - workers: 1 - reruns: 2 - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-llm-other-providers - flag: llm-other-providers - shards: 3 - reruns: 2 - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-core-utils - flag: core-utils - shards: 2 - reruns: 1 - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-integrations - flag: integrations - shards: 2 - reruns: 3 - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-misc - flag: misc - shards: 2 - reruns: 2 - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-proxy-db-proxy-utils - flag: proxy-db-proxy-utils - shards: 1 - reruns: 2 - dist: worksteal - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - unit: - name: unit-proxy-db-key-generation - flag: proxy-db-key-generation - shards: 1 - workers: 0 - reruns: 2 - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - - documentation - - integration: - name: integration-<< matrix.suite >> - matrix: - parameters: - suite: [sdk] - base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> - pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> diff --git a/.github/scripts/assert_ci_coverage.py b/.github/scripts/assert_ci_coverage.py index 59511c7dc7b..e90589d3c95 100644 --- a/.github/scripts/assert_ci_coverage.py +++ b/.github/scripts/assert_ci_coverage.py @@ -9,7 +9,6 @@ import sys import warnings from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass -from types import MappingProxyType from typing import Final import yaml @@ -27,6 +26,7 @@ DOCKERFILE_INPUT_KEYS = frozenset({"file", "dockerfile"}) TEST_RUNNER_RE = re.compile(r"\bpytest\b|\bcircleci tests\b|\bhelm unittest\b|\bplaywright test\b|\bpython[0-9.]*\s") IMAGE_BUILD_RE = re.compile(r"\bdocker\s+(?:buildx\s+)?build\b") TEST_TOKEN_RE = re.compile(r"tests/[A-Za-z0-9_./*?-]+") +IGNORE_ARG_RE: Final = re.compile(r"--ignore(?:-glob)?[= ](\S+)") DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*") COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE) GLOB_CHARS = frozenset("*?") @@ -72,6 +72,17 @@ class Scalar: value: str +@dataclass(frozen=True, slots=True) +class Selection: + included: frozenset[str] + ignored: frozenset[str] + + def covers(self, relative_path: str) -> bool: + return any(_token_covers(token, relative_path) for token in self.included) and not any( + _token_covers(token, relative_path) for token in self.ignored + ) + + @dataclass(frozen=True, slots=True) class Finding: subject: str @@ -111,44 +122,32 @@ def _uncommented(value: str) -> str: return COMMENT_RE.sub("", value) -def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]: - return frozenset( - match.group(0).rstrip("/") +def _selection_for_scalar(scalar: Scalar) -> Selection: + text: Final = _uncommented(scalar.value) + ignored: Final = frozenset().union( + *( + frozenset(token.rstrip("/") for token in TEST_TOKEN_RE.findall(ignored_argument)) + for ignored_argument in IGNORE_ARG_RE.findall(text) + ) + ) + included: Final = frozenset( + match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(IGNORE_ARG_RE.sub("", text)) + ) + return Selection(included=included, ignored=ignored) + + +def _invoked_selections(scalars: Iterable[Scalar]) -> tuple[Selection, ...]: + selected_scalars: Final = tuple( + scalar for scalar in scalars if scalar.key in TEST_PATH_KEYS or TEST_RUNNER_RE.search(scalar.value) - for match in TEST_TOKEN_RE.finditer(_uncommented(scalar.value)) ) + selections: Final = tuple(_selection_for_scalar(scalar) for scalar in selected_scalars) + return tuple(selection for selection in selections if selection.included) -SELECTION_ARM_RE = re.compile(r"(?ms)^\s*([A-Za-z0-9_|*-]+)\)\s*(.*?);;") - - -def _unit_selection_arms(repo_root: pathlib.Path = REPO_ROOT) -> Mapping[str, frozenset[str]]: - script: Final = repo_root / ".circleci/scripts/unit_selection.sh" - if not script.is_file(): - return MappingProxyType({}) - text: Final = _uncommented(script.read_text()) - return MappingProxyType( - { - label: frozenset(match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(body)) - for label, body in SELECTION_ARM_RE.findall(text) - } - ) - - -def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]: - return frozenset(token for tokens in _unit_selection_arms(repo_root).values() for token in tokens) - - -def _wired_unit_flags(scalars: Iterable[Scalar]) -> frozenset[str]: - return frozenset(scalar.value for scalar in scalars if scalar.key == "unit-flag" and "${{" not in scalar.value) - - -def _shard_tokens(scalars: Iterable[Scalar], arms: Mapping[str, frozenset[str]]) -> frozenset[str]: - wired: Final = _wired_unit_flags(scalars) - return _invoked_test_tokens(scalars) | frozenset( - token for label, tokens in arms.items() if label in wired for token in tokens - ) +def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]: + return frozenset().union(*(selection.included for selection in _invoked_selections(scalars))) def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]: @@ -211,11 +210,12 @@ def _dockerfiles() -> tuple[str, ...]: ) -def _uncovered_tests(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]: +def _uncovered_tests(allowlist: Allowlist, selections: tuple[Selection, ...]) -> tuple[Finding, ...]: uncovered = tuple( relative_path for relative_path in _test_files() - if not any(_token_covers(token, relative_path) for token in tokens) and not allowlist.covers_test(relative_path) + if not any(selection.covers(relative_path) for selection in selections) + and not allowlist.covers_test(relative_path) ) directories = tuple(dict.fromkeys(path.rsplit("/", 1)[0] for path in uncovered)) return tuple( @@ -505,7 +505,7 @@ def _check_slices() -> int: def _check_shards() -> int: - findings = _unassigned_shard_children(_shard_tokens(_all_scalars(), _unit_selection_arms())) + findings = _unassigned_shard_children(_invoked_test_tokens(_all_scalars())) if findings: _report( "test directories and files that no shard claims", @@ -678,7 +678,11 @@ def main() -> int: integration_paths, ownership_findings = _integration_ownership() test_findings = ( - _uncovered_tests(allowlist, _invoked_test_tokens(scalars) | _unit_selection_tokens() | integration_paths) + _uncovered_tests( + allowlist, + _invoked_selections(scalars) + + (Selection(included=integration_paths, ignored=frozenset()),), + ) + ownership_findings ) dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars)) diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index a5a7787c569..7db01d721cf 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -13,16 +13,6 @@ on: have its path existence-checked like any other token. required: true type: string - unit-flag: - description: >- - Codecov flag of the `.circleci/tests.yml` job that now owns part of - this shard. The shard also runs the files - `.circleci/scripts/unit_selection.sh` lists for the flag, on every - event, because the CircleCI pipeline is manual-only while the tests - migrate. - required: false - type: string - default: "" workers: description: "Number of pytest-xdist workers" required: false @@ -138,7 +128,7 @@ jobs: timeout-minutes: 8 run: | diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra caching --extra extra_proxy --extra proxy-runtime uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' - name: Cache Prisma binaries @@ -158,7 +148,6 @@ jobs: timeout-minutes: ${{ inputs.timeout-minutes }} env: TEST_PATH: ${{ inputs.test-path }} - UNIT_FLAG: ${{ inputs.unit-flag }} MAX_FAILURES: ${{ inputs.max-failures }} WORKERS: ${{ inputs.workers }} RERUNS: ${{ inputs.reruns }} @@ -168,9 +157,6 @@ jobs: run: | echo "has-coverage=false" >> "$GITHUB_OUTPUT" selection="${TEST_PATH}" - if [ -n "${UNIT_FLAG}" ]; then - selection="${TEST_PATH} $(bash .circleci/scripts/unit_selection.sh "${UNIT_FLAG}" | tr '\n' ' ')" - fi if [ -z "${selection// /}" ]; then echo "shard selection is empty; nothing to run" exit 0 diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index da4477b6947..c4397c2067b 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -20,13 +20,6 @@ concurrency: # rather than alphabetical letter ranges. Adding a new test file means adding it # to whichever group it belongs to, not reshuffling slices. # -# `.circleci/tests.yml` runs each group's files on same-repo events under the -# `proxy-db-` Codecov flag; `.circleci/scripts/unit_selection.sh` holds -# the file lists. That pipeline is manual-only while the tests migrate, so -# `unit-flag` makes the shard run that list on every event. `test-path` keeps -# the files that still reach real providers and never left -# tests/proxy_unit_tests. -# # Design targets: # * Every shard runs in <= 7 minutes of wall-clock on the default runner. # Most of a shard's time is pytest plugin load + xdist worker imports + @@ -78,93 +71,136 @@ jobs: include: # Must run serially — event-loop conflict with the logging worker. - test-group: key-generation - test-path: "" - unit-flag: proxy-db-key-generation + test-path: >- + tests/unit/proxy/management_endpoints/test_key_generate_prisma.py workers: 0 dist: loadscope timeout: 20 # ---- auth: split into 2 shards ---- - test-group: auth-checks - test-path: "" - unit-flag: proxy-db-auth-checks + test-path: >- + tests/unit/proxy/auth/test_auth_checks.py + tests/unit/proxy/auth/test_user_api_key_auth.py + tests/unit/proxy/test_credential_slot_registry.py + tests/unit/proxy/test_deprecated_key_grace_period.py workers: 4 dist: loadscope timeout: 15 - test-group: jwt-and-keys - test-path: "" - unit-flag: proxy-db-jwt-and-keys + test-path: >- + tests/unit/proxy/auth/test_jwt.py + tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py + tests/unit/proxy/test_proxy_custom_auth.py workers: 4 dist: loadscope timeout: 15 # ---- test_proxy_utils.py, single shard, worksteal distribution ---- - test-group: proxy-utils - test-path: "" - unit-flag: proxy-db-proxy-utils + test-path: >- + tests/unit/proxy/test_proxy_utils.py workers: 4 dist: worksteal timeout: 15 # ---- proxy server: split into 2 shards ---- - test-group: proxy-server-core - test-path: "tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py" - unit-flag: proxy-db-proxy-server-core + test-path: >- + tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py + tests/unit/proxy/test__lazy_features.py + tests/unit/proxy/test_aproxy_startup.py + tests/unit/proxy/test_proxy_server.py workers: 4 dist: loadscope timeout: 15 - test-group: proxy-runtime - test-path: "" - unit-flag: proxy-db-proxy-runtime + test-path: >- + tests/unit/proxy/auth/test_multipart_bypass_repro.py + tests/unit/proxy/auth/test_proxy_routes.py + tests/unit/proxy/middleware/test_request_size_limit_middleware.py + tests/unit/proxy/test_proxy_config_unit_test.py + tests/unit/proxy/test_proxy_token_counter.py + tests/unit/proxy/test_server_root_path.py workers: 4 dist: loadscope timeout: 15 # ---- logging: split into 2 shards ---- - test-group: custom-logging - test-path: "tests/proxy_unit_tests/test_proxy_custom_logger.py" - unit-flag: proxy-db-custom-logging + test-path: >- + tests/proxy_unit_tests/test_proxy_custom_logger.py + tests/unit/proxy/test_custom_callback_input.py + tests/unit/proxy/test_custom_logger_s3_gcs.py workers: 4 dist: loadscope timeout: 15 - test-group: logging-misc - test-path: "" - unit-flag: proxy-db-logging-misc + test-path: >- + tests/unit/proxy/management_helpers/test_audit_logs_proxy.py + tests/unit/proxy/spend_tracking/test_search_api_logging.py + tests/unit/proxy/test_proxy_reject_logging.py workers: 4 dist: loadscope timeout: 15 - test-group: db-and-spend - test-path: "" - unit-flag: proxy-db-db-and-spend + test-path: >- + tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py + tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py + tests/unit/proxy/db/test_update_daily_tag_spend.py + tests/unit/proxy/test_db_schema_changes.py + tests/unit/proxy/test_prisma_client_backoff_retry.py + tests/unit/proxy/test_update_spend.py + tests/unit/skills/test_skills_db.py workers: 4 dist: loadscope timeout: 15 # ---- guardrails + budget + hooks: split into 2 ---- - test-group: guardrails-hooks - test-path: "" - unit-flag: proxy-db-guardrails-hooks + test-path: >- + tests/unit/proxy/hooks/test_banned_keyword_list.py + tests/unit/proxy/test_proxy_setting_guardrails.py + tests/unit/proxy/test_unit_test_proxy_hooks.py workers: 4 dist: loadscope timeout: 15 - test-group: budgets - test-path: "" - unit-flag: proxy-db-budgets + test-path: >- + tests/unit/proxy/auth/test_default_end_user_budget_simple.py + tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py + tests/unit/proxy/test_zero_cost_model_budget_bypass.py workers: 4 dist: loadscope timeout: 15 - test-group: endpoints-and-responses - test-path: "tests/proxy_unit_tests/test_proxy_exception_mapping.py" - unit-flag: proxy-db-endpoints-and-responses + test-path: >- + tests/proxy_unit_tests/test_proxy_exception_mapping.py + tests/unit/proxy/lens + tests/unit/proxy/auth/test_models_fallback_endpoint.py + tests/unit/proxy/common_utils/test_check_batch_cost.py + tests/unit/proxy/common_utils/test_check_responses_cost.py + tests/unit/proxy/common_utils/test_realtime_cache.py + tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py + tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py + tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py + tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py + tests/unit/proxy/response_polling + tests/unit/proxy/test_custom_tokenizer_bug.py + tests/unit/proxy/test_get_favicon.py + tests/unit/proxy/test_get_image.py + tests/unit/proxy/test_prompt_test_endpoint.py + tests/unit/proxy/test_reducto_ocr_route.py + tests/unit/proxy/test_response_polling_pre_call_checks.py + tests/unit/proxy/test_ui_path_detection.py workers: 4 dist: loadscope timeout: 15 uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} - unit-flag: ${{ matrix.unit-flag }} workers: ${{ matrix.workers }} reruns: 2 timeout-minutes: ${{ matrix.timeout }} diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 3f99cee1720..5f9ea430069 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -35,10 +35,6 @@ concurrency: # a matrix and carries a shard-coverage guard that reads that file by name. # Folding it in here is a follow-up, together with generalising that guard into # assert_ci_coverage.py. -# -# `unit-flag` names the `.circleci/tests.yml` job that now runs part of the -# shard under the same Codecov flag. That pipeline is manual-only while the -# tests migrate, so the shard also runs those files on every event. jobs: unit: name: ${{ matrix.shard }} @@ -52,8 +48,9 @@ jobs: include: - shard: core-utils artifact-name: core-utils - test-path: tests/unit/decisions - unit-flag: core-utils + test-path: >- + tests/unit/decisions + tests/unit/litellm_core_utils workers: 2 reruns: 1 timeout-minutes: 20 @@ -61,8 +58,22 @@ jobs: - shard: enterprise-routing artifact-name: enterprise-routing - test-path: "" - unit-flag: enterprise-routing + test-path: >- + tests/unit/google_genai + tests/unit/router_strategy + tests/unit/router_utils + tests/unit/proxy/common_utils/test_cache_aware_routing.py + tests/unit/enterprise/enterprise_callbacks/send_emails + tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py + tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py + tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py + tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py + tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py + tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py + tests/unit/enterprise/proxy/test_enterprise_routes.py + tests/unit/enterprise/proxy/test_file_deletion_blocking.py + tests/unit/enterprise/proxy/test_managed_files_access_check.py + tests/unit/enterprise/proxy/test_managed_files_hook.py workers: 2 reruns: 2 timeout-minutes: 20 @@ -73,7 +84,7 @@ jobs: test-path: >- tests/test_litellm/integrations tests/test_litellm/tracing - unit-flag: integrations + tests/unit/integrations workers: 2 reruns: 3 timeout-minutes: 20 @@ -81,8 +92,8 @@ jobs: - shard: Vertex AI artifact-name: llm-vertex-ai - test-path: "" - unit-flag: llm-vertex-ai + test-path: >- + tests/unit/llms/vertex_ai workers: 1 reruns: 2 timeout-minutes: 20 @@ -90,8 +101,10 @@ jobs: - shard: All Other Providers artifact-name: llm-other-providers - test-path: "" - unit-flag: llm-other-providers + test-path: >- + tests/unit/llms + --ignore=tests/unit/llms/vertex_ai + --ignore=tests/unit/llms/base_llm/batches/base_batches_config_test.py workers: 2 reruns: 2 timeout-minutes: 20 @@ -101,7 +114,27 @@ jobs: artifact-name: misc test-path: >- tests/test_litellm/test_*.py - unit-flag: misc + tests/unit/test_*.py + tests/unit/test_router + tests/unit/a2a_protocol + tests/unit/batches + tests/unit/chat_completions + tests/unit/completion_extras + tests/unit/containers + tests/unit/embeddings + tests/unit/endpoints + tests/unit/files + tests/unit/harness + tests/unit/images + tests/unit/interactions + tests/unit/messages + tests/unit/rag + tests/unit/rerank_api + tests/unit/rust_bridge + tests/unit/secret_managers + tests/unit/vector_stores + tests/unit/videos + --ignore=tests/unit/rust_bridge/native_route_wheel_test.py workers: 2 reruns: 2 timeout-minutes: 20 @@ -209,7 +242,10 @@ jobs: tests/unit/proxy/enterprise_billing tests/unit/proxy/types_utils tests/unit/proxy/logging_endpoints - unit-flag: proxy-infra + tests/unit/gateway + tests/unit/proxy/management + tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py + tests/unit/proxy/roi_calculator workers: 4 reruns: 2 timeout-minutes: 20 @@ -251,8 +287,8 @@ jobs: - shard: caching-local artifact-name: caching-local - test-path: "" - unit-flag: caching-local + test-path: >- + tests/unit/caching workers: 2 reruns: 2 timeout-minutes: 20 @@ -260,8 +296,8 @@ jobs: - shard: proxy-extras artifact-name: proxy-extras - test-path: "" - unit-flag: proxy-extras + test-path: >- + tests/unit/litellm_proxy_extras workers: 2 reruns: 2 timeout-minutes: 20 @@ -269,8 +305,15 @@ jobs: - shard: enterprise-package artifact-name: enterprise-package - test-path: "" - unit-flag: enterprise-package + test-path: >- + tests/unit/enterprise/integrations + tests/unit/enterprise/proxy/auth + tests/unit/enterprise/proxy/guardrails + tests/unit/enterprise/proxy/hooks + tests/unit/enterprise/proxy/management_endpoints + tests/unit/enterprise/proxy/test_audit_logging_endpoints.py + tests/unit/enterprise/proxy/test_liteadmin.py + tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py workers: 4 reruns: 2 timeout-minutes: 20 @@ -278,16 +321,39 @@ jobs: - shard: responses-caching-types artifact-name: responses-caching-types - test-path: "" - unit-flag: responses-caching-types + test-path: >- + tests/unit/responses + tests/unit/types + --ignore=tests/unit/responses/mcp workers: 2 reruns: 2 timeout-minutes: 20 job-timeout-minutes: 60 + + - shard: unit + artifact-name: unit + test-path: >- + tests/unit/anthropic_interface + tests/unit/compression + tests/unit/enterprise/enterprise_callbacks/test_callback_controls.py + tests/unit/enterprise/enterprise_callbacks/test_llm_guard.py + tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py + tests/unit/integration_support + tests/unit/models + tests/unit/ocr + tests/unit/passthrough + tests/unit/realtime_api + tests/unit/repositories + tests/unit/sandbox + tests/unit/skills/test_skills_main.py + tests/unit/tracing + workers: 2 + reruns: 0 + timeout-minutes: 20 + job-timeout-minutes: 60 uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} - unit-flag: ${{ matrix.unit-flag || '' }} workers: ${{ matrix.workers }} reruns: ${{ matrix.reruns }} timeout-minutes: ${{ matrix.timeout-minutes }} diff --git a/tests/unit/test_assert_ci_coverage.py b/tests/unit/test_assert_ci_coverage.py index e4233dbe682..4498dbe1498 100644 --- a/tests/unit/test_assert_ci_coverage.py +++ b/tests/unit/test_assert_ci_coverage.py @@ -227,37 +227,39 @@ def test_every_sharded_root_named_in_the_script_exists_on_disk(): def test_the_repo_as_it_stands_has_every_shard_child_assigned(): findings = coverage._unassigned_shard_children( - coverage._shard_tokens(coverage._all_scalars(), coverage._unit_selection_arms()) + coverage._invoked_test_tokens(coverage._all_scalars()) ) assert [f.subject for f in findings] == [] -def test_shard_tokens_credits_only_wired_unit_flags(tmp_path): - root = tmp_path / "tests" / "tree" - (root / "wired").mkdir(parents=True) - (root / "wired" / "test_a.py").write_text("def test_a(): assert True\n") - (root / "unwired").mkdir(parents=True) - (root / "unwired" / "test_b.py").write_text("def test_b(): assert True\n") - script = tmp_path / ".circleci" / "scripts" / "unit_selection.sh" - script.parent.mkdir(parents=True) - script.write_text( - "legacy_paths() {\n" - " case \"$1\" in\n" - " wired-flag) echo tests/tree/wired ;;\n" - " unwired-flag)\n" - " echo tests/tree/unwired ;;\n" - " esac\n" - "}\n" +def test_shards_are_credited_only_from_explicit_test_paths() -> None: + scalars: Final = ( + coverage.Scalar(key="test-path", value="tests/unit/wired\ntests/unit/also_wired"), + coverage.Scalar(key="shard", value="tests/unit/not_a_test_path"), ) + assert coverage._invoked_test_tokens(scalars) == frozenset({"tests/unit/wired", "tests/unit/also_wired"}) - scalars: Final = (coverage.Scalar(key="unit-flag", value="wired-flag"),) - findings = coverage._unassigned_shard_children( - coverage._shard_tokens(scalars, coverage._unit_selection_arms(tmp_path)), - roots=("tests/tree",), - repo_root=tmp_path, + +def test_an_ignored_path_is_not_credited_as_invoked() -> None: + scalars: Final = ( + coverage.Scalar(key="test-path", value="tests/unit/a\n--ignore=tests/unit/b/test_x.py"), ) + assert coverage._invoked_test_tokens(scalars) == frozenset({"tests/unit/a"}) - assert tuple(f.subject for f in findings) == ("tests/tree/unwired",) + +def test_a_file_its_only_shard_ignores_is_not_covered_by_that_shards_glob() -> None: + selections: Final = coverage._invoked_selections( + ( + coverage.Scalar( + key="test-path", + value="tests/unit/proxy/test_*.py --ignore=tests/unit/proxy/test_update_spend.py", + ), + ) + ) + assert len(selections) == 1 + selection: Final = selections[0] + assert selection.covers("tests/unit/proxy/test_other.py") is True + assert selection.covers("tests/unit/proxy/test_update_spend.py") is False def test_check_shards_passes_on_the_repo_as_it_stands(capsys): diff --git a/tests/unit/test_unit_shard_missing_paths.py b/tests/unit/test_unit_shard_missing_paths.py index 0360a227142..4fa9c5bd3c1 100644 --- a/tests/unit/test_unit_shard_missing_paths.py +++ b/tests/unit/test_unit_shard_missing_paths.py @@ -38,7 +38,6 @@ def _run_shard(tmp_path: Path, test_path: str, workers: str) -> subprocess.Compl "PATH": f"{shim_dir}{os.pathsep}{os.environ['PATH']}", "GITHUB_OUTPUT": str(tmp_path / "github_output"), "TEST_PATH": test_path, - "UNIT_FLAG": "", "WORKERS": workers, }, capture_output=True, From 3b2ed831523d77a0919d1fb169e3261c270beda9 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 5 Oct 2026 10:56:08 -0700 Subject: [PATCH 005/191] fix(ui): make mobile sidebar and top bar responsive (#44603) * fix(ui): improve gateway layouts on mobile * fix(ui): limit mobile cleanup to navigation and header * fix(ui): share cached settings with mobile navigation --- .../components/SidebarProvider.test.tsx | 45 ++++++++ .../components/SidebarProvider.tsx | 67 ++--------- .../src/app/(dashboard)/layout.test.tsx | 54 ++++++++- .../src/app/(dashboard)/layout.tsx | 68 ++++++++--- .../DashboardHeader.integration.test.tsx | 109 ++++++++++++++++++ .../src/components/DashboardHeader.test.tsx | 26 +++++ .../src/components/DashboardHeader.tsx | 41 ++++++- .../src/components/shared/Sidebar.tsx | 2 +- 8 files changed, 328 insertions(+), 84 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.test.tsx create mode 100644 ui/litellm-dashboard/src/components/DashboardHeader.integration.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.test.tsx new file mode 100644 index 00000000000..a3d75df9606 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.test.tsx @@ -0,0 +1,45 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, within } from "@testing-library/react"; +import type { ComponentProps } from "react"; +import { expect, it, vi } from "vitest"; +import Sidebar from "@/components/leftnav"; +import SidebarProvider from "./SidebarProvider"; + +const { getSettings } = vi.hoisted(() => ({ + getSettings: vi.fn().mockResolvedValue({ values: { enable_projects_ui: true } }), +})); + +vi.mock("@/components/networking", () => ({ getUiSettings: getSettings, getUISettings: getSettings })); +vi.mock("@/components/leftnav", () => ({ + default: ({ enableProjectsUI }: ComponentProps) => ( + + ), +})); + +it("reuses loaded navigation settings immediately when the mobile sidebar mounts again", async () => { + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const shell = (mobileOpen: boolean) => ( + +
+ +
+ {mobileOpen && ( +
+ +
+ )} +
+ ); + const { rerender } = render(shell(false)); + expect(await screen.findByRole("link", { name: "Projects" })).toBeVisible(); + + for (const open of [true, false, true]) { + rerender(shell(open)); + if (open) { + expect( + within(screen.getByRole("region", { name: "Mobile navigation" })).getByRole("link", { name: "Projects" }), + ).toBeVisible(); + } + } + expect(getSettings).toHaveBeenCalledTimes(1); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx index 6aaf08dae79..c7794233ded 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx @@ -1,9 +1,7 @@ "use client"; import Sidebar from "@/components/leftnav"; -import { getUISettings } from "@/components/networking"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useEffect, useState } from "react"; +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; interface SidebarProviderProps { sidebarCollapsed: boolean; @@ -11,66 +9,19 @@ interface SidebarProviderProps { } const SidebarProvider = ({ sidebarCollapsed, onToggleCollapsed }: SidebarProviderProps) => { - const { accessToken } = useAuthorized(); - const [enabledPagesInternalUsers, setEnabledPagesInternalUsers] = useState(null); - const [enableProjectsUI, setEnableProjectsUI] = useState(false); - const [disableAgentsForInternalUsers, setDisableAgentsForInternalUsers] = useState(false); - const [allowAgentsForTeamAdmins, setAllowAgentsForTeamAdmins] = useState(false); - const [disableVectorStoresForInternalUsers, setDisableVectorStoresForInternalUsers] = useState(false); - const [allowVectorStoresForTeamAdmins, setAllowVectorStoresForTeamAdmins] = useState(false); - - useEffect(() => { - const fetchUISettings = async () => { - if (!accessToken) { - return; - } - - try { - const settings = await getUISettings(accessToken); - - // API returns 'values' not 'settings' - if (settings?.values?.enabled_ui_pages_internal_users !== undefined) { - setEnabledPagesInternalUsers(settings.values.enabled_ui_pages_internal_users); - } else { - } - - if (settings?.values?.enable_projects_ui !== undefined) { - setEnableProjectsUI(Boolean(settings.values.enable_projects_ui)); - } - - if (settings?.values?.disable_agents_for_internal_users !== undefined) { - setDisableAgentsForInternalUsers(Boolean(settings.values.disable_agents_for_internal_users)); - } - - if (settings?.values?.allow_agents_for_team_admins !== undefined) { - setAllowAgentsForTeamAdmins(Boolean(settings.values.allow_agents_for_team_admins)); - } - - if (settings?.values?.disable_vector_stores_for_internal_users !== undefined) { - setDisableVectorStoresForInternalUsers(Boolean(settings.values.disable_vector_stores_for_internal_users)); - } - - if (settings?.values?.allow_vector_stores_for_team_admins !== undefined) { - setAllowVectorStoresForTeamAdmins(Boolean(settings.values.allow_vector_stores_for_team_admins)); - } - } catch (error) { - console.error("[SidebarProvider] Failed to fetch UI settings:", error); - } - }; - - fetchUISettings(); - }, [accessToken]); + const { data: settings } = useUISettings(); + const values = settings?.values; return ( ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 7cdf7aa0489..876a4e24d84 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import { usePathname } from "next/navigation"; import { AuthProvider } from "@/contexts/AuthContext"; import Layout from "./layout"; @@ -19,12 +19,17 @@ vi.mock("@/components/liteadmin/LiteAdmin", () => ({ })); vi.mock("@/components/DashboardHeader", () => ({ - DashboardHeader: () =>
, + DashboardHeader: ({ navigationTrigger }: { navigationTrigger?: React.ReactNode }) => ( +
{navigationTrigger}
+ ), })); vi.mock("@/app/(dashboard)/components/SidebarProvider", () => ({ - default: ({ sidebarCollapsed }: { sidebarCollapsed: boolean }) => ( -
+ default: ({ sidebarCollapsed, onToggleCollapsed }: { sidebarCollapsed: boolean; onToggleCollapsed: () => void }) => ( +
+ + Settings +
), })); @@ -89,6 +94,47 @@ describe("(dashboard) Layout", () => { vi.mocked(usePathname).mockReturnValue("/ui/guardrails"); }); + it("starts mobile navigation closed, opens a modal drawer and closes it after choosing a page", async () => { + render( + + +

Gateway content

+
+
, + ); + pendingUiConfig.resolve(); + + const trigger = await screen.findByRole("button", { name: "Open navigation" }); + expect(screen.queryByRole("dialog", { name: "Navigation" })).not.toBeInTheDocument(); + + fireEvent.click(trigger); + const navigation = await screen.findByRole("dialog", { name: "Navigation" }); + expect(within(navigation).getByRole("button", { name: "Close navigation" })).toBeInTheDocument(); + fireEvent.click(within(navigation).getByRole("link", { name: "Settings" })); + await waitFor(() => expect(screen.queryByRole("dialog", { name: "Navigation" })).not.toBeInTheDocument()); + }); + + it("closes the mobile drawer when navigation changes outside the drawer", async () => { + const dashboard = () => ( + + +

Gateway content

+
+
+ ); + const { rerender } = render(dashboard()); + pendingUiConfig.resolve(); + fireEvent.click(await screen.findByRole("button", { name: "Open navigation" })); + expect(await screen.findByRole("dialog", { name: "Navigation" })).toBeInTheDocument(); + + vi.mocked(usePathname).mockReturnValue("/ui/api-keys"); + rerender(dashboard()); + await waitFor(() => expect(screen.queryByRole("dialog", { name: "Navigation" })).not.toBeInTheDocument()); + vi.mocked(usePathname).mockReturnValue("/ui/guardrails"); + rerender(dashboard()); + expect(screen.queryByRole("dialog", { name: "Navigation" })).not.toBeInTheDocument(); + }); + it("collapses the sidebar on Logs for a full-screen view and expands it again after leaving", async () => { const dashboard = () => ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index d705089cee8..4762dd2fed1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -19,6 +19,10 @@ import { routeSegmentForPathname, uiHref } from "@/utils/uiHref"; import { PluginModeProvider, usePluginMode } from "@/contexts/PluginModeContext"; import { createApiClient } from "@/lib/http/client"; import { getProxyBaseUrl } from "@/components/networking"; +import { Sheet, SheetContent, SheetTitle, SheetTrigger } from "@/components/ui/sheet"; +import { Button } from "@/components/ui/button"; +import { Menu } from "lucide-react"; +import { useMediaQuery } from "usehooks-ts"; const pluginApiClient = createApiClient({ getBaseUrl: () => getProxyBaseUrl() ?? "" }); @@ -104,7 +108,16 @@ const FULL_BLEED_SEGMENTS = new Set(["logs"]); function DashboardShell({ children }: { children: React.ReactNode }) { const { accessToken } = useAuth(); const { mode } = usePluginMode(); - const routeSegment = routeSegmentForPathname(usePathname()); + const pathname = usePathname(); + const routeSegment = routeSegmentForPathname(pathname); + const searchParams = useSearchParams(); + const navigationKey = `${pathname}?${searchParams.toString()}`; + const isDesktop = useMediaQuery("(min-width: 768px)", { initializeWithValue: false }); + const [mobileNavigationKey, setMobileNavigationKey] = useState(null); + if (mobileNavigationKey !== null && (isDesktop || mobileNavigationKey !== navigationKey)) { + setMobileNavigationKey(null); + } + const mobileNavigationOpen = !isDesktop && mobileNavigationKey === navigationKey; const isFullBleed = FULL_BLEED_SEGMENTS.has(routeSegment); // A manual toggle holds only for the route it was made on; full-bleed routes default to collapsed. const [sidebarOverride, setSidebarOverride] = useState<{ segment: string; collapsed: boolean } | null>(null); @@ -138,21 +151,46 @@ function DashboardShell({ children }: { children: React.ReactNode }) { // sidebar owns its own scroll and the content column scrolls independently, // so the page can't be dragged past the end of the nav. return ( -
- - -
- - - - - - - -
{children}
+ setMobileNavigationKey(open ? navigationKey : null)}> +
+
+
- -
+ { + if (event.target instanceof Element && event.target.closest("a[href]")) setMobileNavigationKey(null); + }} + > + Navigation + setMobileNavigationKey(null)} /> + + +
+ + } + > + + + } + /> + + + + + + +
{children}
+
+
+
+ ); } diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.integration.test.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.integration.test.tsx new file mode 100644 index 00000000000..756997b7ede --- /dev/null +++ b/ui/litellm-dashboard/src/components/DashboardHeader.integration.test.tsx @@ -0,0 +1,109 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { PluginModeProvider } from "@/contexts/PluginModeContext"; +import { DashboardHeader } from "./DashboardHeader"; + +vi.mock("next/navigation", () => ({ usePathname: () => "/ui/logs" })); + +beforeEach(() => { + localStorage.clear(); + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith("/get/ui_settings")) return Response.json({ values: {} }); + if (url.endsWith("/litellm/.well-known/litellm-ui-config")) { + const config = { + server_root_path: "", + proxy_base_url: "", + is_control_plane: true, + workers: [ + { worker_id: "primary", name: "Primary worker", url: "http://primary.test" }, + { worker_id: "backup", name: "Backup worker", url: "http://backup.test" }, + ], + }; + return Response.json(config); + } + if (url.endsWith("/api/plugins")) { + return Response.json([{ name: "metrics", display_name: "Metrics", url: "/metrics" }]); + } + if (url.endsWith("/public/litellm_blog_posts")) { + return Response.json({ + posts: [ + { + title: "Gateway update", + description: "Release notes", + date: "2026-01-01", + url: "https://example.com/update", + }, + ], + }); + } + return Response.json({}, { status: 404 }); + }), + ); +}); + +afterEach(() => { + vi.unstubAllGlobals(); + localStorage.clear(); +}); + +async function openTools() { + const user = userEvent.setup(); + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + render( + + + + + , + ); + await user.click(screen.getByRole("button", { name: "More options" })); + const tools = await screen.findByRole("dialog", { name: "Gateway tools" }); + return { user, tools }; +} + +describe("mobile header nested controls", () => { + it.each(["mouse", "touch"] as const)( + "selects a portalled gateway item with %s without dismissing More", + async (pointer) => { + const { user, tools } = await openTools(); + await user.click(within(tools).getByRole("button", { name: "AI Gateway" })); + const item = await screen.findByRole("menuitem", { name: "Metrics" }); + if (pointer === "touch") { + await user.pointer([ + { keys: "[TouchA>]", target: item }, + { keys: "[/TouchA]", target: item }, + ]); + } else { + await user.click(item); + } + expect(within(tools).getByRole("button", { name: "Metrics" })).toBeVisible(); + expect(screen.getByRole("dialog", { name: "Gateway tools" })).toBeVisible(); + }, + ); + + it("lets a nested blog link receive the click before closing its menu", async () => { + const { user, tools } = await openTools(); + await user.click(within(tools).getByRole("button", { name: "Blog" })); + const link = await screen.findByRole("link", { name: /Gateway update/ }); + expect(link).toHaveAttribute("href", "https://example.com/update"); + await user.click(link); + expect(screen.queryByRole("menu", { name: "Blog" })).not.toBeInTheDocument(); + expect(screen.getByRole("dialog", { name: "Gateway tools" })).toBeVisible(); + }); + + it("allows filtering and selecting a worker from the nested combobox", async () => { + localStorage.setItem("litellm_selected_worker_id", "primary"); + const { user, tools } = await openTools(); + const worker = await within(tools).findByRole("combobox", { name: "Worker" }); + await user.clear(worker); + await user.type(worker, "Backup"); + await user.click(await screen.findByRole("option", { name: "Backup worker" })); + expect(localStorage.getItem("litellm_selected_worker_id")).toBeNull(); + expect(screen.getByRole("dialog", { name: "Gateway tools" })).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx index 751a0192b8a..63f7b7aa136 100644 --- a/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx @@ -8,6 +8,7 @@ const { mockUsePluginMode, mockUseUISettings, state } = vi.hoisted(() => { plugins: [] as { name: string; display_name: string; url: string }[], enableChatUI: false, pathname: "/ui/logs", + isDesktop: false, }; return { state, @@ -19,6 +20,7 @@ const { mockUsePluginMode, mockUseUISettings, state } = vi.hoisted(() => { vi.mock("@/contexts/PluginModeContext", () => ({ usePluginMode: mockUsePluginMode })); vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: mockUseUISettings })); vi.mock("next/navigation", () => ({ usePathname: () => state.pathname })); +vi.mock("usehooks-ts", () => ({ useMediaQuery: () => state.isDesktop })); vi.mock("@/hooks/useWorker", () => ({ useWorker: () => ({ isControlPlane: false, selectedWorker: null }) })); vi.mock("@/app/(dashboard)/hooks/useDisableShowPrompts", () => ({ useDisableShowPrompts: () => false })); vi.mock("@/components/Navbar/BlogDropdown/BlogDropdown", () => ({ BlogDropdown: () => null })); @@ -34,6 +36,7 @@ describe("DashboardHeader breadcrumb", () => { state.plugins = []; state.enableChatUI = false; state.pathname = "/ui/logs"; + state.isDesktop = false; }); it("titles the breadcrumb from the current route, not from a sidebar page id", () => { @@ -99,4 +102,27 @@ describe("DashboardHeader breadcrumb", () => { Node.DOCUMENT_POSITION_FOLLOWING, ); }); + + it("keeps the gateway selector and tools available from the compact header menu", async () => { + render(); + fireEvent.click(screen.getByRole("button", { name: "More options" })); + const tools = await screen.findByRole("dialog", { name: "Gateway tools" }); + expect(within(tools).getByRole("button", { name: "AI Gateway" })).toBeInTheDocument(); + expect(within(tools).getByRole("link", { name: "Docs" })).toBeInTheDocument(); + expect(within(tools).getByRole("button", { name: "LiteAdmin" })).toBeInTheDocument(); + }); + + it("closes mobile tools when switching to desktop and keeps them closed when returning", async () => { + const { rerender } = render(); + fireEvent.click(screen.getByRole("button", { name: "More options" })); + expect(await screen.findByRole("dialog", { name: "Gateway tools" })).toBeInTheDocument(); + + state.isDesktop = true; + rerender(); + expect(screen.queryByRole("dialog", { name: "Gateway tools" })).not.toBeInTheDocument(); + + state.isDesktop = false; + rerender(); + expect(screen.queryByRole("dialog", { name: "Gateway tools" })).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.tsx index e4107669336..43440b0c92d 100644 --- a/ui/litellm-dashboard/src/components/DashboardHeader.tsx +++ b/ui/litellm-dashboard/src/components/DashboardHeader.tsx @@ -22,11 +22,19 @@ import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPro import { clearTokenCookies } from "@/utils/cookieUtils"; import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils"; import { usePathname } from "next/navigation"; +import { useState, type ReactNode } from "react"; +import { useMediaQuery } from "usehooks-ts"; +import { Ellipsis } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Popover, PopoverContent, PopoverTitle, PopoverTrigger } from "@/components/ui/popover"; // Top bar for the dashboard shell. Sits only over the content column (the brand // lives in the sidebar header); mirrors the design's breadcrumb-left / tools-right layout. -export function DashboardHeader() { +export function DashboardHeader({ navigationTrigger }: { navigationTrigger?: ReactNode }) { const { title } = getBreadcrumb(usePathname()); + const isDesktop = useMediaQuery("(min-width: 768px)", { initializeWithValue: false }); + const [mobileToolsOpen, setMobileToolsOpen] = useState(false); + if (isDesktop && mobileToolsOpen) setMobileToolsOpen(false); const { isControlPlane, selectedWorker } = useWorker(); const showWorkerSwitch = isControlPlane && selectedWorker !== null; const hideCommunityLinks = useDisableShowPrompts(); @@ -40,20 +48,21 @@ export function DashboardHeader() { }; return ( -
- +
+ {navigationTrigger} + - + - + {title} -
+
{showWorkerSwitch && ( <> @@ -68,6 +77,26 @@ export function DashboardHeader() {
+
+ + + }> + + + + Gateway tools + + {showWorkerSwitch && } +
+ + + + + {!hideCommunityLinks && } +
+
+
+
); } diff --git a/ui/litellm-dashboard/src/components/shared/Sidebar.tsx b/ui/litellm-dashboard/src/components/shared/Sidebar.tsx index 06e19d1b674..d57d54de789 100644 --- a/ui/litellm-dashboard/src/components/shared/Sidebar.tsx +++ b/ui/litellm-dashboard/src/components/shared/Sidebar.tsx @@ -139,7 +139,7 @@ SidebarMenuBadge.displayName = "SidebarMenuBadge"; const sidebarMenuButtonVariants = cva( [ - "group/menu-btn relative flex w-full items-center gap-2.5 overflow-hidden rounded-md px-2.5 text-left text-[13px] font-medium no-underline", + "group/menu-btn relative flex w-full items-center gap-2.5 overflow-hidden rounded-md px-2.5 text-left text-[13px] font-medium no-underline max-md:min-h-11", "text-sidebar-foreground/70 outline-none transition-colors", "hover:bg-sidebar-accent hover:text-sidebar-accent-foreground", "focus-visible:ring-2 focus-visible:ring-sidebar-ring", From d80f8c28ca7e2fba4257b4b97457d3b313ff0d6a Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 5 Oct 2026 18:01:45 +0000 Subject: [PATCH 006/191] refactor(types): replace Any with proven types in 137 files (#44478) * refactor(types): prove runtime types at harness, search, rag and client boundaries Replace Any with adapter-validated types in the litellm.agent() harness, the search provider transformations, RAG ingestion and query, the vector store pre-call hook and registry, the galileo and opik logging integrations and the proxy client CLI. Each boundary gets unit tests for well-formed and malformed payloads. * chore(typing): prove types at more provider boundaries and restore search transformations Second pass over a2a, embedding, rerank, image, audio and small provider modules. Search transformations go back to their previous form because validating their response bodies would change the proxy status for malformed upstream bodies from 400 to 500. * refactor(types): prove types at logging, files, rerank, image, audio and management boundaries Replace Any with validated or annotation-only types in 37 more files: logging integrations, token counters, provider files/rerank/image generation/audio transcription transformations, pass-through logging handlers and management endpoints. No proxy HTTP status or error type changes. * refactor(types): prove types at repository, spend, files and router boundaries Replace Any with repository table accessors, validated mappings and annotation-only types in 33 more files: Prisma repositories, the enterprise batch and responses cost checkers, budget reservation, files endpoints, management endpoints, the policy registry, the adaptive and complexity routers and the secret managers. No proxy HTTP status or error type changes. * test(types): run the aiohttp transformation test in-process and cover repository row conversion The aiohttp chat transformation test no longer starts a server. It feeds the transformation a response whose json() returns the body under test. The proxy unit shards now exercise stored model rows whose params are JSON strings and the object permission create and update paths. --- .../send_emails/base_email.py | 6 +- .../proxy/common_utils/check_batch_cost.py | 18 +- .../common_utils/check_responses_cost.py | 4 +- .../proxy/hooks/managed_files.py | 4 +- litellm/a2a_protocol/providers/base.py | 9 +- .../providers/bedrock_agentcore/config.py | 10 +- .../providers/bedrock_agentcore/handler.py | 10 +- .../bedrock_agentcore/transformation.py | 2 +- .../a2a_protocol/providers/langflow/config.py | 9 +- .../providers/pydantic_ai_agents/config.py | 16 +- .../providers/pydantic_ai_agents/handler.py | 16 +- .../pydantic_ai_agents/transformation.py | 12 +- .../providers/watsonx_orchestrate/config.py | 16 +- .../providers/watsonx_orchestrate/handler.py | 4 +- litellm/caching/_embedding_router.py | 8 +- litellm/files/main.py | 2 +- litellm/files/streaming.py | 4 +- litellm/harness/context.py | 6 +- litellm/harness/handlers/cli_handler.py | 4 +- litellm/harness/options.py | 8 +- litellm/harness/runtime.py | 42 ++-- litellm/harness/sync.py | 44 ++-- litellm/harness/types.py | 6 +- litellm/integrations/arize/_utils.py | 15 +- .../integrations/datadog/datadog_llm_obs.py | 8 +- litellm/integrations/galileo.py | 15 +- litellm/integrations/opentelemetry.py | 8 +- .../opik/opik_payload_builder/api.py | 11 +- .../opik/opik_payload_builder/extractors.py | 6 +- .../opik/opik_payload_builder/types.py | 15 +- litellm/integrations/opik/utils.py | 9 +- litellm/integrations/otel/model/config.py | 4 +- litellm/integrations/otel/plumbing/metrics.py | 4 +- litellm/integrations/s3_v2.py | 4 +- .../vector_store_pre_call_hook.py | 20 +- litellm/llms/a2a/chat/transformation.py | 14 +- .../aiohttp_openai/chat/transformation.py | 10 +- .../anthropic/count_tokens/token_counter.py | 2 +- .../llms/anthropic/files/transformation.py | 15 +- .../audio_transcription/transformation.py | 18 +- .../anthropic/count_tokens/handler.py | 2 +- .../llms/base_llm/harness/transformation.py | 4 +- litellm/llms/bedrock/batches/handler.py | 16 +- .../image_edit/stability_transformation.py | 3 +- .../amazon_titan_transformation.py | 10 +- litellm/llms/bedrock/rerank/handler.py | 9 +- litellm/llms/bedrock/search/transformation.py | 7 +- .../image_generation/transformation.py | 8 +- litellm/llms/clarifai/chat/transformation.py | 6 +- .../claude_code/harness/transformation.py | 8 +- litellm/llms/codex/harness/transformation.py | 16 +- litellm/llms/cohere/rerank/transformation.py | 9 +- .../image_generation/transformation.py | 8 +- .../image_generation/transformation.py | 16 +- .../llms/dashscope/rerank/transformation.py | 23 +- .../llms/deepagents/harness/transformation.py | 6 +- .../audio_transcription/transformation.py | 9 +- litellm/llms/e2b/sandbox/transformation.py | 36 ++- .../audio_transcription/transformation.py | 14 +- .../flux_pro_v11_ultra_transformation.py | 19 +- .../ideogram_v3_transformation.py | 11 +- .../recraft_v3_transformation.py | 8 +- .../stable_diffusion_transformation.py | 21 +- .../fireworks_ai/rerank/transformation.py | 55 +++-- .../gemini/image_generation/transformation.py | 11 +- .../llms/gigachat/embedding/transformation.py | 6 +- .../embedding/transformation.py | 5 +- .../llms/jina_ai/embedding/transformation.py | 6 +- litellm/llms/jina_ai/rerank/transformation.py | 25 +- litellm/llms/manus/files/transformation.py | 40 +-- .../minimax/text_to_speech/transformation.py | 10 +- .../image_generation/transformation.py | 13 +- .../llms/nvidia_nim/rerank/transformation.py | 26 +- .../audio_transcription/transformation.py | 10 +- litellm/llms/oci/embed/transformation.py | 6 +- .../openai/responses/count_tokens/handler.py | 5 +- .../llms/opencode/harness/transformation.py | 8 +- .../image_generation/transformation.py | 20 +- .../audio_transcription/transformation.py | 11 +- .../image_generation/transformation.py | 8 +- litellm/llms/replicate/chat/transformation.py | 6 +- litellm/llms/sagemaker/common_utils.py | 7 +- .../snowflake/embedding/transformation.py | 6 +- .../image_generation/transformation.py | 9 +- .../llms/tinyfish/search/transformation.py | 5 +- litellm/llms/together_ai/rerank/handler.py | 14 +- .../image_generation_handler.py | 11 +- .../vertex_gemini_transformation.py | 12 +- .../text_to_speech/transformation.py | 10 +- litellm/llms/voyage/rerank/transformation.py | 19 +- .../audio_transcription/transformation.py | 13 +- litellm/llms/watsonx/embed/transformation.py | 9 +- litellm/llms/watsonx/rerank/transformation.py | 17 +- litellm/proxy/_experimental/mcp_server/db.py | 11 +- .../proxy/agent_endpoints/a2a_endpoints.py | 7 +- litellm/proxy/client/chat.py | 7 +- litellm/proxy/client/cli/commands/auth.py | 5 +- litellm/proxy/client/cli/commands/debug.py | 3 +- litellm/proxy/db/model_insights_tasks.py | 10 +- .../proxy/fine_tuning_endpoints/endpoints.py | 2 +- .../health_endpoints/_health_endpoints.py | 14 +- .../auto_router_endpoints.py | 8 +- .../callback_management_endpoints.py | 5 +- .../key_management_endpoints.py | 5 +- .../management_v1/budgets.py | 6 +- .../mcp_management_endpoints.py | 6 +- .../model_insights_endpoints.py | 16 +- .../model_management_endpoints.py | 10 +- .../organization_endpoints.py | 18 +- .../policy_endpoints/endpoints.py | 37 ++- .../router_settings_endpoints.py | 2 +- .../tag_management_endpoints.py | 7 +- .../management_endpoints/team_endpoints.py | 2 +- .../file_content_streaming_handler.py | 6 +- .../openai_files_endpoints/files_endpoints.py | 3 +- .../anthropic_passthrough_logging_handler.py | 7 +- .../vertex_passthrough_logging_handler.py | 7 +- .../proxy/policy_engine/policy_registry.py | 4 +- .../public_endpoints/public_endpoints.py | 7 +- .../spend_tracking/budget_reservation.py | 19 +- .../spend_tracking/spend_log_error_logger.py | 4 +- litellm/rag/ingestion/base_ingestion.py | 6 +- litellm/rag/ingestion/gemini_ingestion.py | 16 +- litellm/rag/ingestion/openai_ingestion.py | 6 +- litellm/rag/ingestion/s3_vectors_ingestion.py | 4 +- litellm/rag/main.py | 36 ++- .../autorouter_session_repository.py | 15 +- litellm/repositories/model_repository.py | 28 ++- .../object_permission_repository.py | 22 +- litellm/repositories/project_repository.py | 15 +- litellm/repositories/user_repository.py | 18 +- .../adaptive_router/adaptive_router.py | 3 + .../complexity_router/complexity_router.py | 8 +- .../reasoning_effort_capability.py | 2 +- litellm/secret_managers/aws_secret_manager.py | 6 +- litellm/secret_managers/main.py | 6 +- .../vector_stores/vector_store_registry.py | 8 +- .../test_pydantic_ai_agent_headers.py | 71 +++++- .../watsonx_orchestrate/test_config.py | 161 ++++++++++++ .../send_emails/test_base_email.py | 75 +++++- tests/unit/files/test_main.py | 52 ++++ .../integrations/arize/test_arize_utils.py | 41 ++++ .../datadog/test_datadog_llm_obs.py | 28 +++ .../opik/opik_payload_builder/__init__.py | 0 .../opik/opik_payload_builder/test_api.py | 230 ++++++++++++++++++ .../integrations/opik/test_opik_extractors.py | 32 +++ tests/unit/integrations/test_galileo.py | 79 ++++++ tests/unit/integrations/test_opentelemetry.py | 34 +++ tests/unit/integrations/test_opik_utils.py | 21 +- .../test_vector_store_pre_call_hook.py | 163 ++++++++++++- .../a2a/chat/test_a2a_chat_transformation.py | 36 +++ tests/unit/llms/aiohttp_openai/__init__.py | 0 .../unit/llms/aiohttp_openai/chat/__init__.py | 0 .../chat/test_transformation.py | 101 ++++++++ .../test_anthropic_files_transformation.py | 107 ++++++++ .../test_azure_speech_audio_transcription.py | 20 ++ .../test_stability_transformation.py | 14 ++ .../llms/bedrock/image_generation/__init__.py | 0 .../test_amazon_titan_transformation.py | 57 +++++ .../test_bedrock_rerank_header_forwarding.py | 59 +++++ .../test_agentcore_search_transformation.py | 43 ++++ ...est_bfl_image_generation_transformation.py | 46 ++++ tests/unit/llms/clarifai/__init__.py | 0 tests/unit/llms/clarifai/chat/__init__.py | 0 .../llms/clarifai/chat/test_transformation.py | 70 ++++++ .../harness/test_transformation.py | 23 ++ .../llms/codex/harness/test_transformation.py | 32 ++- .../llms/cohere/rerank/test_transformation.py | 72 ++++++ .../cometapi/image_generation/__init__.py | 0 .../image_generation/test_transformation.py | 54 ++++ .../test_dashscope_rerank_transformation.py | 87 +++++++ ...gram_audio_transcription_transformation.py | 80 ++++++ tests/unit/llms/e2b/__init__.py | 0 tests/unit/llms/e2b/sandbox/__init__.py | 0 .../llms/e2b/sandbox/test_transformation.py | 121 +++++++++ .../audio_transcription/__init__.py | 0 .../test_transformation.py | 100 ++++++++ .../test_flux_pro_v11_ultra_transformation.py | 56 +++++ .../test_ideogram_v3_transformation.py | 71 ++++++ .../test_recraft_v3_transformation.py | 64 +++++ .../test_stable_diffusion_transformation.py | 76 ++++++ ...test_fireworks_ai_rerank_transformation.py | 66 +++++ ..._gemini_image_generation_transformation.py | 48 ++++ .../test_gigachat_embedding_transformation.py | 63 ++++- ...github_copilot_embedding_transformation.py | 57 +++++ .../test_jina_embedding_transformation.py | 61 +++++ tests/unit/llms/jina_ai/rerank/__init__.py | 0 .../jina_ai/rerank/test_transformation.py | 111 +++++++++ tests/unit/llms/manus/files/__init__.py | 0 .../llms/manus/files/test_transformation.py | 128 ++++++++++ .../llms/minimax/text_to_speech/__init__.py | 0 .../text_to_speech/test_transformation.py | 122 ++++++++++ ...est_modelscope_image_gen_transformation.py | 83 +++++++ .../test_nvidia_nim_rerank_transformation.py | 86 +++++++ .../embed/test_oci_embed_transformation.py | 75 ++++++ .../openai/responses/count_tokens/__init__.py | 0 .../responses/count_tokens/test_handler.py | 26 ++ ...est_openrouter_image_gen_transformation.py | 83 +++++++ ...loud_audio_transcription_transformation.py | 42 ++++ .../test_recraft_image_gen_transformation.py | 54 ++++ tests/unit/llms/replicate/__init__.py | 0 tests/unit/llms/replicate/chat/__init__.py | 0 .../replicate/chat/test_transformation.py | 54 ++++ .../sagemaker/test_sagemaker_common_utils.py | 13 + .../embedding/test_snowflake_embedding.py | 64 ++++- .../test_stability_image_generation.py | 33 +++ .../llms/tinyfish/test_tinyfish_search.py | 54 ++++ .../unit/llms/together_ai/rerank/__init__.py | 0 .../llms/together_ai/rerank/test_handler.py | 96 ++++++++ .../test_image_generation_handler.py | 43 ++++ ...rtex_ai_image_generation_transformation.py | 64 +++++ .../text_to_speech/test_transformation.py | 30 +++ .../test_voyage_rerank_transformation.py | 79 ++++++ ...sonx_audio_transcription_transformation.py | 58 +++++ .../test_watsonx_embedding_transformation.py | 90 +++++++ .../watsonx/rerank/test_watsonx_rerank.py | 62 +++++ .../mcp_server/test_db_credentials.py | 43 ++++ .../mcp_server/test_mcp_env_vars.py | 28 +++ .../proxy/client/cli/test_auth_commands.py | 31 +++ tests/unit/proxy/client/test_chat.py | 51 ++++ .../db/test_object_permission_repository.py | 73 ++++++ .../fine_tuning_endpoints/test_endpoints.py | 68 ++++++ .../policy_endpoints/test_endpoints.py | 75 ++++++ .../test_model_management_endpoints.py | 47 ++++ .../test_organization_endpoints.py | 28 +++ .../test_tag_management_endpoints.py | 65 +++++ ...t_anthropic_passthrough_logging_handler.py | 19 ++ ...test_vertex_passthrough_logging_handler.py | 95 ++++++++ .../policy_engine/test_policy_registry.py | 154 ++++++++++++ .../spend_tracking/test_budget_reservation.py | 84 +++++++ .../unit/rag/ingestion/test_base_ingestion.py | 101 ++++++++ .../rag/ingestion/test_gemini_ingestion.py | 111 +++++++++ .../rag/ingestion/test_openai_ingestion.py | 101 ++++++++ tests/unit/rag/test_main.py | 87 +++++++ tests/unit/repositories/test_repositories.py | 75 ++++++ .../adaptive_router/test_adaptive_router.py | 9 + .../test_aws_secret_manager.py | 36 +++ .../test_secret_managers_main.py | 35 +++ tests/unit/test_dashscope_image_generation.py | 96 ++++++++ .../test_vector_store_registry.py | 96 +++++++- 240 files changed, 6820 insertions(+), 556 deletions(-) create mode 100644 tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py create mode 100644 tests/unit/integrations/opik/opik_payload_builder/__init__.py create mode 100644 tests/unit/integrations/opik/opik_payload_builder/test_api.py create mode 100644 tests/unit/llms/aiohttp_openai/__init__.py create mode 100644 tests/unit/llms/aiohttp_openai/chat/__init__.py create mode 100644 tests/unit/llms/aiohttp_openai/chat/test_transformation.py create mode 100644 tests/unit/llms/bedrock/image_edit/test_stability_transformation.py create mode 100644 tests/unit/llms/bedrock/image_generation/__init__.py create mode 100644 tests/unit/llms/bedrock/image_generation/test_amazon_titan_transformation.py create mode 100644 tests/unit/llms/clarifai/__init__.py create mode 100644 tests/unit/llms/clarifai/chat/__init__.py create mode 100644 tests/unit/llms/clarifai/chat/test_transformation.py create mode 100644 tests/unit/llms/cohere/rerank/test_transformation.py create mode 100644 tests/unit/llms/cometapi/image_generation/__init__.py create mode 100644 tests/unit/llms/cometapi/image_generation/test_transformation.py create mode 100644 tests/unit/llms/e2b/__init__.py create mode 100644 tests/unit/llms/e2b/sandbox/__init__.py create mode 100644 tests/unit/llms/e2b/sandbox/test_transformation.py create mode 100644 tests/unit/llms/elevenlabs/audio_transcription/__init__.py create mode 100644 tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py create mode 100644 tests/unit/llms/fal_ai/image_generation/test_flux_pro_v11_ultra_transformation.py create mode 100644 tests/unit/llms/fal_ai/image_generation/test_ideogram_v3_transformation.py create mode 100644 tests/unit/llms/fal_ai/image_generation/test_recraft_v3_transformation.py create mode 100644 tests/unit/llms/fal_ai/image_generation/test_stable_diffusion_transformation.py create mode 100644 tests/unit/llms/jina_ai/rerank/__init__.py create mode 100644 tests/unit/llms/jina_ai/rerank/test_transformation.py create mode 100644 tests/unit/llms/manus/files/__init__.py create mode 100644 tests/unit/llms/manus/files/test_transformation.py create mode 100644 tests/unit/llms/minimax/text_to_speech/__init__.py create mode 100644 tests/unit/llms/minimax/text_to_speech/test_transformation.py create mode 100644 tests/unit/llms/openai/responses/count_tokens/__init__.py create mode 100644 tests/unit/llms/openai/responses/count_tokens/test_handler.py create mode 100644 tests/unit/llms/replicate/__init__.py create mode 100644 tests/unit/llms/replicate/chat/__init__.py create mode 100644 tests/unit/llms/replicate/chat/test_transformation.py create mode 100644 tests/unit/llms/together_ai/rerank/__init__.py create mode 100644 tests/unit/llms/together_ai/rerank/test_handler.py create mode 100644 tests/unit/llms/vertex_ai/image_generation/test_image_generation_handler.py create mode 100644 tests/unit/proxy/db/test_object_permission_repository.py create mode 100644 tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py create mode 100644 tests/unit/proxy/policy_engine/test_policy_registry.py create mode 100644 tests/unit/rag/ingestion/test_base_ingestion.py create mode 100644 tests/unit/rag/ingestion/test_gemini_ingestion.py create mode 100644 tests/unit/rag/ingestion/test_openai_ingestion.py create mode 100644 tests/unit/secret_managers/test_aws_secret_manager.py diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index a29f0a1b43a..c648f7faea0 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -46,6 +46,8 @@ from litellm.proxy._types import ( UserAPIKeyAuth, WebhookEvent, ) +from litellm.repositories.table_repositories import InvitationLinkRepository +from litellm.repositories.user_repository import UserRepository from litellm.secret_managers.main import get_secret_bool from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL @@ -844,7 +846,7 @@ class BaseEmailLogger(CustomLogger): ) return None - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) @@ -929,7 +931,7 @@ class BaseEmailLogger(CustomLogger): try: # Try to get existing invitation existing_invitations = ( - await prisma_client.db.litellm_invitationlink.find_many( + await InvitationLinkRepository(prisma_client).table.find_many( where={"user_id": user_id}, order={"created_at": "desc"}, ) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 35fc3510414..985a9b6de38 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -15,6 +15,10 @@ from litellm.constants import ( MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, ) +from litellm.repositories.table_repositories import ManagedObjectRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository if TYPE_CHECKING: from prisma import models as prisma_models @@ -58,25 +62,19 @@ class _ManagedObjectRow(Protocol): def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]": - table: Final[TableActions[_ManagedObjectRow]] = prisma_client.db.litellm_managedobjecttable - return table + return ManagedObjectRepository(prisma_client).table def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]": - table: Final[TableActions[prisma_models.LiteLLM_UserTable]] = prisma_client.db.litellm_usertable - return table + return UserRepository(prisma_client).table def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]": - table: Final[TableActions[prisma_models.LiteLLM_VerificationToken]] = ( - prisma_client.db.litellm_verificationtoken - ) - return table + return VerificationTokenRepository(prisma_client).table def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]": - table: Final[TableActions[prisma_models.LiteLLM_TeamTable]] = prisma_client.db.litellm_teamtable - return table + return TeamRepository(prisma_client).table class CheckBatchCost: diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index cdeea0d3d4b..738b9ebe1e0 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -16,6 +16,7 @@ from litellm.constants import ( MAX_OBJECTS_PER_POLL_CYCLE, STALE_OBJECT_CLEANUP_BATCH_SIZE, ) +from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN @@ -43,8 +44,7 @@ class _ManagedObjectRow(Protocol): def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]": - table: Final[TableActions[_ManagedObjectRow]] = prisma_client.db.litellm_managedobjecttable - return table + return ManagedObjectRepository(prisma_client).table class CheckResponsesCost: diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index e0a94612646..011bf95defd 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1647,7 +1647,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): owner_filter: Final = build_owner_filter(user_api_key_dict) if owner_filter is None: - return FileListPage(**build_list_page([])) + return FileListPage.model_validate(build_list_page([])) if after: cursor_row = await _managed_file_table(self.prisma_client).find_first( @@ -1686,7 +1686,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): cursor_id = chunk[-1].unified_file_id chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE) - return FileListPage(**build_list_page(matches[:page_size], has_more=len(matches) > page_size)) + return FileListPage.model_validate(build_list_page(matches[:page_size], has_more=len(matches) > page_size)) def _is_batch_polling_enabled(self) -> bool: """ diff --git a/litellm/a2a_protocol/providers/base.py b/litellm/a2a_protocol/providers/base.py index 5a5eff8cf35..6196cc789fc 100644 --- a/litellm/a2a_protocol/providers/base.py +++ b/litellm/a2a_protocol/providers/base.py @@ -4,7 +4,6 @@ Base configuration for A2A protocol providers. from abc import ABC, abstractmethod from collections.abc import AsyncIterator -from typing import Any class BaseA2AProviderConfig(ABC): @@ -19,10 +18,10 @@ class BaseA2AProviderConfig(ABC): async def handle_non_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, **kwargs, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Handle non-streaming A2A request. @@ -40,10 +39,10 @@ class BaseA2AProviderConfig(ABC): async def handle_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, **kwargs, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """ Handle streaming A2A request. diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py index 2b37c0c4906..22dc3aa508a 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py @@ -3,7 +3,7 @@ Bedrock AgentCore A2A provider configuration. """ from collections.abc import AsyncIterator -from typing import Any, Final +from typing import Final from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.bedrock_agentcore.handler import ( @@ -23,10 +23,10 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): async def handle_non_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, **kwargs, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Handle non-streaming request to AgentCore A2A agent.""" litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: @@ -43,10 +43,10 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): async def handle_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, **kwargs, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """Handle streaming request to AgentCore A2A agent.""" litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index da5eb522187..aecc887fbfe 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -7,7 +7,7 @@ completion bridge that would otherwise strip the envelope. import json from collections.abc import AsyncIterator, Mapping -from typing import Any, Final +from typing import Final from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( @@ -30,9 +30,9 @@ class BedrockAgentCoreA2AHandler: async def handle_non_streaming( request_id: str, params: Mapping[str, object], - litellm_params: dict[str, Any], + litellm_params: dict[str, object], agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Handle non-streaming A2A request to AgentCore. @@ -77,9 +77,9 @@ class BedrockAgentCoreA2AHandler: async def handle_streaming( request_id: str, params: Mapping[str, object], - litellm_params: dict[str, Any], + litellm_params: dict[str, object], agent_extra_headers: dict[str, str] | None = None, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """ Handle streaming A2A request to AgentCore. diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index c486f1f6d95..5d4e0d5fb39 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -219,7 +219,7 @@ class BedrockAgentCoreA2ATransformation: return url, signed_headers, signed_body @staticmethod - async def parse_sse_events(response: _SSELineSource) -> AsyncIterator[dict[str, Any]]: + async def parse_sse_events(response: _SSELineSource) -> AsyncIterator[dict[str, object]]: """ Parse SSE events from an httpx streaming response. diff --git a/litellm/a2a_protocol/providers/langflow/config.py b/litellm/a2a_protocol/providers/langflow/config.py index 54d403f88c0..b8819dbe4bc 100644 --- a/litellm/a2a_protocol/providers/langflow/config.py +++ b/litellm/a2a_protocol/providers/langflow/config.py @@ -1,5 +1,4 @@ from collections.abc import AsyncIterator -from typing import Any from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2A_USER_API_KEY_HASH_PARAM, @@ -16,10 +15,10 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): async def handle_non_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, **kwargs, - ) -> dict[str, Any]: + ) -> dict[str, object]: litellm_params = kwargs.get("litellm_params") if not litellm_params: raise ValueError( @@ -39,10 +38,10 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): async def handle_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, **kwargs, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: litellm_params = kwargs.get("litellm_params") if not litellm_params: raise ValueError( diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py index 20404e3702b..f4c758af2d0 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py @@ -2,8 +2,7 @@ Pydantic AI provider configuration. """ -from collections.abc import AsyncIterator -from typing import Any +from collections.abc import AsyncIterator, Mapping from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAIHandler @@ -20,9 +19,12 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): async def handle_non_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, - **kwargs: Any, + *, + timeout: float = 60.0, + agent_extra_headers: Mapping[str, str] | None = None, + **kwargs: object, ) -> dict[str, object]: """Handle non-streaming request to Pydantic AI agent.""" if api_base is None: @@ -31,14 +33,14 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): request_id=request_id, params=params, api_base=api_base, - timeout=kwargs.get("timeout", 60.0), - agent_extra_headers=kwargs.get("agent_extra_headers"), + timeout=timeout, + agent_extra_headers=agent_extra_headers, ) async def handle_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, **kwargs, ) -> AsyncIterator[dict[str, object]]: diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py index c083c0267f7..f3ceb899a92 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py @@ -5,8 +5,8 @@ Pydantic AI agents follow A2A protocol but don't support streaming natively. This handler provides fake streaming by converting non-streaming responses into streaming chunks. """ -from collections.abc import AsyncIterator -from typing import Any, Final +from collections.abc import AsyncIterator, Mapping +from typing import Final from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( @@ -26,11 +26,11 @@ class PydanticAIHandler: @staticmethod async def handle_non_streaming( request_id: str, - params: dict[str, Any], + params: Mapping[str, object], api_base: str | None = None, timeout: float = 60.0, - agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + agent_extra_headers: Mapping[str, str] | None = None, + ) -> dict[str, object]: """ Handle non-streaming request to Pydantic AI agent. @@ -63,13 +63,13 @@ class PydanticAIHandler: @staticmethod async def handle_streaming( request_id: str, - params: dict[str, Any], + params: Mapping[str, object], api_base: str | None = None, timeout: float = 60.0, chunk_size: int = 50, delay_ms: int = 10, - agent_extra_headers: dict[str, str] | None = None, - ) -> AsyncIterator[dict[str, Any]]: + agent_extra_headers: Mapping[str, str] | None = None, + ) -> AsyncIterator[dict[str, object]]: """ Handle streaming request to Pydantic AI agent with fake streaming. diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 024e8c179c2..35fe6e76f4d 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -7,7 +7,7 @@ This module provides fake streaming by converting non-streaming responses into s import asyncio from collections.abc import AsyncIterator, Mapping, Sequence -from typing import Any, Final, Protocol, cast, runtime_checkable +from typing import Final, Protocol, runtime_checkable from uuid import uuid4 from pydantic import TypeAdapter @@ -100,7 +100,7 @@ class PydanticAITransformation: request_id: str, max_attempts: int = 30, poll_interval: float = 0.5, - agent_extra_headers: dict[str, str] | None = None, + agent_extra_headers: Mapping[str, str] | None = None, ) -> dict[str, object]: """ Poll for task completion using tasks/get method. @@ -156,7 +156,7 @@ class PydanticAITransformation: request_id: str, params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, - agent_extra_headers: dict[str, str] | None = None, + agent_extra_headers: Mapping[str, str] | None = None, ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -200,7 +200,7 @@ class PydanticAITransformation: # Send request to Pydantic AI agent using shared async HTTP client client: Final = get_async_httpx_client( - llm_provider=cast(Any, "pydantic_ai_agent"), + llm_provider="pydantic_ai_agent", params={"timeout": timeout}, ) response: Final = await client.post( @@ -242,7 +242,7 @@ class PydanticAITransformation: request_id: str, params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, - agent_extra_headers: dict[str, str] | None = None, + agent_extra_headers: Mapping[str, str] | None = None, ) -> dict[str, object]: """ Send a non-streaming A2A request to Pydantic AI agent and wait for completion. @@ -278,7 +278,7 @@ class PydanticAITransformation: request_id: str, params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, - agent_extra_headers: dict[str, str] | None = None, + agent_extra_headers: Mapping[str, str] | None = None, ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py index 44873edf271..a1a6c3a5957 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -3,11 +3,11 @@ A2A provider configuration for IBM watsonx Orchestrate (WXO). """ from collections.abc import AsyncIterator -from typing import Any, Final from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import ( WatsonxOrchestrateHandler, + WXOLitellmParams, ) @@ -17,12 +17,13 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): async def handle_non_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, - **kwargs: Any, + *, + litellm_params: WXOLitellmParams | None = None, + **kwargs: object, ) -> dict[str, object]: """Handle a non-streaming A2A request via WXO runs API.""" - litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: raise ValueError( "litellm_params is required for WatsonxOrchestrateA2AConfig " @@ -37,12 +38,13 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): async def handle_streaming( self, request_id: str, - params: dict[str, Any], + params: dict[str, object], api_base: str | None = None, - **kwargs: Any, + *, + litellm_params: WXOLitellmParams | None = None, + **kwargs: object, ) -> AsyncIterator[dict[str, object]]: """Handle a streaming A2A request via WXO streaming runs API.""" - litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: raise ValueError( "litellm_params is required for WatsonxOrchestrateA2AConfig " diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index c66b07c321c..0a2871d25cc 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -7,7 +7,7 @@ import hashlib import json import time from collections.abc import AsyncIterator -from typing import Any, Final, NamedTuple, Protocol +from typing import Final, NamedTuple, Protocol import httpx from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -228,7 +228,7 @@ class WatsonxOrchestrateHandler: return run_data @staticmethod - async def _accumulate_wxo_sse_text(response: Any) -> str: + async def _accumulate_wxo_sse_text(response: _SSELineSource) -> str: source: Final[_WXOView] = {"sse_source": response} accumulated_text = "" async for line in source["sse_source"].aiter_lines(): diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py index cec25634bb8..d82e58ecc91 100644 --- a/litellm/caching/_embedding_router.py +++ b/litellm/caching/_embedding_router.py @@ -12,7 +12,7 @@ This module is dependency-injected: callers pass the proxy ``llm_router`` and from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import litellm @@ -39,10 +39,10 @@ def resolve_embedding_router( def build_router_embedding_metadata( - request_metadata: dict[str, Any] | None, -) -> dict[str, Any]: + request_metadata: Mapping[str, object] | None, +) -> Mapping[str, object]: """Forward the caller's full metadata, flagged as a semantic-cache embedding.""" - metadata: Final[dict[str, Any]] = dict(request_metadata or {}) + metadata: Final = dict(request_metadata or {}) metadata["semantic-cache-embedding"] = True return metadata diff --git a/litellm/files/main.py b/litellm/files/main.py index 723784795b0..9057211256d 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -331,7 +331,7 @@ async def afile_retrieve( else: response = init_response - return OpenAIFileObject(**response.model_dump()) + return OpenAIFileObject.model_validate(response.model_dump()) except Exception as e: raise e diff --git a/litellm/files/streaming.py b/litellm/files/streaming.py index 5d23ebf32ae..460050872f4 100644 --- a/litellm/files/streaming.py +++ b/litellm/files/streaming.py @@ -34,7 +34,7 @@ class FileContentStreamingResponse: self.custom_llm_provider = custom_llm_provider self.logging_obj = logging_obj self.standard_logging_object: StandardLoggingPayload | None = None - self._hidden_params: dict[str, Any] = {} + self._hidden_params: dict[str, object] = {} self._logging_completed = False self._close_completed = False self._start_time = ( @@ -121,7 +121,7 @@ class FileContentStreamingResponse: return response def _sync_hidden_params(self) -> None: - litellm_params: dict[str, Any] = {} + litellm_params: dict[str, object] = {} if self.logging_obj is not None: litellm_params = self.logging_obj.model_call_details.get("litellm_params", {}) or {} diff --git a/litellm/harness/context.py b/litellm/harness/context.py index eaae9aafe1f..f75ee7e49f1 100644 --- a/litellm/harness/context.py +++ b/litellm/harness/context.py @@ -4,7 +4,7 @@ from __future__ import annotations from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, TypeAlias +from typing import TYPE_CHECKING, TypeAlias from pydantic import BaseModel @@ -42,7 +42,7 @@ class SessionContext: api_base: str | None = None endpoint: ModelEndpoint | None = None instructions: str | None = None - tools: Sequence[Callable[..., Any]] = () + tools: Sequence[Callable[..., object]] = () skills: Sequence[str] = () disable_tools: Sequence[str] = () permissions: PermissionMode = "full" @@ -50,7 +50,7 @@ class SessionContext: output: type[BaseModel] | None = None max_turns: int | None = None timeout: float | None = None - metadata: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, object] = field(default_factory=dict) options: HarnessOptions | None = None # Set by the handler after each turn. final_text: str = "" diff --git a/litellm/harness/handlers/cli_handler.py b/litellm/harness/handlers/cli_handler.py index 0fda16c34fe..d217d1e5bbb 100644 --- a/litellm/harness/handlers/cli_handler.py +++ b/litellm/harness/handlers/cli_handler.py @@ -13,7 +13,7 @@ import asyncio import os from collections import deque from collections.abc import AsyncIterator, Sequence -from typing import Any, Final +from typing import Final from litellm._logging import verbose_logger from litellm.constants import HARNESS_STDERR_TAIL_LINES, HARNESS_STREAM_READ_CHUNK_BYTES @@ -125,7 +125,7 @@ class CLIHarnessHandler(BaseHarnessHandler): self._proc = proc tail: Final[deque[str]] = deque(maxlen=HARNESS_STDERR_TAIL_LINES) # mutable-ok: bounded stderr ring buffer stderr_task = asyncio.ensure_future(drain_stderr(proc.stderr, tail)) - state: Any = self.config.create_stream_state() + state: Final[object] = self.config.create_stream_state() exit_code: int | None = None try: await send_stdin(proc, request.stdin) diff --git a/litellm/harness/options.py b/litellm/harness/options.py index 18865359ba8..5d7c8ae0807 100644 --- a/litellm/harness/options.py +++ b/litellm/harness/options.py @@ -9,7 +9,7 @@ from typing import Any, Literal @dataclass(frozen=True) class ClaudeCodeOptions: - config: Mapping[str, Any] = field(default_factory=dict) + config: Mapping[str, object] = field(default_factory=dict) env: Mapping[str, str] = field(default_factory=dict) @@ -17,20 +17,20 @@ class ClaudeCodeOptions: class CodexOptions: reasoning_effort: Literal["low", "medium", "high", "xhigh"] | None = None web_search: bool = False - config: Mapping[str, Any] = field(default_factory=dict) + config: Mapping[str, object] = field(default_factory=dict) env: Mapping[str, str] = field(default_factory=dict) @dataclass(frozen=True) class OpenCodeOptions: agent: str = "build" - config: Mapping[str, Any] = field(default_factory=dict) + config: Mapping[str, object] = field(default_factory=dict) env: Mapping[str, str] = field(default_factory=dict) @dataclass(frozen=True) class DeepAgentsOptions: - subagents: Sequence[Any] = () + subagents: Sequence[object] = () recursion_limit: int | None = None diff --git a/litellm/harness/runtime.py b/litellm/harness/runtime.py index da90b28a751..dc24a1f0ef7 100644 --- a/litellm/harness/runtime.py +++ b/litellm/harness/runtime.py @@ -1005,24 +1005,24 @@ def aagent( `await litellm.aagent(...)` returns a Result. With stream=True it returns an async iterator of events instead: `async for event in litellm.aagent(..., stream=True)`. """ - kwargs: dict[str, Any] = { # mutable-ok: forwarded as **kwargs to arun_agent/astream_agent - "sandbox": sandbox, - "model": model, - "api_key": api_key, - "api_base": api_base, - "instructions": instructions, - "tools": tools, - "skills": skills, - "disable_tools": disable_tools, - "permissions": permissions, - "on_approval": on_approval, - "output": output, - "max_turns": max_turns, - "timeout": timeout, - "metadata": metadata, - "options": options, - "install": install, - } - if stream: - return astream_agent(harness, prompt, **kwargs) - return arun_agent(harness, prompt, **kwargs) + call: Final = astream_agent if stream else arun_agent + return call( + harness, + prompt, + sandbox=sandbox, + model=model, + api_key=api_key, + api_base=api_base, + instructions=instructions, + tools=tools, + skills=skills, + disable_tools=disable_tools, + permissions=permissions, + on_approval=on_approval, + output=output, + max_turns=max_turns, + timeout=timeout, + metadata=metadata, + options=options, + install=install, + ) diff --git a/litellm/harness/sync.py b/litellm/harness/sync.py index 9787800833b..18336a3d06e 100644 --- a/litellm/harness/sync.py +++ b/litellm/harness/sync.py @@ -8,7 +8,7 @@ import threading from collections.abc import AsyncIterator, Callable, Coroutine, Mapping, Sequence from concurrent.futures import Future from typing import ( - Any, + Final, TypeVar, ) @@ -417,24 +417,24 @@ def agent( Prefix the model with `litellm_proxy/` to route every model call through your LiteLLM AI Gateway. """ - kwargs: dict[str, Any] = { # mutable-ok: forwarded as **kwargs to _run/_stream - "sandbox": sandbox, - "model": model, - "api_key": api_key, - "api_base": api_base, - "instructions": instructions, - "tools": tools, - "skills": skills, - "disable_tools": disable_tools, - "permissions": permissions, - "on_approval": on_approval, - "output": output, - "max_turns": max_turns, - "timeout": timeout, - "metadata": metadata, - "options": options, - "install": install, - } - if stream: - return _stream(harness, prompt, **kwargs) - return _run(harness, prompt, **kwargs) + call: Final = _stream if stream else _run + return call( + harness, + prompt, + sandbox=sandbox, + model=model, + api_key=api_key, + api_base=api_base, + instructions=instructions, + tools=tools, + skills=skills, + disable_tools=disable_tools, + permissions=permissions, + on_approval=on_approval, + output=output, + max_turns=max_turns, + timeout=timeout, + metadata=metadata, + options=options, + install=install, + ) diff --git a/litellm/harness/types.py b/litellm/harness/types.py index 475568a9979..6b0118f4349 100644 --- a/litellm/harness/types.py +++ b/litellm/harness/types.py @@ -7,7 +7,7 @@ import json from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum -from typing import Any, Literal +from typing import Literal from pydantic import BaseModel @@ -71,7 +71,7 @@ class ToolCall: id: str name: str native_name: str - input: Mapping[str, Any] + input: Mapping[str, object] builtin: bool = True @@ -100,7 +100,7 @@ class Approval: """A request to run a tool. The turn waits until allow() or deny() is called.""" tool: str - input: Mapping[str, Any] + input: Mapping[str, object] _decision: asyncio.Future[tuple[bool, str]] = field( default_factory=lambda: asyncio.get_event_loop().create_future(), compare=False, diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 0271cf1e03c..3f89ac6978f 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -2,6 +2,7 @@ import json from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict, override from litellm._logging import verbose_logger @@ -29,6 +30,8 @@ from litellm.integrations._types.open_inference import ( ToolCallAttributes, ) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) + class ArizeOTELAttributes(BaseLLMObsOTELAttributes): @staticmethod @@ -609,9 +612,7 @@ def _coerce_response_obj_for_attrs(response_obj): text: Final = getattr(response_obj, "text", None) if isinstance(text, str) and text: try: - parsed: Final = json.loads(text) - if isinstance(parsed, dict): - return parsed + return _JSON_OBJECT.validate_python(json.loads(text)) except Exception: pass return response_obj @@ -1062,9 +1063,7 @@ def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs): inner = candidate.get("response") if isinstance(inner, str): try: - parsed = json.loads(inner) - if isinstance(parsed, dict): - return parsed + return _JSON_OBJECT.validate_python(json.loads(inner)) except Exception: continue if isinstance(inner, dict): @@ -1078,9 +1077,7 @@ def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs): return original if isinstance(original, str): try: - parsed = json.loads(original) - if isinstance(parsed, dict): - return parsed + return _JSON_OBJECT.validate_python(json.loads(original)) except Exception: return None return None diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 22bb50cd739..9565a64e69c 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -15,6 +15,7 @@ from types import MappingProxyType from typing import Any, Final, Literal import httpx +from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger @@ -57,6 +58,7 @@ from litellm.types.utils import ( _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) _EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""} _MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024 +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset( {"agent", "assistant", "developer", "function", "model", "system", "tool", "user"} ) @@ -282,8 +284,10 @@ def _to_dd_arguments(raw_arguments: object) -> dict[str, object] | str: return raw_arguments if isinstance(raw_arguments, dict) else str(raw_arguments) if len(raw_arguments) > _MAX_PARSED_TOOL_ARGUMENT_CHARS: return raw_arguments - parsed: Final = safe_json_loads(raw_arguments) - return parsed if isinstance(parsed, dict) else raw_arguments + try: + return _JSON_OBJECT.validate_python(safe_json_loads(raw_arguments)) + except ValidationError: + return raw_arguments def _to_dd_tool_calls(message: Mapping[str, object]) -> tuple[ToolCall, ...]: diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 21906b0d996..9bc324256fb 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -4,12 +4,12 @@ import json import os import re import uuid -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timezone, tzinfo from typing import Any, Final, Protocol, cast import httpx -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -35,6 +35,11 @@ GALILEO_CLOUD_API_BASE_URL: Final = "https://api.galileo.ai" # unavailable, invalid credentials) cannot leak memory unboundedly. GALILEO_MAX_IN_MEMORY_RECORDS: Final = 1000 +_UNTYPED_VALUE: Final = TypeAdapter(object) +_ZERO_ARGUMENT_CALLABLE: Final[TypeAdapter[Callable[[], object]]] = TypeAdapter( + Callable[[], object], config=ConfigDict(hide_input_in_errors=True) +) + class _GalileoLoginBody(TypedDict): """Decoded body of the Galileo login response.""" @@ -473,9 +478,9 @@ class GalileoObserve(CustomLogger): if isinstance(value, str): return value - def _json_default(obj: Any) -> object: + def _json_default(obj: object) -> object: if hasattr(obj, "model_dump"): - return obj.model_dump() + return _ZERO_ARGUMENT_CALLABLE.validate_python(getattr(obj, "model_dump", None))() return str(obj) return json.dumps(value, default=_json_default) @@ -496,7 +501,7 @@ class GalileoObserve(CustomLogger): if hasattr(message, "json"): message_json: Final[object] = message.json() if isinstance(message_json, str): - return json.loads(message_json) + return _UNTYPED_VALUE.validate_python(json.loads(message_json)) return message_json return message return None diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 8d588896b2f..956e754944a 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -8,6 +8,8 @@ from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from pydantic import ConfigDict, TypeAdapter + import litellm from litellm._logging import verbose_logger from litellm.integrations._types.open_inference import ( @@ -95,6 +97,8 @@ class _ResponseWithUsageView(TypedDict, total=False): usage: "_UsageCompletionTokensView | None" +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + # Cap on credential-scoped providers held at once; each one owns an exporter thread. _MAX_DYNAMIC_TRACER_PROVIDERS: Final = 256 @@ -1633,7 +1637,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): attributes = self.config.attributes if attributes is None and self.callback_name in (None, "otel"): otel_settings: Final = (litellm.callback_settings or {}).get("otel") or {} - raw: Final = otel_settings.get("attributes") if isinstance(otel_settings, dict) else None + raw: Final[object] = otel_settings.get("attributes") if isinstance(otel_settings, dict) else None if raw is not None: attributes = _build_metric_attribute_filter(raw) ( @@ -2919,7 +2923,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): import json try: - _parsed: Final[Mapping[str, object]] = json.loads(_raw_response) + _parsed: Final = _JSON_OBJECT.validate_python(json.loads(_raw_response)) for param, val in _parsed.items(): self.safe_set_attribute( span=span, diff --git a/litellm/integrations/opik/opik_payload_builder/api.py b/litellm/integrations/opik/opik_payload_builder/api.py index 4334c81eb6e..d16bc71c416 100644 --- a/litellm/integrations/opik/opik_payload_builder/api.py +++ b/litellm/integrations/opik/opik_payload_builder/api.py @@ -1,12 +1,17 @@ """Public API for Opik payload building.""" +from collections.abc import Mapping from datetime import datetime from typing import Any, Final +from pydantic import ConfigDict, TypeAdapter + from litellm.integrations.opik import utils from . import extractors, payload_builders, types +_STANDARD_LOGGING_FIELDS: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + def build_opik_payload( kwargs: dict[str, Any], @@ -36,12 +41,14 @@ def build_opik_payload( - First element is TracePayload if creating a new trace, None if attaching to existing - Second element is always SpanPayload """ - standard_logging_object: Final = kwargs["standard_logging_object"] + standard_logging_object: Final = _STANDARD_LOGGING_FIELDS.validate_python(kwargs["standard_logging_object"]) # Extract litellm params and metadata litellm_params: Final = kwargs.get("litellm_params", {}) or {} litellm_metadata: Final = litellm_params.get("metadata", {}) or {} - standard_logging_metadata: Final = standard_logging_object.get("metadata", {}) or {} + standard_logging_metadata: Final = _STANDARD_LOGGING_FIELDS.validate_python( + standard_logging_object.get("metadata", {}) or {} + ) # Extract and merge Opik metadata opik_metadata: Final = extractors.extract_opik_metadata(litellm_metadata, standard_logging_metadata) diff --git a/litellm/integrations/opik/opik_payload_builder/extractors.py b/litellm/integrations/opik/opik_payload_builder/extractors.py index 4dd3d40fae3..1ceb9aa763f 100644 --- a/litellm/integrations/opik/opik_payload_builder/extractors.py +++ b/litellm/integrations/opik/opik_payload_builder/extractors.py @@ -4,8 +4,12 @@ import json from collections.abc import Mapping from typing import Any, Final +from pydantic import TypeAdapter + from litellm import _logging +_DECODED_JSON: Final = TypeAdapter(object) + def normalize_provider_name(provider: str | None) -> str | None: """ @@ -149,7 +153,7 @@ def apply_proxy_header_overrides( thread_id = value elif param_key == "tags": try: - parsed_tags: object = json.loads(value) + parsed_tags = _DECODED_JSON.validate_python(json.loads(value)) if isinstance(parsed_tags, list): tags.extend(parsed_tags) except (json.JSONDecodeError, TypeError): diff --git a/litellm/integrations/opik/opik_payload_builder/types.py b/litellm/integrations/opik/opik_payload_builder/types.py index 546ce55f840..2f418b2ef9e 100644 --- a/litellm/integrations/opik/opik_payload_builder/types.py +++ b/litellm/integrations/opik/opik_payload_builder/types.py @@ -1,7 +1,8 @@ """Type definitions for Opik payload building.""" +from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Final, Literal +from typing import Final, Literal @dataclass @@ -13,9 +14,9 @@ class TracePayload: name: str start_time: str end_time: str - input: Any - output: Any - metadata: dict[str, Any] + input: object + output: object + metadata: Mapping[str, object] tags: list[str] thread_id: str | None = None @@ -32,9 +33,9 @@ class SpanPayload: model: str start_time: str end_time: str - input: Any - output: Any - metadata: dict[str, Any] + input: object + output: object + metadata: Mapping[str, object] tags: list[str] usage: dict[str, int] parent_span_id: str | None = None diff --git a/litellm/integrations/opik/utils.py b/litellm/integrations/opik/utils.py index 2caacd1f871..79b6199423e 100644 --- a/litellm/integrations/opik/utils.py +++ b/litellm/integrations/opik/utils.py @@ -2,7 +2,8 @@ import configparser import os import time import uuid -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final CONFIG_FILE_PATH_DEFAULT: Final[str] = "~/.opik.config" @@ -93,14 +94,14 @@ def create_usage_object(usage): return usage_dict -def _remove_nulls(x: dict[str, Any]) -> dict[str, Any]: +def _remove_nulls(x: Mapping[str, object]) -> dict[str, object]: """Remove None values from dict.""" return {k: v for k, v in x.items() if v is not None} def get_traces_and_spans_from_payload( - payload: list[dict[str, Any]], -) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + payload: Sequence[Mapping[str, object]], +) -> tuple[list[dict[str, object]], list[dict[str, object]]]: """ Separate traces and spans from payload. diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index ddb8e127408..6e9038c0cad 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -2,7 +2,7 @@ from enum import Enum from functools import lru_cache -from typing import Annotated, Any, Final +from typing import Annotated, Final from pydantic import AliasChoices, BaseModel, Field, TypeAdapter, ValidationError, field_validator, model_validator from pydantic.fields import FieldInfo @@ -315,7 +315,7 @@ class OpenTelemetryV2Config(BaseSettings): mode="before", ) @classmethod - def _split_csv(cls, value: Any) -> Any: + def _split_csv(cls, value: object) -> object: """Accept a comma-separated string for list fields. Env vars are strings, but these fields are lists. Pydantic-settings would diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index e1623f4697f..76b23467679 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -307,7 +307,7 @@ class GenAIMetricRecorder: return common_attrs - def _bounded_attributes(self, kwargs: Mapping[str, Any]) -> MetricAttributes: + def _bounded_attributes(self, kwargs: Mapping[str, object]) -> MetricAttributes: """The datapoint attributes, capped at :data:`METRIC_ATTRIBUTE_CEILING`. The cap runs BEFORE the operator's include/exclude filter so the filter can @@ -322,7 +322,7 @@ class GenAIMetricRecorder: attributes = None if self._callback_name in (None, "otel"): otel_settings: Final = (litellm.callback_settings or {}).get("otel") or {} - raw: Final = otel_settings.get("attributes") if isinstance(otel_settings, dict) else None + raw: Final[object] = otel_settings.get("attributes") if isinstance(otel_settings, dict) else None if raw is not None: attributes = _build_metric_attribute_filter(raw) # A bad filter (include_list + exclude_list both set, an unfilterable name) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index f504292cb64..08249d33561 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -595,7 +595,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): and not (self.s3_drop_on_terminal_error and _is_terminal(response)) and attempt < max_retries - 1 ): - wait_time = 2**attempt # 1s, 2s + wait_time = 1 << attempt # 1s, 2s verbose_logger.log( logging.DEBUG if _in_flush.get() else logging.WARNING, "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s", @@ -897,7 +897,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): and not (self.s3_drop_on_terminal_error and _is_terminal(response)) and attempt < max_retries - 1 ): - wait_time = 2**attempt # 1s, 2s + wait_time = 1 << attempt # 1s, 2s verbose_logger.warning( "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s", response.status_code, diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 80530622e36..f40a166e67d 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -6,12 +6,12 @@ It searches the vector store for relevant context, runs the request's pre-call g over that context, and appends it to the messages. """ -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from itertools import chain from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_args -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError from typing_extensions import assert_never import litellm @@ -44,10 +44,12 @@ else: LiteLLMLoggingObj = Any SEARCH_FAILURES_FIELD: Final = "vector_store_search_failures" +_PROVIDER_FIELDS_ATTRIBUTE: Final = "provider_specific_fields" _DEFAULT_FAILURE_MODE: Final[VectorStoreSearchFailureMode] = "annotate" _FAILURE_MODE_ADAPTER: Final = TypeAdapter(VectorStoreSearchFailureMode) _OBJECT_ADAPTER: Final = TypeAdapter(object) _STR_KEYED_ADAPTER: Final = TypeAdapter(dict[str, object]) +_ITERABLE_ADAPTER: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) _GUARDRAIL_KEYS_THE_PROXY_MERGES_INTO_METADATA: Final = frozenset( {"guardrails", "guardrail_config", "policies", "include_guardrail_response"} ) @@ -475,7 +477,7 @@ class VectorStorePreCallHook(CustomLogger): async def async_post_call_streaming_deployment_hook( self, request_data: dict, - response_chunk: Any, + response_chunk: object, call_type: CallTypes | None, ) -> object | None: """ @@ -496,15 +498,17 @@ class VectorStorePreCallHook(CustomLogger): return response_chunk # Add search results to streaming chunk - if hasattr(response_chunk, "choices") and response_chunk.choices: - for choice in response_chunk.choices: - if hasattr(choice, "delta") and choice.delta: - provider_fields = getattr(choice.delta, "provider_specific_fields", None) or {} + choices: Final[object] = getattr(response_chunk, "choices", None) + if choices: + for choice in _ITERABLE_ADAPTER.validate_python(choices): + delta: object = getattr(choice, "delta", None) + if delta: + provider_fields = getattr(delta, _PROVIDER_FIELDS_ATTRIBUTE, None) or {} if search_results: provider_fields["search_results"] = search_results if search_failures: provider_fields[SEARCH_FAILURES_FIELD] = search_failures - choice.delta.provider_specific_fields = provider_fields + setattr(delta, _PROVIDER_FIELDS_ATTRIBUTE, provider_fields) # Return modified chunk return response_chunk diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index af4c6f69944..6171b341ca4 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -3,8 +3,8 @@ A2A Protocol Transformation for LiteLLM """ import uuid -from collections.abc import Iterator, Mapping -from typing import TYPE_CHECKING, Any, Final +from collections.abc import AsyncIterator, Iterator, Mapping +from typing import TYPE_CHECKING, Final import httpx @@ -72,9 +72,9 @@ class A2AConfig(BaseConfig): agent_name: str, api_base: str | None, api_key: str | None, - headers: dict[str, Any] | None, - optional_params: dict[str, Any], - ) -> tuple[str | None, str | None, dict[str, Any] | None]: + headers: dict[str, object] | None, + optional_params: dict[str, object], + ) -> tuple[str | None, str | None, dict[str, object] | None]: """ Resolve agent configuration from the registry for a registered agent. @@ -376,7 +376,7 @@ class A2AConfig(BaseConfig): def get_model_response_iterator( self, - streaming_response: Iterator | Any, + streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, ) -> BaseModelResponseIterator: @@ -397,7 +397,7 @@ class A2AConfig(BaseConfig): json_mode=json_mode, ) - def _openai_message_to_a2a_message(self, message: dict[str, Any]) -> dict[str, Any]: + def _openai_message_to_a2a_message(self, message: Mapping[str, object]) -> dict[str, object]: """ Convert OpenAI message to A2A message format. diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index a06c670e3f1..dba7a7405ab 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -7,9 +7,11 @@ https://github.com/BerriAI/litellm/issues/6592 New config to ensure we introduce this without causing breaking changes for users """ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final from aiohttp import ClientResponse +from pydantic import ConfigDict, TypeAdapter from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.types.llms.openai import AllMessageValues @@ -23,6 +25,10 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECTS: Final = TypeAdapter( + Iterable[Mapping[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True) +) + class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): def get_complete_url( @@ -73,7 +79,9 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): ) -> ModelResponse: _json_response: Final = await raw_response.json() model_response.id = _json_response.get("id") - model_response.choices = [Choices(**choice) for choice in _json_response.get("choices")] + model_response.choices = [ + Choices.model_validate(choice) for choice in _JSON_OBJECTS.validate_python(_json_response.get("choices")) + ] model_response.created = _json_response.get("created") model_response.model = _json_response.get("model") model_response.object = _json_response.get("object") diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 920916726f2..401c09b580b 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -27,7 +27,7 @@ class AnthropicTokenCounter(BaseTokenCounter): self, model_to_use: str, messages: list[dict[str, Any]] | None, - contents: list[dict[str, Any]] | None, + contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", tools: list[dict[str, Any]] | None = None, diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index da21f130f1a..082902e8897 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -19,6 +19,7 @@ from typing import Final, cast import httpx from openai.types.file_deleted import FileDeleted +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -47,6 +48,8 @@ ANTHROPIC_FILES_API_BASE: Final = "https://api.anthropic.com" ANTHROPIC_FILES_BETA_HEADER: Final = "files-api-2025-04-14" ANTHROPIC_MESSAGE_BATCH_ID_PREFIX: Final = "msgbatch_" +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) + class AnthropicFilesConfig(BaseFilesConfig): """ @@ -211,7 +214,7 @@ class AnthropicFilesConfig(BaseFilesConfig): "created_at": "2025-01-01T00:00:00Z" } """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) return self._parse_anthropic_file(response_json) def transform_retrieve_file_request( @@ -230,7 +233,7 @@ class AnthropicFilesConfig(BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> OpenAIFileObject: - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) return self._parse_anthropic_file(response_json) def transform_delete_file_request( @@ -249,13 +252,9 @@ class AnthropicFilesConfig(BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> FileDeleted: - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) file_id: Final = response_json.get("id", "") - return FileDeleted( - id=file_id, - deleted=True, - object="file", - ) + return FileDeleted.model_validate({"id": file_id, "deleted": True, "object": "file"}) def transform_list_files_request( self, diff --git a/litellm/llms/azure/audio_transcription/transformation.py b/litellm/llms/azure/audio_transcription/transformation.py index 44623ec778a..17597da3895 100644 --- a/litellm/llms/azure/audio_transcription/transformation.py +++ b/litellm/llms/azure/audio_transcription/transformation.py @@ -5,10 +5,12 @@ Maps OpenAI-compatible audio transcription calls to Azure Speech REST recognition for short audio. """ -from typing import Any, Final +from collections.abc import Mapping +from typing import Final from urllib.parse import urlencode, urlparse import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.llms.base_llm.audio_transcription.transformation import ( @@ -23,6 +25,8 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import FileTypes, TranscriptionResponse +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class AzureSpeechAudioTranscriptionException(BaseLLMException): pass @@ -127,7 +131,8 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): raw_response: httpx.Response, ) -> TranscriptionResponse: response_json: Final = raw_response.json() - recognition_status: Final = response_json.get("RecognitionStatus") + payload: Final = _JSON_OBJECT.validate_python(response_json) + recognition_status: Final = payload.get("RecognitionStatus") if recognition_status is not None and recognition_status != "Success": raise AzureSpeechAudioTranscriptionException( message=(f"Azure AI Speech transcription failed with RecognitionStatus={recognition_status}."), @@ -135,7 +140,7 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): headers=raw_response.headers, ) - text: Final = self._extract_text(response_json) + text: Final = self._extract_text(payload) response: Final = TranscriptionResponse(text=text) response._hidden_params = response_json return response @@ -194,9 +199,10 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return "detailed" return "simple" - def _extract_text(self, response_json: dict[str, Any]) -> str: - if isinstance(response_json.get("DisplayText"), str): - return response_json["DisplayText"] + def _extract_text(self, response_json: Mapping[str, object]) -> str: + display_text: Final = response_json.get("DisplayText") + if isinstance(display_text, str): + return display_text nbest: Final = response_json.get("NBest") if isinstance(nbest, list) and nbest: diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 36d5a56db0d..6d6e10ce1dc 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -30,7 +30,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): messages: list[dict[str, Any]], api_key: str, api_base: str, - litellm_params: dict[str, Any] | None = None, + litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, tools: list[dict[str, Any]] | None = None, system: object = None, diff --git a/litellm/llms/base_llm/harness/transformation.py b/litellm/llms/base_llm/harness/transformation.py index 643f808d4d8..d6db3d2ecb1 100644 --- a/litellm/llms/base_llm/harness/transformation.py +++ b/litellm/llms/base_llm/harness/transformation.py @@ -18,7 +18,7 @@ from __future__ import annotations from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar +from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar from litellm.harness.errors import HarnessError, OptionsMismatch from litellm.harness.types import Capabilities, Event, Harness @@ -133,7 +133,7 @@ class BaseCLIHarnessConfig(BaseHarnessConfig[OptionsT], Generic[OptionsT, Stream """Fresh per-turn parser state.""" @abstractmethod - def transform_stream_line(self, line: Mapping[str, Any], state: StreamStateT) -> Sequence[Event]: + def transform_stream_line(self, line: Mapping[str, object], state: StreamStateT) -> Sequence[Event]: """One decoded JSON line from stdout to zero or more events. Pure.""" @abstractmethod diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index fd4c3dc1659..e87c5c35fed 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,6 +1,6 @@ from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Final, Literal from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata @@ -17,7 +17,12 @@ if TYPE_CHECKING: # AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses. # Mirrors the mapping used by `BedrockBatchesConfig.transform_create_batch_response` # so create / retrieve return consistent statuses. -_BEDROCK_MIJ_STATUS_TO_OPENAI: Final = { +_BEDROCK_MIJ_STATUS_TO_OPENAI: Final[ + Mapping[ + str, + Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], + ] +] = { "Submitted": "validating", "Validating": "validating", "Scheduled": "validating", @@ -92,7 +97,7 @@ def _record_counts_from_response(response: Mapping[str, object]) -> BatchRequest ) -def _to_epoch(value: Any) -> int | None: +def _to_epoch(value: object) -> int | None: if value is None: return None if isinstance(value, (int, float)): @@ -349,10 +354,7 @@ class BedrockBatchesHandler: ) bedrock_status: Final = str(response.get("status", "")) - openai_status: Final = cast( - Any, - _BEDROCK_MIJ_STATUS_TO_OPENAI.get(bedrock_status, "in_progress"), - ) + openai_status: Final = _BEDROCK_MIJ_STATUS_TO_OPENAI.get(bedrock_status, "in_progress") input_uri: Final = response.get("inputDataConfig", {}).get("s3InputDataConfig", {}).get("s3Uri", "") output_prefix: Final = response.get("outputDataConfig", {}).get("s3OutputDataConfig", {}).get("s3Uri", "") diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index 01e25f4671e..f66dcf130b8 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -25,6 +25,7 @@ import base64 from typing import TYPE_CHECKING, Any, Final import httpx +from httpx._types import RequestFiles from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.bedrock.common_utils import BedrockError @@ -165,7 +166,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): image_edit_optional_request_params: dict, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[dict, Any]: + ) -> tuple[dict, RequestFiles]: """ Transform OpenAI-style request to Bedrock Stability request format. diff --git a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py index e1b06791c9d..12ace32f43b 100644 --- a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py @@ -80,9 +80,7 @@ class AmazonTitanImageGenerationConfig: non_default_params: dict, optional_params: dict, ): - from typing import Any - - image_generation_config: Final[dict[str, Any]] = {} + image_generation_config: Final[dict[str, object]] = {} for k, v in non_default_params.items(): if k == "size" and v is not None: width, height = v.split("x") @@ -106,11 +104,9 @@ class AmazonTitanImageGenerationConfig: text: str, optional_params: dict, ) -> AmazonTitanImageGenerationRequestBody: - from typing import Any - image_generation_config = optional_params.pop("imageGenerationConfig", {}) negative_text: Final = optional_params.pop("negativeText", None) - text_to_image_params: Final[dict[str, Any]] = {"text": text} + text_to_image_params: Final = AmazonTitanTextToImageParams(text=text) if negative_text: text_to_image_params["negativeText"] = negative_text task_type: Final = optional_params.pop("taskType", "TEXT_IMAGE") @@ -121,7 +117,7 @@ class AmazonTitanImageGenerationConfig: } return AmazonTitanImageGenerationRequestBody( taskType=task_type, - textToImageParams=AmazonTitanTextToImageParams(**text_to_image_params), + textToImageParams=text_to_image_params, imageGenerationConfig=AmazonNovaCanvasImageGenerationConfig(**image_generation_config), ) diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 8847381cbc9..4a8308bb46a 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -2,6 +2,7 @@ import json from typing import TYPE_CHECKING, Any, Final, cast import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -24,6 +25,8 @@ if TYPE_CHECKING: else: AWSPreparedRequest = Any +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) + class BedrockRerankHandler(BaseAWSLLM): async def arerank( @@ -55,13 +58,13 @@ class BedrockRerankHandler(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return BedrockRerankConfig()._transform_response(response.json()) + return BedrockRerankConfig()._transform_response(_JSON_DICT.validate_python(response.json())) def rerank( self, model: str, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], optional_params: dict, logging_obj: LitellmLogging, top_n: int | None = None, @@ -136,7 +139,7 @@ class BedrockRerankHandler(BaseAWSLLM): api_key="", ) - response_json: Final = response.json() + response_json: Final = _JSON_DICT.validate_python(response.json()) return BedrockRerankConfig()._transform_response(response_json) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 19ba7d5673b..565fc12aa1e 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -37,6 +37,7 @@ from collections.abc import Iterator, Mapping, Sequence from typing import Final import httpx +from pydantic import TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -81,6 +82,8 @@ _SSE_EVENT_SEPARATOR: Final = re.compile(r"\r?\n[ \t]*\r?\n") _SSE_LINE_PREFIXES: Final = ("event:", "data:", ":", "id:", "retry:") +_JSON_VALUE: Final = TypeAdapter(object) + def _gateway_host_match(api_base: str) -> re.Match[str] | None: return _GATEWAY_HOST_PATTERN.fullmatch(httpx.URL(api_base).host) @@ -128,7 +131,7 @@ def _parse_result_items(raw_text: object) -> tuple[Mapping[str, object], ...]: if not isinstance(raw_text, str): return () try: - parsed: Final = json.loads(raw_text) + parsed: Final = _JSON_VALUE.validate_python(json.loads(raw_text)) except json.JSONDecodeError: return () return _result_items(parsed) @@ -147,7 +150,7 @@ def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]: if not payload: continue try: - parsed = json.loads(payload) + parsed = _JSON_VALUE.validate_python(json.loads(payload)) except json.JSONDecodeError: continue if isinstance(parsed, dict): diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index e5c2bdf59e9..9d586fba26a 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -8,9 +8,11 @@ API Reference: https://docs.bfl.ai/ """ import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -36,6 +38,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): """ @@ -218,7 +222,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): https://docs.bfl.ai/flux_models/flux_1_1_pro """ # Build request body with prompt - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "prompt": prompt, } @@ -275,7 +279,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): message=f"Error parsing BFL response: {e}", ) - result: Final = response_data.get("result", {}) + result: Final = _JSON_OBJECT.validate_python(response_data).get("result", {}) if not model_response.data: model_response.data = [] diff --git a/litellm/llms/clarifai/chat/transformation.py b/litellm/llms/clarifai/chat/transformation.py index a0946254de0..cea6e6c7440 100644 --- a/litellm/llms/clarifai/chat/transformation.py +++ b/litellm/llms/clarifai/chat/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import OpenAIError @@ -20,6 +22,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) + class ClarifaiConfig(OpenAIGPTConfig): """ @@ -111,7 +115,7 @@ class ClarifaiConfig(OpenAIGPTConfig): headers=raw_response.headers, ) from e - response: Final = ModelResponse(**completion_response) + response: Final = ModelResponse(**_JSON_OBJECT.validate_python(completion_response)) if response.model is not None: response.model = "clarifai/" + model diff --git a/litellm/llms/claude_code/harness/transformation.py b/litellm/llms/claude_code/harness/transformation.py index ea3cdd67ffb..f75d86c7e63 100644 --- a/litellm/llms/claude_code/harness/transformation.py +++ b/litellm/llms/claude_code/harness/transformation.py @@ -16,6 +16,8 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final +from pydantic import ConfigDict, TypeAdapter + from litellm.harness.errors import HarnessError, OptionsMismatch from litellm.harness.options import ClaudeCodeOptions from litellm.harness.types import ( @@ -48,6 +50,7 @@ if TYPE_CHECKING: CLAUDE_BINARY: Final = "claude" SYNTHETIC_MODEL: Final = "" +_BLOCK: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(hide_input_in_errors=True)) BASE_COMMAND: Final = ("-p", "--output-format", "stream-json", "--verbose", "--input-format", "text") @@ -157,7 +160,7 @@ def _stringify_block(block: object) -> str: return json.dumps(block, ensure_ascii=False) -def _message_blocks(event: Mapping[str, object]) -> Sequence[Any]: +def _message_blocks(event: Mapping[str, object]) -> Sequence[object]: message: Final = event.get("message") content: Final = message.get("content") if isinstance(message, Mapping) else None if isinstance(content, str): @@ -211,8 +214,7 @@ def _user_events(event: Mapping[str, object]) -> Sequence[Event]: output=stringify_tool_output(block.get("content")), is_error=bool(block.get("is_error", False)), ) - for block in _message_blocks(event) - if _is_tool_result(block) + for block in map(_BLOCK.validate_python, filter(_is_tool_result, _message_blocks(event))) ) ) diff --git a/litellm/llms/codex/harness/transformation.py b/litellm/llms/codex/harness/transformation.py index ab17cf3d869..86a5f8274e6 100644 --- a/litellm/llms/codex/harness/transformation.py +++ b/litellm/llms/codex/harness/transformation.py @@ -17,6 +17,8 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final +from pydantic import ConfigDict, TypeAdapter + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.harness.errors import HarnessError, OptionsMismatch from litellm.harness.options import CodexOptions @@ -61,6 +63,7 @@ MANAGED_CONFIG_KEYS: Final = frozenset( ) _BARE_TOML_KEY: Final = re.compile(r"^[A-Za-z0-9_-]+$") _TOOL_ITEM_TYPES: Final = frozenset({"command_execution", "file_change", "web_search", "mcp_tool_call"}) +_CHANGES: Final = TypeAdapter(tuple[Mapping[object, object], ...], config=ConfigDict(hide_input_in_errors=True)) @dataclass @@ -92,7 +95,7 @@ def _tool_input(item: Mapping[str, Any]) -> tuple[str, str, Mapping[str, Any], b return name, tool, tool_args, False -def _tool_output(item: Mapping[str, Any]) -> tuple[str, bool]: +def _tool_output(item: Mapping[str, object]) -> tuple[str, bool]: """(output text, is_error) for a completed tool-like item.""" item_type = item.get("type") status = item.get("status") @@ -101,7 +104,8 @@ def _tool_output(item: Mapping[str, Any]) -> tuple[str, bool]: is_error = status == "failed" or (exit_code is not None and exit_code != 0) return str(item.get("aggregated_output") or ""), is_error if item_type == "file_change": - lines = (f"{c.get('kind', '')} {c.get('path', '')}".strip() for c in item.get("changes") or ()) + changes: Final = _CHANGES.validate_python(item.get("changes") or ()) + lines: Final = (f"{c.get('kind', '')} {c.get('path', '')}".strip() for c in changes) return "\n".join(lines), status == "failed" if item_type == "web_search": return "", status == "failed" @@ -118,7 +122,7 @@ def _tool_output(item: Mapping[str, Any]) -> tuple[str, bool]: def _tool_item_events( - item_id: str, item: Mapping[str, Any], completed: bool, state: CodexStreamState + item_id: str, item: Mapping[str, object], completed: bool, state: CodexStreamState ) -> Iterator[Event]: if item_id not in state.started: state.started.add(item_id) @@ -129,7 +133,7 @@ def _tool_item_events( yield ToolResult(id=item_id, output=output, is_error=is_error) -def _item_events(event_type: str, item: Mapping[str, Any], state: CodexStreamState) -> Sequence[Event]: +def _item_events(event_type: str, item: Mapping[str, object], state: CodexStreamState) -> Sequence[Event]: item_type = item.get("type") item_id = str(item.get("id") or "") completed = event_type == "item.completed" @@ -181,7 +185,7 @@ def _config_override(key: object, value: object) -> str: return f"{dotted}={toml_value(value)}" -def config_overrides(config: Mapping[str, Any]) -> Sequence[str]: +def config_overrides(config: Mapping[str, object]) -> Sequence[str]: """`-c` override strings for CodexOptions.config, rejecting managed keys.""" overrides: Final = (_config_override(key, value) for key, value in config.items()) return list(overrides) # mutable-ok: public helper; tests compare to a list @@ -310,7 +314,7 @@ class CodexHarnessConfig(BaseCLIHarnessConfig): def create_stream_state(self) -> CodexStreamState: return CodexStreamState() - def transform_stream_line(self, line: Mapping[str, Any], state: CodexStreamState) -> Sequence[Event]: + def transform_stream_line(self, line: Mapping[str, object], state: CodexStreamState) -> Sequence[Event]: """turn.completed usage is ignored on purpose: the session endpoint accounts it.""" event_type = line.get("type") if event_type == "thread.started": diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index a8e755406d8..b6d94d4853a 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -1,7 +1,8 @@ from collections.abc import Mapping -from typing import Any, Final +from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -12,6 +13,8 @@ from litellm.types.rerank import OptionalRerankParams, RerankRequest, RerankResp from ..common_utils import CohereError +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class CohereRerankConfig(BaseRerankConfig): """ @@ -51,7 +54,7 @@ class CohereRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, @@ -148,7 +151,7 @@ class CohereRerankConfig(BaseRerankConfig): except Exception: raise CohereError(message=raw_response.text, status_code=raw_response.status_code) - return RerankResponse(**raw_response_json) + return RerankResponse.model_validate(_JSON_OBJECT.validate_python(raw_response_json)) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return CohereError(message=error_message, status_code=status_code) diff --git a/litellm/llms/cometapi/image_generation/transformation.py b/litellm/llms/cometapi/image_generation/transformation.py index 4432c151a64..567dbfd308e 100644 --- a/litellm/llms/cometapi/image_generation/transformation.py +++ b/litellm/llms/cometapi/image_generation/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -20,6 +22,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) + class CometAPIImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://api.cometapi.com" @@ -155,7 +160,8 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): # CometAPI returns OpenAI-compatible format # Expected format: {"created": timestamp, "data": [{"url": "...", "b64_json": "..."}]} if "data" in response_data: - for image_data in response_data["data"]: + payload: Final = _JSON_OBJECT.validate_python(response_data) + for image_data in _JSON_OBJECTS.validate_python(payload["data"]): image_obj = ImageObject( b64_json=image_data.get("b64_json"), url=image_data.get("url"), diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index ffa60a3d9bc..6a172e1e074 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -23,9 +23,11 @@ Response format: } """ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -45,6 +47,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) + DEFAULT_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation" CHAT_COMPATIBLE_MODE_PATH: Final = "/compatible-mode/v1" @@ -192,9 +197,10 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): # DashScope can return API-level errors in a 200 response body. # Example: {"code": "InvalidParameter", "message": "Size not supported"} - if "code" in response_data and "output" not in response_data: + response_object: Final = _JSON_OBJECT.validate_python(response_data) + if "code" in response_object and "output" not in response_object: raise self.get_error_class( - error_message=str(response_data.get("message", response_data)), + error_message=str(response_object.get("message", response_object)), status_code=raw_response.status_code, headers=raw_response.headers, ) @@ -202,9 +208,11 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): if not model_response.data: model_response.data = [] - choices: Final = response_data.get("output", {}).get("choices", []) + output: Final = _JSON_OBJECT.validate_python(response_object.get("output", {})) + choices: Final = _JSON_OBJECTS.validate_python(output.get("choices", [])) for choice in choices: - content_list = choice.get("message", {}).get("content", []) + message = _JSON_OBJECT.validate_python(choice.get("message", {})) + content_list = _JSON_OBJECTS.validate_python(message.get("content", [])) for content_item in content_list: image_url = content_item.get("image") if image_url: diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 14ad756ec9c..4ec7d697f02 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -26,10 +26,11 @@ as supported only for gte-rerank-v2 / qwen3-vl-rerank. Docs - https://help.aliyun.com/zh/model-studio/text-rerank-api """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Iterable, Mapping +from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -48,6 +49,11 @@ from ..common_utils import DashScopeError, resolve_dashscope_family_rerank_api_b DEFAULT_RERANK_URL: Final = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_OPTIONAL_INT: Final[TypeAdapter[int | None]] = TypeAdapter(int | None) +_STR: Final = TypeAdapter(str) + class DashScopeRerankConfig(BaseRerankConfig): """ @@ -117,7 +123,7 @@ class DashScopeRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, @@ -197,7 +203,8 @@ class DashScopeRerankConfig(BaseRerankConfig): message=response_json.get("message", str(response_json)), ) - results: Final = response_json.get("results") + payload: Final = _JSON_OBJECT.validate_python(response_json) + results: Final = payload.get("results") if results is None: raise DashScopeError( status_code=raw_response.status_code, @@ -210,7 +217,7 @@ class DashScopeRerankConfig(BaseRerankConfig): # "document": {"text": "..."} # which already matches LiteLLM's RerankResponseDocument shape. transformed_results: Final[list[dict]] = [] - for r in results: + for r in _JSON_OBJECTS.validate_python(results): item: dict[str, object] = { "index": r["index"], "relevance_score": r["relevance_score"], @@ -223,14 +230,14 @@ class DashScopeRerankConfig(BaseRerankConfig): item["document"] = {"text": doc} transformed_results.append(item) - usage: Final = response_json.get("usage") or {} - total_tokens: Final = usage.get("total_tokens") + usage: Final = _JSON_OBJECT.validate_python(payload.get("usage") or {}) + total_tokens: Final = _OPTIONAL_INT.validate_python(usage.get("total_tokens")) billed_units: Final = RerankBilledUnits(total_tokens=total_tokens) tokens: Final = RerankTokens(input_tokens=total_tokens) meta: Final = RerankResponseMeta(billed_units=billed_units, tokens=tokens) return RerankResponse( - id=response_json.get("id") or str(uuid.uuid4()), + id=_STR.validate_python(payload.get("id") or str(uuid.uuid4())), results=transformed_results, meta=meta, ) diff --git a/litellm/llms/deepagents/harness/transformation.py b/litellm/llms/deepagents/harness/transformation.py index 14f87f0753b..1f8728230bc 100644 --- a/litellm/llms/deepagents/harness/transformation.py +++ b/litellm/llms/deepagents/harness/transformation.py @@ -173,7 +173,7 @@ def stream_events( ] -def tool_call_event(call: Mapping[str, Any]) -> ToolCall: +def tool_call_event(call: Mapping[str, object]) -> ToolCall: native = str(call.get("name") or "") args = call.get("args") return ToolCall( @@ -185,7 +185,7 @@ def tool_call_event(call: Mapping[str, Any]) -> ToolCall: ) -def _node_messages(update: Mapping[Any, Any]) -> Iterator[object]: +def _node_messages(update: Mapping[object, object]) -> Iterator[object]: for node, delta in update.items(): if node not in _EVENT_NODES or not isinstance(delta, Mapping): continue @@ -231,7 +231,7 @@ def interrupts_in( return list(items) # mutable-ok: list return; callers/tests compare to lists -def final_ai_text(messages: Sequence[Any]) -> str: +def final_ai_text(messages: Sequence[object]) -> str: for message in reversed(messages): if getattr(message, "type", None) == "ai": text = content_text(getattr(message, "content", "")) diff --git a/litellm/llms/deepgram/audio_transcription/transformation.py b/litellm/llms/deepgram/audio_transcription/transformation.py index 95bdd406c3f..c5c9fe12423 100644 --- a/litellm/llms/deepgram/audio_transcription/transformation.py +++ b/litellm/llms/deepgram/audio_transcription/transformation.py @@ -2,10 +2,12 @@ Translates from OpenAI's `/v1/audio/transcriptions` to Deepgram's `/v1/listen` """ +from collections.abc import Iterable, Mapping from typing import Final from urllib.parse import urlencode from httpx import Headers, Response +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -22,6 +24,9 @@ from ...base_llm.audio_transcription.transformation import ( ) from ..common_utils import DeepgramException +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) + class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): def get_supported_openai_params(self, model: str) -> list[OpenAIAudioTranscriptionOptionalParams]: @@ -105,7 +110,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): response["task"] = "transcribe" # Use detected_language if available, otherwise default to "en" - detected_language: Final = first_channel.get("detected_language") + detected_language: Final = _JSON_OBJECT.validate_python(first_channel).get("detected_language") response["language"] = detected_language if detected_language else "en" response["duration"] = response_json["metadata"]["duration"] @@ -114,7 +119,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): if "words" in first_alternative: response["words"] = [ {"word": word["word"], "start": word["start"], "end": word["end"]} - for word in first_alternative["words"] + for word in _JSON_OBJECTS.validate_python(first_alternative["words"]) ] # Store full response in hidden params diff --git a/litellm/llms/e2b/sandbox/transformation.py b/litellm/llms/e2b/sandbox/transformation.py index 4928ca0c092..d2702b116dc 100644 --- a/litellm/llms/e2b/sandbox/transformation.py +++ b/litellm/llms/e2b/sandbox/transformation.py @@ -8,9 +8,11 @@ Talks to e2b's REST API directly over httpx (no e2b SDK dependency): """ import json +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.sandbox.transformation import ( SANDBOX_MAX_OUTPUT_BYTES, @@ -32,6 +34,12 @@ JUPYTER_PORT: Final = 49999 DEFAULT_SANDBOX_TIMEOUT: Final = 300 MAX_OUTPUT_BYTES: Final = SANDBOX_MAX_OUTPUT_BYTES +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_MESSAGE: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter( + Mapping[str, object] | None, config=ConfigDict(hide_input_in_errors=True) +) +_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True)) + class E2BSandboxConfig(BaseSandboxConfig): def _http(self, client: AsyncHTTPHandler | None) -> AsyncHTTPHandler: @@ -73,12 +81,14 @@ class E2BSandboxConfig(BaseSandboxConfig): headers={"X-API-Key": key, "Content-Type": "application/json"}, json=body, ) - data: Final = response.json() + data: Final = _JSON_OBJECT.validate_python(response.json()) - handle: Final = ContainerHandle( - id=data["sandboxID"], - provider="e2b", - domain=data.get("domain") or E2B_DEFAULT_DOMAIN, + handle: Final = ContainerHandle.model_validate( + { + "id": data["sandboxID"], + "provider": "e2b", + "domain": data.get("domain") or E2B_DEFAULT_DOMAIN, + } ) handle._hidden_params = { "envd_access_token": data.get("envdAccessToken"), @@ -158,7 +168,7 @@ class E2BSandboxConfig(BaseSandboxConfig): def _parse_lines(lines: list[str]) -> CodeExecutionResult: def _try_parse(stripped: str): try: - return json.loads(stripped) + return _MESSAGE.validate_python(json.loads(stripped)) except json.JSONDecodeError: return None @@ -178,10 +188,12 @@ class E2BSandboxConfig(BaseSandboxConfig): None, ) - return CodeExecutionResult( - stdout="".join(m.get("text", "") for m in of_type("stdout")), - stderr="".join(m.get("text", "") for m in of_type("stderr")), - results=[{k: v for k, v in m.items() if k != "type"} for m in of_type("result")], - error=error, - execution_count=execution_count, + return CodeExecutionResult.model_validate( + { + "stdout": "".join(_TEXT.validate_python(m.get("text", "")) for m in of_type("stdout")), + "stderr": "".join(_TEXT.validate_python(m.get("text", "")) for m in of_type("stderr")), + "results": [{k: v for k, v in m.items() if k != "type"} for m in of_type("result")], + "error": error, + "execution_count": execution_count, + } ) diff --git a/litellm/llms/elevenlabs/audio_transcription/transformation.py b/litellm/llms/elevenlabs/audio_transcription/transformation.py index a2b25760aac..0dbf61917e2 100644 --- a/litellm/llms/elevenlabs/audio_transcription/transformation.py +++ b/litellm/llms/elevenlabs/audio_transcription/transformation.py @@ -2,9 +2,11 @@ Translates from OpenAI's `/v1/audio/transcriptions` to ElevenLabs's `/v1/speech-to-text` """ +from collections.abc import Iterable, Mapping from typing import Final from httpx import Headers, Response +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file @@ -22,6 +24,9 @@ from ...base_llm.audio_transcription.transformation import ( ) from ..common_utils import ElevenLabsException +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) + class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @property @@ -115,21 +120,22 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): """ try: response_json: Final = raw_response.json() + response_object: Final = _JSON_OBJECT.validate_python(response_json) # Extract the main transcript text - text: Final = response_json.get("text", "") + text: Final = response_object.get("text", "") # Create TranscriptionResponse object response: Final = TranscriptionResponse(text=text) # Add additional metadata matching OpenAI format response["task"] = "transcribe" - response["language"] = response_json.get("language_code", "unknown") + response["language"] = response_object.get("language_code", "unknown") # Map ElevenLabs words to OpenAI format - if "words" in response_json: + if "words" in response_object: response["words"] = [] - for word_data in response_json["words"]: + for word_data in _JSON_OBJECTS.validate_python(response_object["words"]): # Only include actual words, skip spacing and audio events if word_data.get("type") == "word": response["words"].append( diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py index 7c63e1077f1..d4d2941db43 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageResponse @@ -15,6 +17,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class FalAIFluxProV11UltraConfig(FalAIBaseConfig): """ @@ -228,16 +232,17 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): if not model_response.data: model_response.data = [] - images: Final = response_data.get("images", []) + response_object: Final = _JSON_OBJECT.validate_python(response_data) + images: Final = response_object.get("images", []) model_response.data.extend(fal_images_to_image_objects(images)) # Add additional metadata from Flux Pro response if hasattr(model_response, "_hidden_params"): - if "seed" in response_data: - model_response._hidden_params["seed"] = response_data["seed"] - if "timings" in response_data: - model_response._hidden_params["timings"] = response_data["timings"] - if "has_nsfw_concepts" in response_data: - model_response._hidden_params["has_nsfw_concepts"] = response_data["has_nsfw_concepts"] + if "seed" in response_object: + model_response._hidden_params["seed"] = response_object["seed"] + if "timings" in response_object: + model_response._hidden_params["timings"] = response_object["timings"] + if "has_nsfw_concepts" in response_object: + model_response._hidden_params["has_nsfw_concepts"] = response_object["has_nsfw_concepts"] return model_response diff --git a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py index ad1852a622b..970addbd5a2 100644 --- a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageObject, ImageResponse @@ -15,6 +17,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class FalAIIdeogramV3Config(FalAIBaseConfig): """ @@ -169,7 +173,8 @@ class FalAIIdeogramV3Config(FalAIBaseConfig): if not model_response.data: model_response.data = [] - images: Final = response_data.get("images", []) + response_object: Final = _JSON_OBJECT.validate_python(response_data) + images: Final = response_object.get("images", []) if isinstance(images, list): for image_entry in images: if isinstance(image_entry, dict): @@ -184,7 +189,7 @@ class FalAIIdeogramV3Config(FalAIBaseConfig): ) ) - if hasattr(model_response, "_hidden_params") and "seed" in response_data: - model_response._hidden_params["seed"] = response_data["seed"] + if hasattr(model_response, "_hidden_params") and "seed" in response_object: + model_response._hidden_params["seed"] = response_object["seed"] return model_response diff --git a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py index 934ce420d53..038a06c5749 100644 --- a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageObject, ImageResponse @@ -15,6 +17,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class FalAIRecraftV3Config(FalAIBaseConfig): """ @@ -89,7 +93,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): return optional_params - def _map_image_size(self, size: str) -> Any: + def _map_image_size(self, size: str) -> str | Mapping[str, int]: """ Map OpenAI size format to Recraft v3 image_size format. @@ -203,7 +207,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): model_response.data = [] # Handle Recraft v3 response format - images: Final = response_data.get("images", []) + images: Final = _JSON_OBJECT.validate_python(response_data).get("images", []) if isinstance(images, list): for image_data in images: if isinstance(image_data, dict): diff --git a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py index 79d8800773b..659011ae537 100644 --- a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py +++ b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageObject, ImageResponse @@ -15,6 +17,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class FalAIStableDiffusionConfig(FalAIBaseConfig): """ @@ -124,7 +128,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): return optional_params - def _map_image_size(self, size: str) -> Any: + def _map_image_size(self, size: str) -> str | Mapping[str, int]: """ Map OpenAI size format to Stable Diffusion image_size format. @@ -243,7 +247,8 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): model_response.data = [] # Handle Stable Diffusion response format - images: Final = response_data.get("images", []) + response_object: Final = _JSON_OBJECT.validate_python(response_data) + images: Final = response_object.get("images", []) if isinstance(images, list): for image_data in images: if isinstance(image_data, dict): @@ -264,11 +269,11 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): # Add additional metadata from Stable Diffusion response if hasattr(model_response, "_hidden_params"): - if "seed" in response_data: - model_response._hidden_params["seed"] = response_data["seed"] - if "timings" in response_data: - model_response._hidden_params["timings"] = response_data["timings"] - if "has_nsfw_concepts" in response_data: - model_response._hidden_params["has_nsfw_concepts"] = response_data["has_nsfw_concepts"] + if "seed" in response_object: + model_response._hidden_params["seed"] = response_object["seed"] + if "timings" in response_object: + model_response._hidden_params["timings"] = response_object["timings"] + if "has_nsfw_concepts" in response_object: + model_response._hidden_params["has_nsfw_concepts"] = response_object["has_nsfw_concepts"] return model_response diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 509dbd5ff24..3d9d813c6e9 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -4,10 +4,11 @@ Fireworks AI Rerank API transformation Reference: https://docs.fireworks.ai/inference-api-reference/rerank """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Iterable, Mapping, Sequence +from typing import Final import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -23,6 +24,26 @@ from litellm.types.rerank import ( ) +class _FireworksAIUsageFields(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + total_tokens: int | None = 0 + prompt_tokens: int | None = 0 + completion_tokens: int | None = 0 + + +class _FireworksAIResultFields(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, hide_input_in_errors=True) + + index: int | float | str + relevance_score: int | float | str + + +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_STR: Final = TypeAdapter(str) + + class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): """ Fireworks AI Rerank API configuration @@ -59,7 +80,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: Sequence[str | Mapping[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, @@ -204,23 +225,18 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): # } # Extract usage information - usage: Final = raw_response_json.get("usage", {}) - _billed_units: Final = RerankBilledUnits(search_units=usage.get("total_tokens", 0)) - _tokens: Final = RerankTokens( - input_tokens=usage.get("prompt_tokens", 0), - output_tokens=usage.get("completion_tokens", 0), - ) - rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) + response_json: Final = _JSON_OBJECT.validate_python(raw_response_json) + usage: Final = _JSON_OBJECT.validate_python(response_json.get("usage", {})) # Extract results - Fireworks AI uses "data" instead of "results" - _results: Final[list[dict] | None] = raw_response_json.get("data") or raw_response_json.get("results") + _results: Final = response_json.get("data") or response_json.get("results") if _results is None: raise ValueError(f"No results found in the response={raw_response_json}") rerank_results: Final[list[RerankResponseResult]] = [] - for result in _results: + for result in _JSON_OBJECTS.validate_python(_results): # Validate required fields exist if not all(key in result for key in ["index", "relevance_score"]): raise ValueError(f"Missing required fields in the result={result}") @@ -239,9 +255,10 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): document = RerankResponseDocument(text=str(text)) # Create typed result + fields = _FireworksAIResultFields.model_validate(result) rerank_result = RerankResponseResult( - index=int(result["index"]), - relevance_score=float(result["relevance_score"]), + index=int(fields.index), + relevance_score=float(fields.relevance_score), ) # Only add document if it exists @@ -250,7 +267,15 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): rerank_results.append(rerank_result) - response_id: Final = raw_response_json.get("id") or str(uuid.uuid4()) + usage_fields: Final = _FireworksAIUsageFields.model_validate(usage) + _billed_units: Final = RerankBilledUnits(search_units=usage_fields.total_tokens) + _tokens: Final = RerankTokens( + input_tokens=usage_fields.prompt_tokens, + output_tokens=usage_fields.completion_tokens, + ) + rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) + + response_id: Final = _STR.validate_python(response_json.get("id") or str(uuid.uuid4())) return RerankResponse( id=response_id, diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index bb8d7455031..4d13d0d2e6b 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -31,6 +33,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) + class GoogleImageGenConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" @@ -216,13 +221,15 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): # Extract usage metadata for Gemini models if "usageMetadata" in response_data: - model_response.usage = transform_gemini_image_usage(response_data["usageMetadata"]) + model_response.usage = transform_gemini_image_usage( + _JSON_DICT.validate_python(response_data["usageMetadata"]) + ) web_search_requests: Final = get_gemini_image_web_search_requests(response_data) if web_search_requests and model_response.usage is not None: setattr(model_response.usage, "web_search_requests", web_search_requests) else: # Original Imagen format - predictions with generated images - predictions: Final = response_data.get("predictions", []) + predictions: Final = _JSON_OBJECTS.validate_python(response_data.get("predictions", [])) for prediction in predictions: # Google AI returns base64 encoded images in the prediction model_response.data.append( diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 927f5e944b6..321f5525433 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -8,9 +8,11 @@ API Documentation: https://developers.sber.ru/docs/ru/gigachat/api/reference/res from __future__ import annotations import types +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -22,6 +24,8 @@ from litellm.types.utils import EmbeddingResponse from ..authenticator import get_access_token +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class GigaChatEmbeddingError(BaseLLMException): """GigaChat Embedding API error.""" @@ -165,7 +169,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): "total_tokens": total_tokens, } - return EmbeddingResponse(**response_json) + return EmbeddingResponse.model_validate(_JSON_OBJECT.validate_python(response_json)) def validate_environment( self, diff --git a/litellm/llms/github_copilot/embedding/transformation.py b/litellm/llms/github_copilot/embedding/transformation.py index 7ea7a89b4ca..c1acf6fccaa 100644 --- a/litellm/llms/github_copilot/embedding/transformation.py +++ b/litellm/llms/github_copilot/embedding/transformation.py @@ -11,6 +11,7 @@ import os from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError @@ -33,6 +34,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_RESPONSE_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) + class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): """ @@ -154,7 +157,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): logging_obj.post_call(original_response=raw_response.text) # GitHub Copilot returns standard OpenAI-compatible embedding response - response_json: Final = raw_response.json() + response_json: Final = _RESPONSE_OBJECT.validate_python(raw_response.json()) return convert_to_model_response_object( response_object=response_json, diff --git a/litellm/llms/jina_ai/embedding/transformation.py b/litellm/llms/jina_ai/embedding/transformation.py index 260d9e6e494..e184ad2628b 100644 --- a/litellm/llms/jina_ai/embedding/transformation.py +++ b/litellm/llms/jina_ai/embedding/transformation.py @@ -7,9 +7,11 @@ Docs - https://jina.ai/embeddings/ """ import types +from collections.abc import Mapping from typing import Final, cast import httpx +from pydantic import ConfigDict, TypeAdapter from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -22,6 +24,8 @@ from litellm.utils import is_base64_encoded from ..common_utils import JinaAIError +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class JinaAIEmbeddingConfig(BaseEmbeddingConfig): """ @@ -139,7 +143,7 @@ class JinaAIEmbeddingConfig(BaseEmbeddingConfig): additional_args={"complete_input_dict": request_data}, original_response=response_json, ) - return EmbeddingResponse(**response_json) + return EmbeddingResponse.model_validate(_JSON_OBJECT.validate_python(response_json)) def validate_environment( self, diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 2a81c38fe34..f4a3a9f0f09 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -6,10 +6,11 @@ Why separate file? Make it easy to see how transformation works Docs - https://jina.ai/reranker """ -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from typing import Final from httpx import URL, Response +from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj @@ -23,6 +24,12 @@ from litellm.types.rerank import ( ) from litellm.types.utils import ModelInfo +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_BILLED_UNITS: Final = TypeAdapter(RerankBilledUnits) +_TOKENS: Final = TypeAdapter(RerankTokens) +_STR: Final = TypeAdapter(str) + class JinaAIRerankConfig(BaseRerankConfig): def get_supported_cohere_rerank_params(self, model: str) -> list: @@ -100,13 +107,11 @@ class JinaAIRerankConfig(BaseRerankConfig): logging_obj.post_call(original_response=raw_response.text) - _json_response: Final = raw_response.json() + _json_response: Final = _JSON_OBJECT.validate_python(raw_response.json()) - _billed_units: Final = RerankBilledUnits(**_json_response.get("usage", {})) - _tokens: Final = RerankTokens(**_json_response.get("usage", {})) - rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) + usage: Final = _JSON_OBJECT.validate_python(_json_response.get("usage", {})) - _results: Final[list[dict] | None] = _json_response.get("results") + _results: Final = _json_response.get("results") if _results is None: raise ValueError(f"No results found in the response={_json_response}") @@ -115,7 +120,7 @@ class JinaAIRerankConfig(BaseRerankConfig): # Jina AI returns: {"index": 0, "relevance_score": 0.72, "document": "hello"} # LiteLLM expects: {"index": 0, "relevance_score": 0.72, "document": {"text": "hello"}} transformed_results: Final = [] - for result in _results: + for result in _JSON_OBJECTS.validate_python(_results): transformed_result = { "index": result["index"], "relevance_score": result["relevance_score"], @@ -128,8 +133,12 @@ class JinaAIRerankConfig(BaseRerankConfig): transformed_result["document"] = result["document"] transformed_results.append(transformed_result) + _billed_units: Final = _BILLED_UNITS.validate_python(usage) + _tokens: Final = _TOKENS.validate_python(usage) + rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) + return RerankResponse( - id=_json_response.get("id") or str(uuid.uuid4()), + id=_STR.validate_python(_json_response.get("id") or str(uuid.uuid4())), results=transformed_results, meta=rerank_meta, ) # Return response diff --git a/litellm/llms/manus/files/transformation.py b/litellm/llms/manus/files/transformation.py index 8c7b02f8d92..94b0e62199d 100644 --- a/litellm/llms/manus/files/transformation.py +++ b/litellm/llms/manus/files/transformation.py @@ -11,10 +11,12 @@ Reference: https://open.manus.im/docs/openai-compatibility#file-management """ import time -from typing import Any, Final +from collections.abc import Iterable, Mapping +from typing import Final import httpx from openai.types.file_deleted import FileDeleted +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -39,6 +41,10 @@ from litellm.types.utils import LlmProviders MANUS_API_BASE: Final = "https://api.manus.im" +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True)) + class ManusFilesConfig(BaseFilesConfig): """ @@ -337,8 +343,7 @@ class ManusFilesConfig(BaseFilesConfig): litellm_params: dict, ) -> FileDeleted: """Transform delete file response.""" - response_json: Final = raw_response.json() - return FileDeleted(**response_json) + return FileDeleted.model_validate(_JSON_OBJECT.validate_python(raw_response.json())) def transform_list_files_request( self, @@ -366,19 +371,20 @@ class ManusFilesConfig(BaseFilesConfig): litellm_params: dict, ) -> list[OpenAIFileObject]: """Transform list files response.""" - response_json: Final = raw_response.json() - files_data: Final = response_json.get("data", []) + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) + files_data: Final = _JSON_OBJECTS.validate_python(response_json.get("data", [])) return [self._parse_file_dict(f) for f in files_data] - def _parse_file_dict(self, file_dict: dict[str, Any]) -> OpenAIFileObject: + def _parse_file_dict(self, file_dict: Mapping[str, object]) -> OpenAIFileObject: """Parse a file dict into OpenAIFileObject.""" created_at_str: Final = file_dict.get("created_at", "") if created_at_str: + created_at_text: Final = _TEXT.validate_python(created_at_str) try: created_at = int( time.mktime( time.strptime( - created_at_str.replace("Z", "+00:00")[:19], + created_at_text.replace("Z", "+00:00")[:19], "%Y-%m-%dT%H:%M:%S", ) ) @@ -388,15 +394,17 @@ class ManusFilesConfig(BaseFilesConfig): else: created_at = int(time.time()) - return OpenAIFileObject( - id=file_dict.get("id", ""), - bytes=file_dict.get("bytes", 0), - created_at=created_at, - filename=file_dict.get("filename", ""), - object="file", - purpose=file_dict.get("purpose", "assistants"), - status=file_dict.get("status", "uploaded"), - status_details=file_dict.get("status_details"), + return OpenAIFileObject.model_validate( + { + "id": file_dict.get("id", ""), + "bytes": file_dict.get("bytes", 0), + "created_at": created_at, + "filename": file_dict.get("filename", ""), + "object": "file", + "purpose": file_dict.get("purpose", "assistants"), + "status": file_dict.get("status", "uploaded"), + "status_details": file_dict.get("status_details"), + } ) def transform_file_content_request( diff --git a/litellm/llms/minimax/text_to_speech/transformation.py b/litellm/llms/minimax/text_to_speech/transformation.py index e38a8a2c3a3..1cabc43eb1c 100644 --- a/litellm/llms/minimax/text_to_speech/transformation.py +++ b/litellm/llms/minimax/text_to_speech/transformation.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from httpx import Headers +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -26,6 +27,9 @@ else: LiteLLMLoggingObj = Any HttpxBinaryResponseContent = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_STR: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True)) + class MinimaxException(BaseLLMException): """Custom exception for MiniMax API errors""" @@ -299,7 +303,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): try: # Parse JSON response - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # MiniMax API response format check # The API can return different structures: @@ -320,7 +324,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): # Extract audio data # MiniMax returns audio in "data" field - data: Final = response_json.get("data", {}) + data: Final = _JSON_OBJECT.validate_python(response_json.get("data", {})) # Check if response contains a URL (output_format='url') audio_url: Final = data.get("audio_url", None) @@ -334,7 +338,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get hex-encoded audio data - audio_hex: Final = data.get("audio", "") or response_json.get("audio_file", "") + audio_hex: Final = _STR.validate_python(data.get("audio", "") or response_json.get("audio_file", "") or "") if not audio_hex: raise MinimaxException( diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index 3a8a37307d6..08cdbacdad0 100644 --- a/litellm/llms/modelscope/image_generation/transformation.py +++ b/litellm/llms/modelscope/image_generation/transformation.py @@ -6,9 +6,11 @@ Handles transformation between OpenAI-compatible format and ModelScope API forma API Reference: https://modelscope.cn/docs/model-service/API-Inference/intro """ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter from typing_extensions import override from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -29,6 +31,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = object +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) + class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): """ @@ -177,8 +182,10 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): ) # Check for errors in response - if "error" in response_data: - error_msg: Final = response_data["error"].get("message", str(response_data["error"])) + response_object: Final = _JSON_OBJECT.validate_python(response_data) + if "error" in response_object: + error: Final = _JSON_OBJECT.validate_python(response_object["error"]) + error_msg: Final = error.get("message", str(error)) raise self.get_error_class( error_message=f"ModelScope error: {error_msg}", status_code=raw_response.status_code, @@ -186,7 +193,7 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): ) # Extract images from response - data_list: Final = response_data.get("data", []) + data_list: Final = _JSON_OBJECTS.validate_python(response_object.get("data", [])) if not model_response.data: model_response.data = [] diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index 15cfdb6bece..c8e1a331c8c 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -1,7 +1,8 @@ from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import Final, Literal import httpx +from pydantic import ConfigDict, TypeAdapter from typing_extensions import Required, TypedDict import litellm @@ -44,6 +45,14 @@ class NvidiaNimRerankResponse(TypedDict): rankings: Required[list[NvidiaNimRankingResult]] +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_NUMBER: Final[TypeAdapter[bool | int | float]] = TypeAdapter( + bool | int | float, config=ConfigDict(strict=True, hide_input_in_errors=True) +) +_BILLED_UNITS: Final = TypeAdapter(RerankBilledUnits) +_STR: Final = TypeAdapter(str) + + class NvidiaNimRerankConfig(BaseRerankConfig): """ Reference: https://docs.api.nvidia.com/nim/reference/nvidia-llama-3_2-nv-rerankqa-1b-v2-infer @@ -115,7 +124,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, @@ -133,7 +142,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): Nvidia NIM specific params (passed through as-is from non_default_params): - truncate: How to truncate input if too long (NONE, END) """ - optional_nvidia_nim_rerank_params: Final[dict[str, Any]] = { + optional_nvidia_nim_rerank_params: Final[dict[str, object]] = { "query": query, "documents": documents, } @@ -327,15 +336,18 @@ class NvidiaNimRerankConfig(BaseRerankConfig): # Construct metadata with billed_units # Nvidia NIM uses "usage" field with "total_tokens" - usage: Final = raw_response_json.get("usage", {}) - total_tokens: Final = usage.get("total_tokens", 0) + payload: Final = _JSON_OBJECT.validate_python(raw_response_json) + usage: Final = _JSON_OBJECT.validate_python(payload.get("usage", {})) + total_tokens: Final = _NUMBER.validate_python(usage.get("total_tokens", 0)) - billed_units: Final[RerankBilledUnits] = {"total_tokens": total_tokens if total_tokens > 0 else len(results)} + billed_units: Final = _BILLED_UNITS.validate_python( + {"total_tokens": total_tokens if total_tokens > 0 else len(results)} + ) meta: Final[RerankResponseMeta] = {"billed_units": billed_units} return RerankResponse( - id=raw_response_json.get("id") or str(uuid.uuid4()), + id=_STR.validate_python(payload.get("id") or str(uuid.uuid4())), results=results, meta=meta, ) diff --git a/litellm/llms/nvidia_riva/audio_transcription/transformation.py b/litellm/llms/nvidia_riva/audio_transcription/transformation.py index cc78f5b99ef..262dfb7bd78 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/transformation.py +++ b/litellm/llms/nvidia_riva/audio_transcription/transformation.py @@ -102,7 +102,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): if endpointing_config is not None: recognition_config["endpointing_config"] = endpointing_config - request_payload: Final[dict[str, Any]] = { + request_payload: Final[dict[str, object]] = { "recognition_config": recognition_config, "response_format": optional_params.get("response_format") or "json", "timestamp_granularities": optional_params.get("timestamp_granularities"), @@ -135,7 +135,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # gRPC auth is constructed in the handler, not via HTTP headers. return headers - def _build_recognition_config_dict(self, model: str, optional_params: dict) -> dict[str, Any]: + def _build_recognition_config_dict(self, model: str, optional_params: dict) -> dict[str, object]: """ Build the Riva ``RecognitionConfig`` shape as a plain dict. @@ -159,7 +159,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): "profanity_filter": optional_params.get("profanity_filter", False), } - def _build_endpointing_config_dict(self, optional_params: dict) -> dict[str, Any] | None: + def _build_endpointing_config_dict(self, optional_params: dict) -> dict[str, object] | None: """ Translate an OpenAI-style ``chunking_strategy`` into Riva's ``EndpointingConfig`` shape, or pass through an explicit @@ -177,7 +177,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return None if isinstance(chunking, dict) and chunking.get("type") == "server_vad": - config: Final[dict[str, Any]] = {} + config: Final[dict[str, object]] = {} if "threshold" in chunking: threshold: Final = float(chunking["threshold"]) config["start_threshold"] = threshold @@ -245,7 +245,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): response["task"] = "transcribe" if response_format == "verbose_json": - words: Final[list[dict[str, Any]]] = [] + words: Final[list[dict[str, object]]] = [] if timestamp_granularities and "word" in timestamp_granularities: for item in final_results: for word in item.get("words", []) or []: diff --git a/litellm/llms/oci/embed/transformation.py b/litellm/llms/oci/embed/transformation.py index dec43717387..edadc293771 100644 --- a/litellm/llms/oci/embed/transformation.py +++ b/litellm/llms/oci/embed/transformation.py @@ -22,9 +22,11 @@ Supported models: Reference: https://docs.oracle.com/en-us/iaas/api/#/en/generative-ai-inference/latest/EmbedTextResult/EmbedText """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -52,6 +54,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + # OCI sends up to 96 texts per embedText request (Cohere limit). OCI_EMBED_BATCH_LIMIT: Final = 96 @@ -275,7 +279,7 @@ class OCIEmbedConfig(BaseEmbeddingConfig): ) try: - parsed: Final = OCIEmbedResponse(**json_response) + parsed: Final = OCIEmbedResponse.model_validate(_JSON_OBJECT.validate_python(json_response)) except Exception as e: raise OCIError( status_code=500, diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 66782a83a19..0d7a553d633 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -5,6 +5,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint. """ import json +from collections.abc import Sequence from typing import Any, Final import httpx @@ -26,11 +27,11 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): async def handle_count_tokens_request( self, model: str, - input: str | list[Any], + input: str | Sequence[object], api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - tools: list[dict[str, Any]] | None = None, + tools: list[dict[str, object]] | None = None, instructions: str | None = None, ) -> dict[str, Any]: """ diff --git a/litellm/llms/opencode/harness/transformation.py b/litellm/llms/opencode/harness/transformation.py index af5fa1ae71d..158af0fae0a 100644 --- a/litellm/llms/opencode/harness/transformation.py +++ b/litellm/llms/opencode/harness/transformation.py @@ -165,7 +165,7 @@ def _as_dict(value: object) -> Mapping[str, Any]: return value if isinstance(value, dict) else MappingProxyType({}) -def _tool_events(part: Mapping[str, Any]) -> Sequence[Event]: +def _tool_events(part: Mapping[str, object]) -> Sequence[Event]: native = str(part.get("tool") or "") call_id = str(part.get("callID") or part.get("id") or "") state = _as_dict(part.get("state")) @@ -195,7 +195,7 @@ def _error_message(error: object) -> str: return str(error.get("name") or "opencode reported an error") -def validate_user_config(config: Mapping[str, Any]) -> None: +def validate_user_config(config: Mapping[str, object]) -> None: """Reject OpenCodeOptions.config keys LiteLLM manages (or that bypass permissions).""" for key in config: if key in MANAGED_CONFIG_KEYS: @@ -241,7 +241,7 @@ def build_opencode_config( user_config: Mapping[str, Any] | None = None, instructions_path: str | None = None, skills_path: str | None = None, -) -> Mapping[str, Any]: +) -> Mapping[str, object]: """The full opencode config: user config underneath, LiteLLM-managed keys on top.""" user: Final = user_config or MappingProxyType({}) validate_user_config(user) @@ -379,7 +379,7 @@ class OpenCodeHarnessConfig(BaseCLIHarnessConfig): def create_stream_state(self) -> OpenCodeStreamState: return OpenCodeStreamState() - def transform_stream_line(self, line: Mapping[str, Any], state: OpenCodeStreamState) -> Sequence[Event]: + def transform_stream_line(self, line: Mapping[str, object], state: OpenCodeStreamState) -> Sequence[Event]: """step_finish token counts are ignored on purpose: the session endpoint accounts usage.""" session_id = line.get("sessionID") if session_id and state.session_id is None: diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 67d90d027ec..cfca853e280 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -27,9 +27,11 @@ Response format: } """ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -55,6 +57,11 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_STR: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True)) + class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): """ @@ -355,15 +362,16 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): model_response.data = [] try: - choices: Final = response_json.get("choices", []) + response_object: Final = _JSON_DICT.validate_python(response_json) + choices: Final = _JSON_OBJECTS.validate_python(response_object.get("choices", [])) for choice in choices: - message = choice.get("message", {}) - images = message.get("images", []) + message = _JSON_OBJECT.validate_python(choice.get("message", {})) + images = _JSON_OBJECTS.validate_python(message.get("images", [])) for image_data in images: - image_url_obj = image_data.get("image_url", {}) - image_url = image_url_obj.get("url") + image_url_obj = _JSON_OBJECT.validate_python(image_data.get("image_url", {})) + image_url = _STR.validate_python(image_url_obj.get("url") or "") if image_url: if image_url.startswith("data:"): @@ -389,7 +397,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ) # Extract and set usage and cost information - self._set_usage_and_cost(model_response, response_json, model) + self._set_usage_and_cost(model_response, response_object, model) return model_response diff --git a/litellm/llms/ovhcloud/audio_transcription/transformation.py b/litellm/llms/ovhcloud/audio_transcription/transformation.py index 6fc56ebb61f..086c1afbcd4 100644 --- a/litellm/llms/ovhcloud/audio_transcription/transformation.py +++ b/litellm/llms/ovhcloud/audio_transcription/transformation.py @@ -5,9 +5,11 @@ Our unified API follows the OpenAI standard. More information on our website: https://endpoints.ai.cloud.ovh.net """ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.llms.base_llm.audio_transcription.transformation import ( @@ -24,6 +26,8 @@ from litellm.types.utils import FileTypes, TranscriptionResponse from ..utils import OVHCloudException +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): def get_supported_openai_params(self, model: str) -> list[OpenAIAudioTranscriptionOptionalParams]: @@ -145,7 +149,8 @@ class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): headers=raw_response.headers, ) - text: Final = response_json.get("text") or response_json.get("transcript") or "" + payload: Final = _JSON_OBJECT.validate_python(response_json) + text: Final = payload.get("text") or payload.get("transcript") or "" response: Final = TranscriptionResponse(text=text) # OVHCloud field migration (deadline: 2026-05-11): @@ -153,9 +158,7 @@ class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # Prefer `seconds`, fall back to `duration`, normalize to `duration` # so downstream consumers see a consistent key. duration: Final = ( - response_json["seconds"] - if "seconds" in response_json and response_json["seconds"] is not None - else response_json.get("duration") + payload["seconds"] if "seconds" in payload and payload["seconds"] is not None else payload.get("duration") ) if duration is not None: response_json["duration"] = duration diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py index f65bf1e7292..0fcfafb79fb 100644 --- a/litellm/llms/recraft/image_generation/transformation.py +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -21,6 +23,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) + class RecraftImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://external.api.recraft.ai" @@ -141,7 +146,8 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): if not model_response.data: model_response.data = [] - for image_data in response_data["data"]: + payload: Final = _JSON_OBJECT.validate_python(response_data) + for image_data in _JSON_OBJECTS.validate_python(payload["data"]): model_response.data.append( ImageObject( url=image_data.get("url", None), diff --git a/litellm/llms/replicate/chat/transformation.py b/litellm/llms/replicate/chat/transformation.py index f7e09b7bec0..e8dd8816997 100644 --- a/litellm/llms/replicate/chat/transformation.py +++ b/litellm/llms/replicate/chat/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Iterable from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH @@ -26,6 +28,8 @@ if TYPE_CHECKING: else: LoggingClass = Any +_TEXTS: Final = TypeAdapter(Iterable[str], config=ConfigDict(strict=True, hide_input_in_errors=True)) + class ReplicateConfig(BaseConfig): """ @@ -253,7 +257,7 @@ class ReplicateConfig(BaseConfig): message=f"LiteLLM Error - prediction not succeeded - {raw_response_json}", headers=raw_response.headers, ) - outputs: Final = raw_response_json.get("output", []) + outputs: Final = _TEXTS.validate_python(raw_response_json.get("output", [])) response_str = "".join(outputs) if len(response_str) == 0: # edge case, where result from replicate is empty response_str = " " diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 8e8f7ea61aa..0d86752fc5f 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -1,9 +1,10 @@ import functools import json -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm import verbose_logger @@ -11,6 +12,8 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.types.utils import StreamingChatCompletionChunk +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + def _load_sagemaker_response_stream_shape(): try: @@ -18,7 +21,7 @@ def _load_sagemaker_response_stream_shape(): from botocore.model import ServiceModel loader: Final = Loader() - service_dict: Final = loader.load_service_model("sagemaker-runtime", "service-2") + service_dict: Final = _JSON_OBJECT.validate_python(loader.load_service_model("sagemaker-runtime", "service-2")) return ServiceModel(service_dict).shape_for("InvokeEndpointWithResponseStreamOutput") except Exception as e: verbose_logger.warning( diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py index 75f35a8f379..6aa66de1db5 100644 --- a/litellm/llms/snowflake/embedding/transformation.py +++ b/litellm/llms/snowflake/embedding/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -10,6 +12,8 @@ from litellm.types.utils import EmbeddingResponse from ..utils import SnowflakeBaseConfig, SnowflakeException +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): """ @@ -53,7 +57,7 @@ class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): # convert embeddings to 1d array for item in response_json["data"]: item["embedding"] = item["embedding"][0] - returned_response: Final = EmbeddingResponse(**response_json) + returned_response: Final = EmbeddingResponse.model_validate(_JSON_OBJECT.validate_python(response_json)) returned_response.model = "snowflake/" + (returned_response.model or "") diff --git a/litellm/llms/stability/image_generation/transformation.py b/litellm/llms/stability/image_generation/transformation.py index 656ffe395c8..1856523ce49 100644 --- a/litellm/llms/stability/image_generation/transformation.py +++ b/litellm/llms/stability/image_generation/transformation.py @@ -6,9 +6,11 @@ Handles transformation between OpenAI-compatible format and Stability AI API for API Reference: https://platform.stability.ai/docs/api-reference """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -33,6 +35,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class StabilityImageGenerationConfig(BaseImageGenerationConfig): """ @@ -234,7 +238,8 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): ) # Check finish_reason - finish_reason: Final = response_data.get("finish_reason", "") + payload: Final = _JSON_OBJECT.validate_python(response_data) + finish_reason: Final = payload.get("finish_reason", "") if finish_reason == "CONTENT_FILTERED": raise self.get_error_class( error_message="Content was filtered by Stability AI safety systems", @@ -246,7 +251,7 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): model_response.data = [] # Extract image from response - image_b64: Final = response_data.get("image") + image_b64: Final = payload.get("image") if image_b64: model_response.data.append( ImageObject( diff --git a/litellm/llms/tinyfish/search/transformation.py b/litellm/llms/tinyfish/search/transformation.py index d5ed7da3815..c7976456d87 100644 --- a/litellm/llms/tinyfish/search/transformation.py +++ b/litellm/llms/tinyfish/search/transformation.py @@ -26,6 +26,7 @@ from litellm.secret_managers.main import get_secret_str _UrlEncodableParams: Final = TypeAdapter(dict[str, str | int | float | bool]) _StrList: Final = TypeAdapter(list[str]) _StrFrozenSet: Final = TypeAdapter(frozenset[str]) +_DecodedJson: Final = TypeAdapter(object) _TINYFISH_PARAMS_KEY: Final = "_tinyfish_params" _TINYFISH_DOCS_URL: Final = "https://docs.tinyfish.ai/search-api" @@ -218,7 +219,7 @@ class TinyfishSearchConfig(BaseSearchConfig): ) try: - raw_json: Final[object] = raw_response.json() # any-ok: httpx Response.json() -> Any + raw_json: Final = _DecodedJson.validate_python(raw_response.json()) except json.JSONDecodeError: raise self._wrap_error( error_message=f"Expected JSON response, got: {raw_response.text[:200]}", @@ -276,7 +277,7 @@ class TinyfishSearchConfig(BaseSearchConfig): # for other envelope shapes (CDN HTML pages, other JSON envelopes, plain text). inner_message = error_message try: - body: Final[object] = json.loads(error_message) # any-ok: json.loads -> Any + body: Final = _DecodedJson.validate_python(json.loads(error_message)) if isinstance(body, dict): error_obj: Final[object] = body.get("error") # any-ok: untyped dict if isinstance(error_obj, dict): diff --git a/litellm/llms/together_ai/rerank/handler.py b/litellm/llms/together_ai/rerank/handler.py index b8079e52c97..2fff0854f10 100644 --- a/litellm/llms/together_ai/rerank/handler.py +++ b/litellm/llms/together_ai/rerank/handler.py @@ -4,7 +4,9 @@ Re rank api LiteLLM supports the re rank API format, no paramter transformation occurs """ -from typing import Any, Final +from typing import Final + +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.llms.base import BaseLLM @@ -15,6 +17,8 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.llms.together_ai.rerank.transformation import TogetherAIRerankConfig from litellm.types.rerank import RerankRequest, RerankResponse +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) + def _rerank_url(api_base: str) -> str: return f"{api_base.rstrip('/')}/rerank" @@ -27,7 +31,7 @@ class TogetherAIRerank(BaseLLM): api_key: str, api_base: str, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], top_n: int | None = None, rank_fields: list[str] | None = None, return_documents: bool | None = True, @@ -66,13 +70,13 @@ class TogetherAIRerank(BaseLLM): if response.status_code != 200: raise Exception(response.text) - _json_response: Final = response.json() + _json_response: Final = _JSON_DICT.validate_python(response.json()) return TogetherAIRerankConfig()._transform_response(_json_response) async def async_rerank( # New async method self, - request_data_dict: dict[str, Any], + request_data_dict: dict[str, object], api_key: str, api_base: str, ) -> RerankResponse: @@ -91,6 +95,6 @@ class TogetherAIRerank(BaseLLM): if response.status_code != 200: raise Exception(response.text) - _json_response: Final = response.json() + _json_response: Final = _JSON_DICT.validate_python(response.json()) return TogetherAIRerankConfig()._transform_response(_json_response) diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index 6a5bb484540..7ba7b060991 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -1,8 +1,10 @@ import json +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx from openai.types.image import Image +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.llms.custom_httpx.http_handler import ( @@ -17,11 +19,13 @@ from litellm.types.utils import ImageResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +_PREDICTIONS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) + class VertexImageGeneration(VertexLLM): def process_image_generation_response( self, - json_response: dict[str, Any], + json_response: Mapping[str, object], model_response: ImageResponse, model: str | None = None, ) -> ImageResponse: @@ -32,12 +36,11 @@ class VertexImageGeneration(VertexLLM): model=model, ) - predictions: Final = json_response["predictions"] + predictions: Final = _PREDICTIONS.validate_python(json_response["predictions"]) response_data: Final[list[Image]] = [] for prediction in predictions: - bytes_base64_encoded = prediction["bytesBase64Encoded"] - image_object = Image(b64_json=bytes_base64_encoded) + image_object = Image.model_validate({"b64_json": prediction["bytesBase64Encoded"]}) response_data.append(image_object) model_response.data = response_data diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 17ccf16837b..9f3126fd4ff 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -1,7 +1,9 @@ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -31,6 +33,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): """ @@ -321,10 +326,11 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): ) ) - if usage_metadata := response_data.get("usageMetadata", None): - model_response.usage = self._transform_image_usage(usage_metadata) + response_object: Final = _JSON_OBJECT.validate_python(response_data) + if usage_metadata := response_object.get("usageMetadata", None): + model_response.usage = self._transform_image_usage(_JSON_DICT.validate_python(usage_metadata)) - web_search_requests: Final = get_gemini_image_web_search_requests(response_data) + web_search_requests: Final = get_gemini_image_web_search_requests(response_object) if web_search_requests and model_response.usage is not None: setattr(model_response.usage, "web_search_requests", web_search_requests) diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index a7b079fb89c..a78145fdcf4 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -6,11 +6,12 @@ Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/s """ import base64 -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, Union import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.exceptions import UnsupportedParamsError @@ -45,6 +46,9 @@ else: _LyriaVoice: TypeAlias = str | dict | None +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_STR: Final = TypeAdapter(str, config=ConfigDict(hide_input_in_errors=True)) + class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): """ @@ -465,14 +469,14 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): from litellm.types.llms.openai import HttpxBinaryResponseContent # Parse JSON response - _json_response: Final = raw_response.json() + _json_response: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Get base64-encoded audio content response_content: Final = _json_response.get("audioContent") if not response_content: raise ValueError("No audioContent in Vertex AI TTS response") - binary_data: Final = base64.b64decode(response_content) + binary_data: Final = base64.b64decode(_STR.validate_python(response_content)) media_type: Final = speech_media_type_from_audio_bytes(binary_data) response: Final = httpx.Response( status_code=200, diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index b48efe229b3..3acb2f2ed58 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -4,10 +4,11 @@ Transformation logic for Voyage AI's /v1/rerank endpoint. Docs - https://docs.voyageai.com/docs/reranker """ -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj @@ -23,6 +24,11 @@ from litellm.types.utils import ModelInfo from ..embedding.transformation import VoyageError +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_OPTIONAL_INT: Final[TypeAdapter[int | None]] = TypeAdapter(int | None) +_STR: Final = TypeAdapter(str) + class VoyageRerankConfig(BaseRerankConfig): def get_supported_cohere_rerank_params(self, model: str) -> list: @@ -103,13 +109,14 @@ class VoyageRerankConfig(BaseRerankConfig): ) # Voyage AI returns results in "data" key, not "results" - _results: Final[list[dict] | None] = _json_response.get("data") + payload: Final = _JSON_OBJECT.validate_python(_json_response) + _results: Final = payload.get("data") if _results is None: raise ValueError(f"No results found in the response={_json_response}") # Transform to LiteLLM format transformed_results: Final = [] - for result in _results: + for result in _JSON_OBJECTS.validate_python(_results): transformed_result: dict[str, object] = { "index": result["index"], "relevance_score": result["relevance_score"], @@ -121,14 +128,14 @@ class VoyageRerankConfig(BaseRerankConfig): transformed_result["document"] = result["document"] transformed_results.append(transformed_result) - usage: Final = _json_response.get("usage", {}) - total_tokens: Final = usage.get("total_tokens", 0) + usage: Final = _JSON_OBJECT.validate_python(payload.get("usage", {})) + total_tokens: Final = _OPTIONAL_INT.validate_python(usage.get("total_tokens", 0)) _billed_units: Final = RerankBilledUnits(total_tokens=total_tokens) _tokens: Final = RerankTokens(input_tokens=total_tokens, output_tokens=0) rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) return RerankResponse( - id=_json_response.get("id") or str(uuid.uuid4()), + id=_STR.validate_python(payload.get("id") or str(uuid.uuid4())), results=transformed_results, meta=rerank_meta, ) diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 2169b9bf49a..77ef03b8e10 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -4,9 +4,11 @@ Translates from OpenAI's `/v1/audio/transcriptions` to IBM WatsonX's `/ml/v1/aud WatsonX follows the OpenAI spec for audio transcription. """ +from collections.abc import Mapping from typing import Final from httpx import Response +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file @@ -25,6 +27,8 @@ from ...openai.transcriptions.whisper_transformation import ( ) from ..common_utils import IBMWatsonXMixin +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTranscriptionConfig): """ @@ -174,8 +178,9 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran # Extract only valid fields for TranscriptionResponse.__init__() # TranscriptionResponse only accepts 'text' and 'usage' in __init__() - text: Final = raw_response_json.get("text") - usage: Final = raw_response_json.get("usage") + response_object: Final = _JSON_OBJECT.validate_python(raw_response_json) + text: Final = response_object.get("text") + usage: Final = response_object.get("usage") # Create response with only valid fields response_kwargs: Final = {} @@ -187,14 +192,14 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran if not response_kwargs: raise ValueError( "Invalid response format. Received response does not match the expected format. Got: ", - raw_response_json, + response_object, ) response: Final = TranscriptionResponse(**response_kwargs) # Add other fields using dictionary-style assignment (like duration, task, etc.) # Skip fields that TranscriptionResponse doesn't accept in __init__() - for key, value in raw_response_json.items(): + for key, value in response_object.items(): if key not in [ "text", "usage", diff --git a/litellm/llms/watsonx/embed/transformation.py b/litellm/llms/watsonx/embed/transformation.py index 80cd28a058d..609eeba90c0 100644 --- a/litellm/llms/watsonx/embed/transformation.py +++ b/litellm/llms/watsonx/embed/transformation.py @@ -2,9 +2,11 @@ Translates from OpenAI's `/v1/embeddings` to IBM's `/text/embeddings` route. """ +from collections.abc import Iterable, Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.embedding.transformation import ( BaseEmbeddingConfig, @@ -16,6 +18,9 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import IBMWatsonXMixin, _get_api_params +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_TOKEN_COUNT: Final = TypeAdapter(int) + class IBMWatsonXEmbeddingConfig(IBMWatsonXMixin, BaseEmbeddingConfig): def get_supported_openai_params(self, model: str) -> list: @@ -95,7 +100,7 @@ class IBMWatsonXEmbeddingConfig(IBMWatsonXMixin, BaseEmbeddingConfig): json_resp: Final = raw_response.json() if model_response is None: model_response = EmbeddingResponse(model=json_resp.get("model_id", None)) - results: Final = json_resp.get("results", []) + results: Final = _JSON_OBJECTS.validate_python(json_resp.get("results", [])) embedding_response: Final = [] for idx, result in enumerate(results): embedding_response.append( @@ -107,7 +112,7 @@ class IBMWatsonXEmbeddingConfig(IBMWatsonXMixin, BaseEmbeddingConfig): ) model_response.object = "list" model_response.data = embedding_response - input_tokens: Final = json_resp.get("input_token_count", 0) + input_tokens: Final = _TOKEN_COUNT.validate_python(json_resp.get("input_token_count", 0) or 0) setattr( model_response, "usage", diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index ff95f14951a..c2cfd6597f3 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -5,10 +5,11 @@ Docs - https://cloud.ibm.com/apidocs/watsonx-ai#text-rerank """ import uuid -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from typing import Final, cast import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -24,6 +25,11 @@ from litellm.types.rerank import ( from ..common_utils import IBMWatsonXMixin, _generate_watsonx_token, _get_api_params +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_OPTIONAL_INT: Final[TypeAdapter[int | None]] = TypeAdapter(int | None) +_STR: Final = TypeAdapter(str) + class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): """ @@ -171,13 +177,14 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): headers=raw_response.headers, ) - _results: Final[list[dict] | None] = raw_response_json.get("results") + payload: Final = _JSON_OBJECT.validate_python(raw_response_json) + _results: Final = payload.get("results") if _results is None: raise ValueError(f"No results found in the response={raw_response_json}") transformed_results: Final = [] - for result in _results: + for result in _JSON_OBJECTS.validate_python(_results): transformed_result: dict[str, object] = { "index": result["index"], "relevance_score": result["score"], @@ -191,11 +198,11 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): transformed_results.append(transformed_result) - response_id: Final = raw_response_json.get("id") or str(uuid.uuid4()) + response_id: Final = _STR.validate_python(payload.get("id") or str(uuid.uuid4())) # Extract usage information _tokens: Final = RerankTokens( - input_tokens=raw_response_json.get("input_token_count", 0), + input_tokens=_OPTIONAL_INT.validate_python(payload.get("input_token_count", 0)), ) rerank_meta: Final = RerankResponseMeta(tokens=_tokens) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index b85ff2ce50e..a2aeef3f67f 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -8,7 +8,7 @@ from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast from fastapi import HTTPException -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger @@ -162,6 +162,8 @@ _CLIENT_FORWARDED_AUTH_TYPES: Final["frozenset[str]"] = frozenset({"true_passthr # Minted token material that must never survive a client rotation on a persisted row. _MINTED_TOKEN_CREDENTIAL_FIELDS: Final["frozenset[str]"] = frozenset({"access_token", "refresh_token", "expires_in"}) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + def _bind_submitted_oauth_client( credentials: Mapping[str, object], issuer: str | None, url: str | None @@ -1916,7 +1918,7 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]: creds: Final = getattr(server, "credentials", None) if isinstance(creds, str): try: - parsed: dict[str, object] | None = json.loads(creds) + parsed: Mapping[str, object] | None = _JSON_OBJECT.validate_python(json.loads(creds)) except ValueError: parsed = None else: @@ -2382,13 +2384,10 @@ def _decode_user_env_vars(stored: str) -> dict[str, str]: "re-enter them rather than silently forwarding ciphertext" ) return {} - parsed: dict[str, object] | None try: - parsed = json.loads(decrypted) + parsed: Final = _JSON_OBJECT.validate_python(json.loads(decrypted)) except (ValueError, TypeError): return {} - if not isinstance(parsed, dict): - return {} return {str(k): str(v) for k, v in parsed.items()} diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index c88c6f2570a..dca48babc81 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -19,7 +19,7 @@ from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, StreamingResponse -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -78,6 +78,9 @@ _PASCAL_TO_WIRE: Final[Mapping[str, str]] = { } +_DECODED_JSON: Final = TypeAdapter(object) + + def _sse_event(payload: object) -> str: """Frame a JSON-RPC object as a single A2A SSE event (``data: \\n\\n``).""" return f"data: {json.dumps(payload)}\n\n" @@ -91,7 +94,7 @@ def _to_jsonrpc_object(chunk: object) -> object: """ if isinstance(chunk, (str, bytes, bytearray)): try: - return json.loads(chunk) + return _DECODED_JSON.validate_python(json.loads(chunk)) except (json.JSONDecodeError, UnicodeDecodeError): return chunk if hasattr(chunk, "model_dump"): diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index a330d057490..cf48bee6958 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -3,9 +3,14 @@ from collections.abc import Iterator from typing import Any, Final import requests +from pydantic import ConfigDict, TypeAdapter from .exceptions import UnauthorizedError +_SSE_LINE: Final[TypeAdapter[bytes | bytearray]] = TypeAdapter( + bytes | bytearray, config=ConfigDict(arbitrary_types_allowed=True, strict=True, hide_input_in_errors=True) +) + class ChatClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 600): @@ -172,7 +177,7 @@ class ChatClient: # Parse SSE stream for line in response.iter_lines(): if line: - line = line.decode("utf-8") + line = _SSE_LINE.validate_python(line).decode("utf-8") if line.startswith("data: "): data_str = line[6:] # Remove 'data: ' prefix if data_str.strip() == "[DONE]": diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 98af32fa7aa..684449acd0c 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -8,6 +8,7 @@ from urllib.parse import urlencode import click import requests +from pydantic import ConfigDict, TypeAdapter from rich.console import Console from rich.table import Table from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never @@ -121,6 +122,8 @@ class CliAuthResult(TypedDict): _TeamMapping: Final = TypeVar("_TeamMapping", bound=Mapping[str, object]) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" KEYRING_ENABLE_HINT: Final = "keyring --enable (or unset PYTHON_KEYRING_BACKEND)" @@ -485,7 +488,7 @@ def prompt_team_selection_fallback( def _response_error_detail(response: requests.Response) -> str | None: try: - body: Final[dict[str, object] | list[object] | str | int | float | bool | None] = response.json() + body: Final = _JSON_OBJECT.validate_python(response.json()) except ValueError: return None detail: Final = body.get("detail") if isinstance(body, dict) else None diff --git a/litellm/proxy/client/cli/commands/debug.py b/litellm/proxy/client/cli/commands/debug.py index 4e3914143d8..2aca60ca810 100644 --- a/litellm/proxy/client/cli/commands/debug.py +++ b/litellm/proxy/client/cli/commands/debug.py @@ -115,6 +115,7 @@ class RequestResponsePayload(BaseModel): _SESSION_PAGE: Final = TypeAdapter(SessionLogsPage) _PAYLOAD: Final[TypeAdapter[RequestResponsePayload | None]] = TypeAdapter(RequestResponsePayload | None) _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_BACKTICK_RUNS: Final = TypeAdapter(tuple[str, ...]) _SESSION_PAGE_SIZE: Final = 100 _TRANSPORT_BODY_CHARS: Final = 500 @@ -192,7 +193,7 @@ def _fmt_json(value: JsonValue, max_chars: int) -> str: def _fenced(text: str, info: str = "") -> tuple[str, str, str]: - longest_run: Final = max((len(run) for run in re.findall(r"`+", text)), default=0) + longest_run: Final = max((len(run) for run in _BACKTICK_RUNS.validate_python(re.findall(r"`+", text))), default=0) fence: Final = "`" * max(3, longest_run + 1) return (f"{fence}{info}", text, fence) diff --git a/litellm/proxy/db/model_insights_tasks.py b/litellm/proxy/db/model_insights_tasks.py index 865965dcf75..e117afb881e 100644 --- a/litellm/proxy/db/model_insights_tasks.py +++ b/litellm/proxy/db/model_insights_tasks.py @@ -1,14 +1,20 @@ import json +from collections.abc import Mapping from functools import lru_cache from pathlib import Path from typing import Final +from pydantic import ConfigDict, TypeAdapter + from litellm.types.model_insights import ModelInsightTask _TASKS_FILE: Final = Path(__file__).resolve().parent.parent / "model_insights_tasks.json" +_TASK_ENTRIES: Final = TypeAdapter( + Mapping[str, Mapping[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True) +) @lru_cache(maxsize=1) def load_model_insight_tasks() -> dict[str, ModelInsightTask]: - raw: Final = json.loads(_TASKS_FILE.read_text()) - return {name: ModelInsightTask(task_type=name, **entry) for name, entry in raw.items()} + raw: Final = _TASK_ENTRIES.validate_python(json.loads(_TASKS_FILE.read_text())) + return {name: ModelInsightTask.model_validate(dict(task_type=name, **entry)) for name, entry in raw.items()} diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index a13ad00713d..e09f1ec8ba8 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -426,7 +426,7 @@ async def list_fine_tuning_jobs( route_type=CallTypes.alist_fine_tuning_jobs.value, ) - response: Any | None = None + response: object = None if target_model_names and isinstance(target_model_names, str): target_model_names_list: Final = target_model_names.split(",") if len(target_model_names_list) != 1: diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 00f0ef1756f..2a0b528a50b 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -8,7 +8,7 @@ import time import traceback from collections.abc import Iterable, Mapping from datetime import datetime, timedelta, timezone -from typing import Any, Final, Literal, TypedDict, cast +from typing import Final, Literal, TypedDict import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status @@ -1431,7 +1431,7 @@ async def shared_health_check_status_endpoint( ) -def _read_license_data() -> dict[str, Any] | None: +def _read_license_data() -> EnterpriseLicenseData | None: from litellm.proxy.proxy_server import _license_check, premium_user_data license_data: EnterpriseLicenseData | None = premium_user_data or _license_check.airgapped_license_data @@ -1453,10 +1453,10 @@ def _read_license_data() -> dict[str, Any] | None: if license_data is None: return None - return cast(dict[str, Any], license_data) + return license_data -def _read_allowed_features(license_data: dict[str, Any]) -> list: +def _read_allowed_features(license_data: Mapping[str, object]) -> list: raw_allowed_features: Final = license_data.get("allowed_features") if isinstance(raw_allowed_features, list): return list(raw_allowed_features) @@ -1707,7 +1707,7 @@ def _show_env_credential_login_warning() -> bool: async def _get_health_readiness_details( response: Response | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Detailed health payload for authenticated diagnostics. """ @@ -1726,7 +1726,7 @@ async def _get_health_readiness_details( success_callback_names = litellm.success_callback # check Cache - cache_type: Any = None + cache_type: object = None if litellm.cache is not None: from litellm.caching.caching import RedisSemanticCache @@ -1735,7 +1735,7 @@ async def _get_health_readiness_details( if isinstance(litellm.cache.cache, RedisSemanticCache): # ping the cache # TODO: @ishaan-jaff - we should probably not ping the cache on every /health/readiness check - index_info: Any + index_info: object try: index_info = await litellm.cache.cache._index_info() except Exception as e: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 9806a824c63..5db0fce7d92 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -54,6 +54,8 @@ from litellm.repositories.autorouter_session_repository import AutoRouterSession from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.daily_activity_sql import build_where_clause from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository from litellm.router_strategy.complexity_router import ComplexityRouter from litellm.router_utils.auto_router_model_naming import ( StrategyRouterDependencyRole, @@ -187,15 +189,15 @@ def _team_table(prisma_client: "PrismaClient") -> _TeamTable: def _verification_tokens(prisma_client: "PrismaClient") -> _VerificationTokenTable: - return prisma_client.db.litellm_verificationtoken + return VerificationTokenRepository(prisma_client).table def _team_rows(prisma_client: "PrismaClient") -> _TeamRowsTable: - return prisma_client.db.litellm_teamtable + return TeamRepository(prisma_client).table def _user_rows(prisma_client: "PrismaClient") -> _UserRowsTable: - return prisma_client.db.litellm_usertable + return UserRepository(prisma_client).table def _shadow_eval_jobs(prisma_client: "PrismaClient") -> _ShadowEvalJobTable: diff --git a/litellm/proxy/management_endpoints/callback_management_endpoints.py b/litellm/proxy/management_endpoints/callback_management_endpoints.py index 4f9f46af08f..d910b475b57 100644 --- a/litellm/proxy/management_endpoints/callback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/callback_management_endpoints.py @@ -7,12 +7,15 @@ import os from typing import Final from fastapi import APIRouter, Depends +from pydantic import TypeAdapter from litellm.litellm_core_utils.logging_callback_manager import CallbacksByType from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router: Final = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) + @router.get( "/callbacks/list", @@ -51,6 +54,6 @@ async def get_callback_configs(): ) with open(config_path, "r") as f: - configs: Final = json.load(f) + configs: Final = _DECODED_JSON.validate_python(json.load(f)) return configs diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 323b9e434a8..037b9082cd4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -355,9 +355,12 @@ _KEY_METADATA_REQUEST_FIELDS: Final = frozenset( ) +_DECODED_JSON: Final = TypeAdapter(object) + + def _decode_json_string_column(column: str, value: object) -> object: if column in _KEY_UPDATE_JSON_STRING_COLUMNS and isinstance(value, str): - return json.loads(value) + return _DECODED_JSON.validate_python(json.loads(value)) return value diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py index 106dbfaf7b7..f1760b857c5 100644 --- a/litellm/proxy/management_endpoints/management_v1/budgets.py +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -85,8 +85,7 @@ class PrismaBudgetListExecutor: async def count(self, where: tuple[Predicate, ...]) -> int: clauses, params = where_sql(where) sql: Final = f"SELECT COUNT(*) AS count FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "") - rows: Final = await self.prisma_client.db.query_raw(sql, *params) - counted: Final = _ROW_COUNTS.validate_python(rows) + counted: Final = _ROW_COUNTS.validate_python(await self.prisma_client.db.query_raw(sql, *params)) return counted[0].count if counted else 0 async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]: @@ -97,8 +96,7 @@ class PrismaBudgetListExecutor: + f" ORDER BY {order_by_sql(plan.order)}" + f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" ) - rows: Final = await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip) - return _BUDGET_ROWS.validate_python(rows) + return _BUDGET_ROWS.validate_python(await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip)) def _serialize(row: BudgetListItem) -> BudgetListItem: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 2aed0e2afb8..e5bea71ac7c 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -533,6 +533,8 @@ if MCP_AVAILABLE: except Exception as e: verbose_proxy_logger.debug("Failed to write temporary MCP server to Redis cache: %s", e) + _CACHED_VALUE: Final = TypeAdapter(object) + @with_service_target(MCP_SERVERS_TARGET) async def _get_temporary_mcp_server_from_redis( server_id: str, @@ -550,8 +552,8 @@ if MCP_AVAILABLE: return None try: - cached_server: Final = await cache_backend.async_get_cache( - key=f"{TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX}:{server_id}" + cached_server: Final = _CACHED_VALUE.validate_python( + await cache_backend.async_get_cache(key=f"{TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX}:{server_id}") ) except Exception as e: verbose_proxy_logger.debug("Failed reading temporary MCP server from Redis cache: %s", e) diff --git a/litellm/proxy/management_endpoints/model_insights_endpoints.py b/litellm/proxy/management_endpoints/model_insights_endpoints.py index b787aac2d9f..53e8dfcc71e 100644 --- a/litellm/proxy/management_endpoints/model_insights_endpoints.py +++ b/litellm/proxy/management_endpoints/model_insights_endpoints.py @@ -115,7 +115,7 @@ def _deployment_filter(rows: list[_GroupedModel]) -> list[dict[str, str]]: def _daily_metric(row: _GroupedDaily) -> ModelInsightDailyMetric: - return ModelInsightDailyMetric(date=row.date, **_metric(row).model_dump()) + return ModelInsightDailyMetric.model_validate({**_metric(row).model_dump(), "date": row.date}) def _daily_total(row: _GroupedDate) -> ModelInsightDailyTotal: @@ -146,12 +146,14 @@ def _summarize_tasks(rows: list[_GroupedTask], metric: ModelInsightsMetric) -> l } grand: Final = sum(totals.values()) return [ - ModelInsightTaskSummary( - **(catalog.get(task) or _UNCATEGORIZED_TASK).model_copy(update={"task_type": task}).model_dump(), - value=value, - share=value / grand * 100 if grand else 0.0, - leader=leaders[task].model_group, - provider=leaders[task].custom_llm_provider, + ModelInsightTaskSummary.model_validate( + { + **(catalog.get(task) or _UNCATEGORIZED_TASK).model_copy(update={"task_type": task}).model_dump(), + "value": value, + "share": value / grand * 100 if grand else 0.0, + "leader": leaders[task].model_group, + "provider": leaders[task].custom_llm_provider, + } ) for task, value in sorted(totals.items(), key=lambda item: item[1], reverse=True) ] diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 11a0075ffd0..253013a2ebf 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -2547,8 +2547,8 @@ async def add_new_model( enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)), ) - clean_model_info: Final = ModelInfo( - **without_server_derived_pricing(model_params.model_info.model_dump(exclude_none=True)) + clean_model_info: Final = ModelInfo.model_validate( + dict(without_server_derived_pricing(model_params.model_info.model_dump(exclude_none=True))) ) model_params.model_info = ( # rebind-ok: downstream team-model handling mutates this same object clean_model_info.model_copy(update=MappingProxyType({"member_auto_router": True})) @@ -3194,6 +3194,9 @@ def _deduplicate_litellm_router_models(models: list[dict]) -> list[dict]: return unique_models +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + + def model_info_as_mapping(model_info: object) -> Mapping[str, object] | None: """A DB row's model_info column arrives as a dict or as its JSON string depending on the query path, and every consumer needs the mapping. Single owner of that parse: @@ -3204,10 +3207,9 @@ def model_info_as_mapping(model_info: object) -> Mapping[str, object] | None: if not isinstance(model_info, str): return None try: - parsed: Final = json.loads(model_info) + return _JSON_OBJECT.validate_python(json.loads(model_info)) except (TypeError, ValueError): return None - return parsed if isinstance(parsed, Mapping) else None def _expects_liveness_on_this_pod(model_info: object) -> bool: diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index c8ae7af41db..26abe35b49a 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -27,7 +27,7 @@ from typing import ( import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, status -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -306,7 +306,7 @@ async def _verify_org_access( ) -_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _BUDGET_SETTABLE_FIELDS: Final = frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"} _ORG_COLUMN_FIELDS: Final = frozenset({"organization_alias", "models"}) _ORG_METADATA_FIELDS: Final = tuple( @@ -720,7 +720,7 @@ async def update_organization( ) # Transform UI payload to expected format - raw_data: Final[dict[str, object]] = await request.json() + raw_data: Final = _STR_OBJECT_DICT_ADAPTER.validate_python(await request.json()) raw_data_with_flat_budget_fields: Final = handle_nested_budget_structure_in_organization_update_request(raw_data) # Create validated data model @@ -767,7 +767,9 @@ async def update_organization( # Merge metadata from existing organization with updated metadata if updated_organization_row_json.get("metadata") is not None: existing_metadata: Final = existing_organization_row.metadata or {} - updated_metadata: Final[dict[str, object]] = updated_organization_row_json.get("metadata", {}) + updated_metadata: Final = _STR_OBJECT_DICT_ADAPTER.validate_python( + updated_organization_row_json.get("metadata", {}) + ) merged_metadata: Final[Mapping[str, object]] = _update_dictionary( existing_dict=cast( # cast-ok: prisma de-serializes a Json column to the plain python dict it stores "dict[str, object]", existing_metadata @@ -786,12 +788,16 @@ async def update_organization( ) budget_fields: Final = { - k: v for k, v in data.model_dump().items() if k in _BUDGET_SETTABLE_FIELDS and k in data.model_fields_set + k: v + for k, v in _STR_OBJECT_DICT_ADAPTER.validate_python(data.model_dump()).items() + if k in _BUDGET_SETTABLE_FIELDS and k in data.model_fields_set } if budget_fields and existing_organization_row.budget_id: await update_budget( - budget_obj=BudgetNewRequest(budget_id=existing_organization_row.budget_id, **budget_fields), + budget_obj=BudgetNewRequest.model_validate( + {"budget_id": existing_organization_row.budget_id, **budget_fields} + ), user_api_key_dict=user_api_key_dict, ) diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 69356922ea1..c3c499eeec1 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -12,12 +12,12 @@ All /policy management endpoints import copy import json import os -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Mapping from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import Response, StreamingResponse -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TypedDict import litellm @@ -449,6 +449,23 @@ async def validate_policy( return result +class _LoadedPolicy(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, hide_input_in_errors=True) + + inherit: str | None = None + scope: PolicyScopeResponse = Field(default_factory=PolicyScopeResponse) + guardrails: PolicyGuardrailsResponse = Field(default_factory=PolicyGuardrailsResponse) + resolved_guardrails: tuple[str, ...] = () + inheritance_chain: tuple[str, ...] = () + + +class _LoadedPolicies(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, hide_input_in_errors=True) + + policies: Mapping[str, _LoadedPolicy] = Field(default_factory=dict) + total_count: int = 0 + + @router.get( "/policy/list", tags=["policy management"], @@ -472,19 +489,19 @@ async def list_policies( """ from litellm.proxy.policy_engine.init_policies import get_policies_summary - summary: Final = get_policies_summary() + summary: Final = _LoadedPolicies.model_validate(get_policies_summary()) return PolicyListResponse( policies={ name: PolicySummaryItem( - inherit=data.get("inherit"), - scope=PolicyScopeResponse(**data.get("scope", {})), - guardrails=PolicyGuardrailsResponse(**data.get("guardrails", {})), - resolved_guardrails=data.get("resolved_guardrails", []), - inheritance_chain=data.get("inheritance_chain", []), + inherit=policy.inherit, + scope=policy.scope, + guardrails=policy.guardrails, + resolved_guardrails=list(policy.resolved_guardrails), + inheritance_chain=list(policy.inheritance_chain), ) - for name, data in summary.get("policies", {}).items() + for name, policy in summary.policies.items() }, - total_count=summary.get("total_count", 0), + total_count=summary.total_count, ) diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index a557e3a6082..d8ea96e6a46 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -116,7 +116,7 @@ async def get_router_settings( config: Final = await proxy_config.get_config() router_settings_from_config: Final = config.get("router_settings", {}) - current_values: Final[dict[str, Any]] = {} + current_values: Final[dict[str, object]] = {} if llm_router is not None: # Router exposes routing groups as private `_routing_groups`; the # generic `hasattr` loop below would miss them. diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index b1094684389..9a1cf701142 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -17,6 +17,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, TypedDict, overload from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view @@ -58,6 +59,8 @@ if TYPE_CHECKING: router: Final = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) + class _TagRecord(Protocol): tag_name: str @@ -523,7 +526,7 @@ async def info_tag( model_info: object = {} if tag_record.model_info: if isinstance(tag_record.model_info, str): - model_info = json.loads(tag_record.model_info) + model_info = _DECODED_JSON.validate_python(json.loads(tag_record.model_info)) else: model_info = tag_record.model_info @@ -646,7 +649,7 @@ async def list_tags( model_info: object = {} if tag_record.model_info: if isinstance(tag_record.model_info, str): - model_info = json.loads(tag_record.model_info) + model_info = _DECODED_JSON.validate_python(json.loads(tag_record.model_info)) else: model_info = tag_record.model_info diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fe976c861e5..e452afccef9 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1721,7 +1721,7 @@ async def new_team( complete_team_data_dict = complete_team_data.model_dump(exclude_none=True) # Serialize router_settings to JSON (matching key creation pattern) - router_settings_value: Final = getattr(data, "router_settings", None) + router_settings_value: Final = data.router_settings router_settings_json: Final = ( safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({}) ) diff --git a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py index 2381a5cc2db..3d408accba3 100644 --- a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py +++ b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py @@ -18,11 +18,11 @@ class FileContentStreamingHandler: *, custom_llm_provider: str, file_id: str, - data: dict[str, Any], + data: dict[str, object], should_route: bool, original_file_id: str | None, - credentials: dict[str, Any] | None, - ) -> tuple[str, str, dict[str, Any]]: + credentials: dict[str, object] | None, + ) -> tuple[str, str, dict[str, object]]: """ Resolve the provider, file ID, and request payload to use for streaming. diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index f5e62da5962..dfec77a6c16 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1652,7 +1652,8 @@ async def delete_file( def _as_file_list_page(response: object) -> object: if not isinstance(response, list): return response - return FileListPage(**build_list_page(_LISTED_FILES_ADAPTER.validate_python(response))) + page: Final = build_list_page(_LISTED_FILES_ADAPTER.validate_python(response)) + return FileListPage.model_validate(page) @router.get( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 7cec3bac207..dfb731972ac 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -5,6 +5,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -48,6 +49,8 @@ else: PassThroughEndpointLogging = Any EndpointType = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class AnthropicPassthroughLoggingHandler: @staticmethod @@ -343,11 +346,9 @@ class AnthropicPassthroughLoggingHandler: if not line.startswith("data:"): continue try: - data = json.loads(line[len("data:") :].strip()) + data = _JSON_OBJECT.validate_python(json.loads(line[len("data:") :].strip())) except (json.JSONDecodeError, ValueError): continue - if not isinstance(data, dict): - continue etype = data.get("type") if etype == "message_delta": return False diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 24c1b865db6..be40e225cde 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import urlparse import httpx -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -55,6 +55,7 @@ else: _VERTEX_INTERACTIONS_PATH: Final = re.compile(r"/projects/[^/]+/locations/[^/]+/interactions/?$") _INTERACTIONS_RESPONSE_BODY: Final = TypeAdapter(dict[str, object]) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) def _interactions_model( @@ -99,7 +100,7 @@ class VertexPassthroughLoggingHandler: litellm_model_response: Final = ModelResponse( model=model, usage=InteractionsUsageObjectTransformation.transform_interactions_usage_object( - cast(Mapping[str, Any], usage_object) + _JSON_OBJECT.validate_python(usage_object) ), ) logging_obj.custom_llm_provider = custom_llm_provider @@ -349,7 +350,7 @@ class VertexPassthroughLoggingHandler: model: Final = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) - _json_response: Final[dict[str, object]] = httpx_response.json() + _json_response: Final = _JSON_OBJECT.validate_python(httpx_response.json()) litellm_prediction_response: ModelResponse | EmbeddingResponse | ImageResponse = ModelResponse() if VertexPassthroughLoggingHandler._is_audio_predict_response( diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index f55eb4f7863..efe7a5002e0 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -421,7 +421,7 @@ class PolicyRegistry: if policy_request.condition is not None: data["condition"] = json.dumps(policy_request.condition.model_dump()) if policy_request.pipeline is not None: - validated_pipeline: Final = GuardrailPipeline(**policy_request.pipeline) + validated_pipeline: Final = GuardrailPipeline.model_validate(policy_request.pipeline) data["pipeline"] = json.dumps(validated_pipeline.model_dump()) created_policy: Final = await _policy_table(prisma_client).create(data=data) @@ -496,7 +496,7 @@ class PolicyRegistry: if policy_request.condition is not None: update_data["condition"] = json.dumps(policy_request.condition.model_dump()) if policy_request.pipeline is not None: - validated_pipeline: Final = GuardrailPipeline(**policy_request.pipeline) + validated_pipeline: Final = GuardrailPipeline.model_validate(policy_request.pipeline) update_data["pipeline"] = json.dumps(validated_pipeline.model_dump()) updated_policy: Final = await _policy_table(prisma_client).update( diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 7814903e975..da9a3033187 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -496,10 +496,11 @@ _AUTOROUTER_PRESETS_ADAPTER: Final = TypeAdapter(dict[str, AutoRouterPresetRecor def _load_bundled_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord]: - raw: Final = json.loads( - files("litellm.proxy.public_endpoints").joinpath("autorouter_presets.json").read_text(encoding="utf-8") + return _AUTOROUTER_PRESETS_ADAPTER.validate_python( + json.loads( + files("litellm.proxy.public_endpoints").joinpath("autorouter_presets.json").read_text(encoding="utf-8") + ) ) - return _AUTOROUTER_PRESETS_ADAPTER.validate_python(raw) async def _fetch_remote_autorouter_presets(url: str) -> Mapping[str, AutoRouterPresetRecord]: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 2dd1c9419f1..f9235a33b42 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -11,6 +11,7 @@ from types import MappingProxyType from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast from fastapi import HTTPException, status +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._internal_context import with_service_target @@ -71,6 +72,9 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "Project": Litellm_EntityType.PROJECT.value, } +_CACHED_VALUE: Final = TypeAdapter(object) +_WINDOW_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class _CounterReservationUnavailable(Exception): def __init__( @@ -762,8 +766,8 @@ async def _get_team_member_budget_counter( else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): - default_budget: Final = await user_api_key_cache.async_get_cache( - key=f"team_member_default_budget:{default_budget_id}", + default_budget: Final = _CACHED_VALUE.validate_python( + await user_api_key_cache.async_get_cache(key=f"team_member_default_budget:{default_budget_id}") ) default_cap: Final = _to_float(_get_value(default_budget, "max_budget")) if default_cap is not None and default_cap > 0: @@ -799,8 +803,8 @@ async def _get_org_budget_counter( if org_id is None: return None - org_table: Final = await user_api_key_cache.async_get_cache( - key=f"org_id:{org_id}:with_budget", + org_table: Final = _CACHED_VALUE.validate_python( + await user_api_key_cache.async_get_cache(key=f"org_id:{org_id}:with_budget") ) if org_table is None: return None @@ -833,7 +837,9 @@ async def _get_project_budget_counter( return None source_cache_key: Final = project_cache_key(valid_token.project_id) - project_object: Final = await user_api_key_cache.async_get_cache(key=source_cache_key) + project_object: Final = _CACHED_VALUE.validate_python( + await user_api_key_cache.async_get_cache(key=source_cache_key) + ) if project_object is None: return None @@ -901,10 +907,9 @@ def _coerce_window(window: object) -> Mapping[str, object]: return window if isinstance(window, str): try: - parsed: Final[object] = json.loads(window) + return _WINDOW_OBJECT.validate_python(json.loads(window)) except Exception: return {} - return parsed if isinstance(parsed, Mapping) else {} model_dump: Final = getattr(window, "model_dump", None) if not callable(model_dump): return {} diff --git a/litellm/proxy/spend_tracking/spend_log_error_logger.py b/litellm/proxy/spend_tracking/spend_log_error_logger.py index 03037ba854f..28cd535127f 100644 --- a/litellm/proxy/spend_tracking/spend_log_error_logger.py +++ b/litellm/proxy/spend_tracking/spend_log_error_logger.py @@ -25,7 +25,7 @@ troubleshoot. The UI suppression follows the same gate. import logging import os -from typing import Any, Final +from typing import Final from litellm._logging import verbose_proxy_logger from litellm.secret_managers.main import str_to_bool @@ -58,7 +58,7 @@ def should_suppress_spend_log_tracebacks() -> bool: def spend_log_error( message: str, - *args: Any, + *args: object, exc: BaseException | None = None, ) -> None: """Log a spend-tracking error, with the traceback gated on the env var. diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index da1bc0a1feb..fb799546b37 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -23,6 +23,7 @@ from litellm.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE from litellm.litellm_core_utils.url_utils import async_safe_get from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, + header_value, httpxSpecialProvider, ) from litellm.rag.ingestion.file_parsers import extract_text_from_pdf @@ -79,7 +80,7 @@ class BaseRAGIngestion(ABC): from litellm.litellm_core_utils.credential_accessor import CredentialAccessor credential_name: Final = self.vector_store_config.get("litellm_credential_name") - if credential_name and litellm.credential_list: + if isinstance(credential_name, str) and credential_name and litellm.credential_list: credential_values: Final = CredentialAccessor.get_credential_values(credential_name) if not credential_values: return @@ -125,7 +126,8 @@ class BaseRAGIngestion(ABC): response.raise_for_status() file_content = response.content filename = file_url.split("/")[-1] or "document" - content_type = response.headers.get("content-type", "application/octet-stream") + content_type_header: Final = header_value(response.headers, "content-type") + content_type = "application/octet-stream" if content_type_header is None else content_type_header return filename, file_content, content_type, None if file_id: diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index 1cf5db549e4..a015f3cf614 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -9,6 +9,8 @@ from __future__ import annotations from typing import TYPE_CHECKING, Final, cast +from pydantic import BaseModel, ConfigDict + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -23,6 +25,13 @@ if TYPE_CHECKING: from litellm.types.rag import RAGIngestOptions +class _WhiteSpaceConfig(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, hide_input_in_errors=True) + + max_tokens_per_chunk: object = 800 + max_overlap_tokens: object = 400 + + class GeminiRAGIngestion(BaseRAGIngestion): """ Gemini-specific RAG ingestion using File Search API. @@ -234,12 +243,13 @@ class GeminiRAGIngestion(BaseRAGIngestion): # Add chunking configuration if provided chunking_strategy: Final = self.chunking_strategy if chunking_strategy and isinstance(chunking_strategy, dict): - white_space_config: Final = chunking_strategy.get("white_space_config") + white_space_config: Final[object] = chunking_strategy.get("white_space_config") if white_space_config: + white_space: Final = _WhiteSpaceConfig.model_validate(white_space_config) request_body["chunkingConfig"] = { "whiteSpaceConfig": { - "maxTokensPerChunk": white_space_config.get("max_tokens_per_chunk", 800), - "maxOverlapTokens": white_space_config.get("max_overlap_tokens", 400), + "maxTokensPerChunk": white_space.max_tokens_per_chunk, + "maxOverlapTokens": white_space.max_overlap_tokens, } } diff --git a/litellm/rag/ingestion/openai_ingestion.py b/litellm/rag/ingestion/openai_ingestion.py index 925a49d0f7a..6c7c619237d 100644 --- a/litellm/rag/ingestion/openai_ingestion.py +++ b/litellm/rag/ingestion/openai_ingestion.py @@ -7,7 +7,7 @@ so this implementation skips the embedding step and directly uploads files. from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Final import litellm from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion @@ -106,7 +106,7 @@ class OpenAIRAGIngestion(BaseRAGIngestion): vector_store_id=vector_store_id, file_id=existing_file_id, custom_llm_provider="openai", - chunking_strategy=cast(dict[str, Any] | None, self.chunking_strategy), + chunking_strategy=self.chunking_strategy, api_key=api_key, api_base=api_base, ) @@ -134,7 +134,7 @@ class OpenAIRAGIngestion(BaseRAGIngestion): vector_store_id=vector_store_id, file_id=result_file_id, custom_llm_provider="openai", - chunking_strategy=cast(dict[str, Any] | None, self.chunking_strategy), + chunking_strategy=self.chunking_strategy, api_key=api_key, api_base=api_base, ) diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index e2aa5555eec..e3eafd680de 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -18,7 +18,7 @@ from __future__ import annotations import hashlib import uuid from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, TypedDict +from typing import TYPE_CHECKING, Final, TypedDict import litellm from litellm._logging import verbose_logger @@ -194,7 +194,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): url: str, data: str | None = None, headers: dict[str, str] | None = None, - ) -> Any: + ) -> httpx.Response: """ Helper to sign and execute AWS API requests using httpx + SigV4. diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 1a5301f0579..5ab39bec080 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -18,6 +18,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._internal_context import is_internal_call @@ -71,6 +72,13 @@ _SEARCH_ARGS_SET_BY_PIPELINE: Final = frozenset( ) +_VECTOR_STORE_OPTIONS: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(hide_input_in_errors=True)) + + +def _vector_store_provider(ingest_options: Mapping[str, object]) -> object: + return _VECTOR_STORE_OPTIONS.validate_python(ingest_options.get("vector_store", {})).get("custom_llm_provider") + + def get_ingestion_class(provider: str) -> type[BaseRAGIngestion]: """ Get the ingestion class for a given provider. @@ -137,7 +145,7 @@ async def _execute_ingest_pipeline( @client async def aingest( - ingest_options: dict[str, Any], + ingest_options: Mapping[str, object], file_data: tuple[str, bytes, str] | None = None, file: dict[str, str] | None = None, file_url: str | None = None, @@ -197,7 +205,7 @@ async def aingest( except Exception as e: raise litellm.exception_type( model=None, - custom_llm_provider=ingest_options.get("vector_store", {}).get("custom_llm_provider"), + custom_llm_provider=_vector_store_provider(ingest_options), original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, @@ -226,7 +234,7 @@ def _suppressed_sub_call_billing() -> Iterator[None]: async def _execute_query_pipeline( model: str, messages: list[AllMessageValues], - retrieval_config: dict[str, Any], + retrieval_config: Mapping[str, object], rerank: dict[str, Any] | None = None, stream: bool = False, vector_store_params: Mapping[str, object] | None = None, @@ -276,12 +284,16 @@ async def _execute_query_pipeline( search_provider: Final = retrieval_config.get("custom_llm_provider", "openai") try: - search_cost = sum( - vector_store_search_cost( - model=search_provider if "/" in search_provider else None, - custom_llm_provider=search_provider, - response=search_response, + search_cost = ( + sum( + vector_store_search_cost( + model=search_provider if "/" in search_provider else None, + custom_llm_provider=search_provider, + response=search_response, + ) ) + if isinstance(search_provider, str) + else 0.0 ) except Exception: # noqa: BLE001 - cost accounting must never break the query path search_cost = 0.0 @@ -354,7 +366,7 @@ async def _execute_query_pipeline( async def aquery( model: str, messages: list[AllMessageValues], - retrieval_config: dict[str, Any], + retrieval_config: Mapping[str, object], rerank: dict[str, Any] | None = None, stream: bool = False, vector_store_params: Mapping[str, object] | None = None, @@ -403,7 +415,7 @@ async def aquery( def query( model: str, messages: list[AllMessageValues], - retrieval_config: dict[str, Any], + retrieval_config: Mapping[str, object], rerank: dict[str, Any] | None = None, stream: bool = False, vector_store_params: Mapping[str, object] | None = None, @@ -450,7 +462,7 @@ def query( @client def ingest( - ingest_options: dict[str, Any], + ingest_options: Mapping[str, object], file_data: tuple[str, bytes, str] | None = None, file: dict[str, str] | None = None, file_url: str | None = None, @@ -517,7 +529,7 @@ def ingest( except Exception as e: raise litellm.exception_type( model=None, - custom_llm_provider=ingest_options.get("vector_store", {}).get("custom_llm_provider"), + custom_llm_provider=_vector_store_provider(ingest_options), original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, diff --git a/litellm/repositories/autorouter_session_repository.py b/litellm/repositories/autorouter_session_repository.py index 82e2c091728..640fd26ffe3 100644 --- a/litellm/repositories/autorouter_session_repository.py +++ b/litellm/repositories/autorouter_session_repository.py @@ -2,7 +2,7 @@ Repository for the auto-router per-session rollup (LiteLLM_AutoRouterSession). """ -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol from litellm.models.autorouter_session import LiteLLM_AutoRouterSession from litellm.repositories.base_repository import BaseRepository @@ -12,10 +12,21 @@ if TYPE_CHECKING: from prisma import models as prisma_models +class _AutoRouterSessionDb(Protocol): + @property + def litellm_autoroutersession(self) -> TableActions["prisma_models.LiteLLM_AutoRouterSession"]: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _AutoRouterSessionDb: ... + + class AutoRouterSessionRepository(BaseRepository[LiteLLM_AutoRouterSession]): @property def table(self) -> TableActions["prisma_models.LiteLLM_AutoRouterSession"]: - return self.prisma_client.db.litellm_autoroutersession + client: Final[_PrismaClientView] = self.prisma_client + return client.db.litellm_autoroutersession @property def model_class(self) -> type[LiteLLM_AutoRouterSession]: diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index cbacb6e8f90..6541db4c00b 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -4,20 +4,24 @@ Model repository for database operations on LiteLLM_ProxyModelTable. import json from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final + +from pydantic import ConfigDict, TypeAdapter from litellm.models.model import LiteLLM_ProxyModelTable from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) -from litellm.repositories.base_repository import BaseRepository +from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) + class _ProxyModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_ProxyModelTable"]): table_name = "litellm_proxymodeltable" @@ -60,22 +64,26 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): decrypted[key] = value return decrypted - def _to_model(self, record: Any) -> LiteLLM_ProxyModelTable | None: + def _to_model(self, record: DbRecord | None) -> LiteLLM_ProxyModelTable | None: """Convert a database record to a Model with decryption.""" if record is None: return None - data: Final = record.dict() if hasattr(record, "dict") else dict(record) + data: Final = dict(record_to_dict(record)) - if isinstance(data.get("litellm_params"), str): - data["litellm_params"] = json.loads(data["litellm_params"]) - if isinstance(data.get("model_info"), str): - data["model_info"] = json.loads(data["model_info"]) + litellm_params: Final = data.get("litellm_params") + if isinstance(litellm_params, str): + data["litellm_params"] = json.loads(litellm_params) + model_info: Final = data.get("model_info") + if isinstance(model_info, str): + data["model_info"] = json.loads(model_info) if data.get("litellm_params"): - data["litellm_params"] = self._decrypt_litellm_params(data["litellm_params"]) + data["litellm_params"] = self._decrypt_litellm_params( + _LITELLM_PARAMS.validate_python(data["litellm_params"]) + ) - return LiteLLM_ProxyModelTable(**data) + return LiteLLM_ProxyModelTable.model_validate(data) async def find_by_id(self, model_id: str, id_field: str = "model_id") -> LiteLLM_ProxyModelTable | None: return await super().find_by_id(model_id, id_field) diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py index b732d2ff94c..99d6aeed88b 100644 --- a/litellm/repositories/object_permission_repository.py +++ b/litellm/repositories/object_permission_repository.py @@ -2,7 +2,7 @@ ObjectPermission repository for database operations on LiteLLM_ObjectPermissionTable. """ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.repositories.base_repository import BaseRepository @@ -12,6 +12,19 @@ if TYPE_CHECKING: from prisma import models as prisma_models +class _ObjectPermissionDb(Protocol): + @property + def litellm_objectpermissiontable(self) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _ObjectPermissionDb: ... + + @property + def writer_db(self) -> _ObjectPermissionDb: ... + + class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): """Repository for object permission database operations.""" @@ -21,7 +34,8 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]: - database: Final = self.prisma_client.writer_db if self._use_writer else self.prisma_client.db + client: Final[_PrismaClientView] = self.prisma_client + database: Final = client.writer_db if self._use_writer else client.db return database.litellm_objectpermissiontable @property @@ -48,7 +62,7 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): skills: list[str] | None = None, ) -> LiteLLM_ObjectPermissionTable: """Create a new object permission record.""" - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} if mcp_servers is not None: data["mcp_servers"] = mcp_servers if mcp_access_groups is not None: @@ -90,7 +104,7 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): skills: list[str] | None = None, ) -> LiteLLM_ObjectPermissionTable | None: """Update an object permission record.""" - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} if mcp_servers is not None: data["mcp_servers"] = mcp_servers if mcp_access_groups is not None: diff --git a/litellm/repositories/project_repository.py b/litellm/repositories/project_repository.py index 905e813f35e..0ea5009a5b6 100644 --- a/litellm/repositories/project_repository.py +++ b/litellm/repositories/project_repository.py @@ -3,7 +3,7 @@ Project repository for database operations on LiteLLM_ProjectTable. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol from litellm.models.project import LiteLLM_ProjectTable from litellm.repositories.base_repository import BaseRepository @@ -13,12 +13,23 @@ if TYPE_CHECKING: from prisma import models as prisma_models +class _ProjectDb(Protocol): + @property + def litellm_projecttable(self) -> TableActions["prisma_models.LiteLLM_ProjectTable"]: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _ProjectDb: ... + + class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): """Repository for project database operations.""" @property def table(self) -> TableActions["prisma_models.LiteLLM_ProjectTable"]: - return self.prisma_client.db.litellm_projecttable + client: Final[_PrismaClientView] = self.prisma_client + return client.db.litellm_projecttable @property def model_class(self) -> type[LiteLLM_ProjectTable]: diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index b201dbc566b..5ed1574c6bc 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -5,7 +5,7 @@ User repository for database operations on LiteLLM_UserTable. import json from collections.abc import Mapping, Sequence from itertools import chain -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter @@ -37,6 +37,19 @@ ORDER BY p.user_id _PLACEHOLDER_ROWS_ADAPTER: Final = TypeAdapter(tuple[SCIMPlaceholder, ...]) +class _UserDb(Protocol): + @property + def litellm_usertable(self) -> TableActions["prisma_models.LiteLLM_UserTable"]: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _UserDb: ... + + @property + def writer_db(self) -> _UserDb: ... + + class UserRepository(BaseRepository[LiteLLM_UserTable]): """Repository for user database operations.""" @@ -46,7 +59,8 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_UserTable"]: - database: Final = self.prisma_client.writer_db if self._use_writer else self.prisma_client.db + client: Final[_PrismaClientView] = self.prisma_client + database: Final = client.writer_db if self._use_writer else client.db return database.litellm_usertable @property diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 9aa881fc4c9..95c981502e7 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -123,6 +123,9 @@ class AdaptiveRouter: prefs = self.model_to_prefs.get(model) or _default_prefs() self._cells[(rt, model)] = initial_cell(prefs, rt) + def cell(self, request_type: RequestType, model: str) -> BanditCell: + return self._cells[(request_type, model)] + async def load_state_from_db(self, prisma_client: object) -> None: """Add each row's persisted delta to a freshly computed cold-start prior. diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index fc0865654c7..43200271f9f 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2945,7 +2945,7 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"No candidate models left for tier {tier_key} after routing-plugin filtering") return self._pick_from_tier_value(context.candidate_models, tier_key) - def _ensure_adaptive_router(self) -> Any | None: + def _ensure_adaptive_router(self) -> AdaptiveRouter | None: if not self.config.adaptive: return None if self.adaptive_router is not None: @@ -3087,7 +3087,7 @@ class ComplexityRouter(CustomLogger): pools: Final = self._tier_pools() classified_candidates: Final = _allowed(tuple(pools.get(_tier_name(classified_tier), ())), fit_filter) cold_start_candidates: Final = tuple( - model for model in classified_candidates if adaptive._cells[(request_type, model)].total_samples == 0 + model for model in classified_candidates if adaptive.cell(request_type, model).total_samples == 0 ) if cold_start_candidates: chosen_model: Final = random.choice(cold_start_candidates) @@ -3106,7 +3106,7 @@ class ComplexityRouter(CustomLogger): "candidates": [ { "model": model, - "total_samples": adaptive._cells[(request_type, model)].total_samples, + "total_samples": adaptive.cell(request_type, model).total_samples, } for model in cold_start_candidates ], @@ -3123,7 +3123,7 @@ class ComplexityRouter(CustomLogger): best_score = float("-inf") candidate_scores: Final[list[dict[str, object]]] = [] for model in self._adaptive_candidate_models(classified_tier, hard_floor, hard_ceiling, fit_filter): - cell = adaptive._cells[(request_type, model)] + cell = adaptive.cell(request_type, model) quality_sample = thompson_sample(cell) cost_score = normalized_cost(adaptive.model_to_cost.get(model, 0.0), all_costs) if self.config.adaptive_eligible == "classified_tier": diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index 1d7656f253e..200de29b664 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -34,7 +34,7 @@ from typing import Final, get_args import litellm from litellm.types.llms.openai import REASONING_EFFORT -REASONING_EFFORT_ADVERTISEMENT_ORDER: Final = get_args(REASONING_EFFORT) +REASONING_EFFORT_ADVERTISEMENT_ORDER: Final[tuple[str, ...]] = get_args(REASONING_EFFORT) _EMPTY_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) _EFFORT_FLAGS: Final = ( diff --git a/litellm/secret_managers/aws_secret_manager.py b/litellm/secret_managers/aws_secret_manager.py index 1aed44e9a31..b3a9fbf31c8 100644 --- a/litellm/secret_managers/aws_secret_manager.py +++ b/litellm/secret_managers/aws_secret_manager.py @@ -14,9 +14,13 @@ import os import re from typing import Any, Final +from pydantic import TypeAdapter + import litellm from litellm.proxy._types import KeyManagementSystem +_PARSED_LITERAL: Final = TypeAdapter(object) + def validate_environment(): if "AWS_REGION_NAME" not in os.environ: @@ -107,7 +111,7 @@ class AWSKeyManagementService_V2: if isinstance(secret, str): secret = secret.strip() try: - secret_value_as_bool: Final = ast.literal_eval(secret) + secret_value_as_bool: Final = _PARSED_LITERAL.validate_python(ast.literal_eval(secret)) if isinstance(secret_value_as_bool, bool): return secret_value_as_bool except Exception: diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index fbed4fb8e75..275c029b0b7 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -6,7 +6,7 @@ import traceback from typing import Final import httpx -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger @@ -19,6 +19,8 @@ from litellm.secret_managers.get_azure_ad_token_provider import ( oidc_cache: Final = DualCache() +_PARSED_LITERAL: Final = TypeAdapter(object) + _OIDC_TOKEN_EXPIRY_MARGIN_SECONDS: Final = 60 @@ -348,7 +350,7 @@ def get_secret( secret = os.getenv(secret_name) try: if isinstance(secret, str): - secret_value_as_bool = ast.literal_eval(secret) + secret_value_as_bool = _PARSED_LITERAL.validate_python(ast.literal_eval(secret)) if isinstance(secret_value_as_bool, bool): return secret_value_as_bool else: diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index e78d2aa5f6a..c0cb3796934 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -10,6 +10,8 @@ from typing import ( get_args, ) +from pydantic import ConfigDict, TypeAdapter + from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import remove_items_at_indices from litellm.repositories.table_repositories import ( @@ -30,6 +32,8 @@ if TYPE_CHECKING: else: PrismaClient = Any +_DB_ROW: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) + class VectorStoreIndexRegistry: def __init__(self, vector_store_indexes: list[LiteLLM_ManagedVectorStoreIndex] = []): @@ -95,7 +99,9 @@ class VectorStoreIndexRegistry: ) for vector_store in _vector_stores_from_db: _dict_vector_store = dict(vector_store) - _litellm_managed_vector_store = LiteLLM_ManagedVectorStoreIndex(**_dict_vector_store) + _litellm_managed_vector_store = LiteLLM_ManagedVectorStoreIndex.model_validate( + _DB_ROW.validate_python(_dict_vector_store) + ) vector_stores_from_db.append(_litellm_managed_vector_store) return vector_stores_from_db diff --git a/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py index db561fd1dc2..8d5e15b4b71 100644 --- a/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py +++ b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py @@ -2,10 +2,15 @@ Tests for Pydantic AI agents header forwarding via agent_extra_headers. """ -from unittest.mock import AsyncMock, MagicMock, patch +from typing import Final +from unittest.mock import ANY, AsyncMock, MagicMock, patch +import httpx import pytest +import respx +import litellm +from litellm.a2a_protocol.providers.pydantic_ai_agents.config import PydanticAIProviderConfig from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( PydanticAITransformation, ) @@ -198,3 +203,67 @@ async def test_provider_config_threads_agent_extra_headers(): sent_headers = mock_client.post.await_args.kwargs["headers"] assert sent_headers["x-trace-id"] == "abc-123" assert sent_headers["Content-Type"] == "application/json" + + +COMPLETED_TASK: Final = { + "jsonrpc": "2.0", + "id": "req-5", + "result": { + "id": "task-5", + "kind": "task", + "status": {"state": "completed"}, + "history": [], + "artifacts": [{"artifactId": "a-5", "parts": [{"kind": "text", "text": "ok"}]}], + }, +} + + +def _user_params() -> dict[str, object]: + return {"message": {"role": "user", "parts": [{"kind": "text", "text": "hello"}], "messageId": "msg-user-5"}} + + +@pytest.fixture +def agent_route(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock.post("http://agent.test/").mock(return_value=httpx.Response(200, json=COMPLETED_TASK)) + + +@pytest.mark.asyncio +async def test_provider_config_non_streaming_defaults_to_sixty_second_timeout(agent_route: respx.Route) -> None: + response: Final = await PydanticAIProviderConfig().handle_non_streaming( + "req-5", _user_params(), "http://agent.test" + ) + + request: Final = agent_route.calls.last.request + assert response == { + "jsonrpc": "2.0", + "id": "req-5", + "result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": "ok"}], "messageId": ANY}, + } + assert request.extensions["timeout"] == {"connect": 60.0, "read": 60.0, "write": 60.0, "pool": 60.0} + assert request.headers["content-type"] == "application/json" + + +@pytest.mark.asyncio +async def test_provider_config_non_streaming_forwards_timeout_and_ignores_unrelated_keywords( + agent_route: respx.Route, +) -> None: + response: Final = await PydanticAIProviderConfig().handle_non_streaming( + request_id="req-5", + params=_user_params(), + api_base="http://agent.test", + timeout=5.5, + agent_extra_headers={"x-trace-id": "abc-123"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + + request: Final = agent_route.calls.last.request + assert response["id"] == "req-5" + assert request.extensions["timeout"] == {"connect": 5.5, "read": 5.5, "write": 5.5, "pool": 5.5} + assert request.headers["x-trace-id"] == "abc-123" + + +@pytest.mark.asyncio +async def test_provider_config_non_streaming_requires_api_base() -> None: + with pytest.raises(ValueError, match="api_base is required for PydanticAIProviderConfig"): + await PydanticAIProviderConfig().handle_non_streaming("req-5", _user_params()) diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py new file mode 100644 index 00000000000..38b896723f7 --- /dev/null +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py @@ -0,0 +1,161 @@ +import json +import re +from typing import Final +from unittest.mock import ANY + +import httpx +import pytest +import respx + +import litellm +from litellm.a2a_protocol.providers.watsonx_orchestrate.config import WatsonxOrchestrateA2AConfig +from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import WXOLitellmParams + +MISSING_LITELLM_PARAMS: Final = re.escape( + "litellm_params is required for WatsonxOrchestrateA2AConfig " + "(must contain cp4d_host, instance_id, wxo_agent_id, api_key)" +) +MISSING_HOST: Final = re.escape("'cp4d_host' is required in litellm_params for WXO agents") +RUNS_URL: Final = "https://wxo-config.test/orchestrate/cpd/instances/inst-1/v1/orchestrate/runs" +LITELLM_PARAMS: Final[WXOLitellmParams] = { + "cp4d_host": "https://wxo-config.test", + "instance_id": "inst-1", + "wxo_agent_id": "agent-1", + "api_key": "config-test-key", + "auth_mode": "ibm_cloud", +} +EXPECTED_RUN_BODY: Final = { + "agent_id": "agent-1", + "message": {"role": "user", "content": [{"response_type": "text", "text": "hello"}]}, +} + + +def _a2a_params() -> dict[str, object]: + return {"message": {"role": "user", "parts": [{"kind": "text", "text": "hello"}], "messageId": "m-1"}} + + +def _mock_wxo_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: + respx_mock.post("https://iam.cloud.ibm.com/identity/token").mock( + return_value=httpx.Response(200, json={"access_token": "tok", "expires_in": 0}) + ) + return respx_mock.post(url).mock( + return_value=httpx.Response(200, json={"status": "completed", "results": "wxo says hi"}) + ) + + +@pytest.fixture +def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +@pytest.mark.parametrize("litellm_params", [None, {}]) +async def test_handle_non_streaming_rejects_empty_litellm_params(litellm_params: WXOLitellmParams | None) -> None: + with pytest.raises(ValueError, match=MISSING_LITELLM_PARAMS): + await WatsonxOrchestrateA2AConfig().handle_non_streaming("req-1", _a2a_params(), litellm_params=litellm_params) + + +async def test_handle_non_streaming_requires_litellm_params_when_only_other_keywords_are_given() -> None: + with pytest.raises(ValueError, match=MISSING_LITELLM_PARAMS): + await WatsonxOrchestrateA2AConfig().handle_non_streaming( + "req-1", _a2a_params(), "https://ignored.test", agent_extra_headers={"x-tenant-id": "acme"} + ) + + +async def test_handle_non_streaming_hands_litellm_params_to_the_wxo_handler() -> None: + with pytest.raises(ValueError, match=MISSING_HOST): + await WatsonxOrchestrateA2AConfig().handle_non_streaming( + "req-1", _a2a_params(), litellm_params={"instance_id": "inst-1"} + ) + + +@pytest.mark.usefixtures("httpx_transport") +async def test_handle_non_streaming_runs_the_agent_and_ignores_unrelated_keywords( + respx_mock: respx.MockRouter, +) -> None: + runs_route: Final = _mock_wxo_route(respx_mock, RUNS_URL) + + response: Final = await WatsonxOrchestrateA2AConfig().handle_non_streaming( + request_id="req-1", + params=_a2a_params(), + api_base="https://ignored.test", + litellm_params=LITELLM_PARAMS, + agent_extra_headers={"x-tenant-id": "acme"}, + ) + + run_request: Final = runs_route.calls.last.request + assert response == { + "jsonrpc": "2.0", + "id": "req-1", + "result": { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": "wxo says hi"}], + "messageId": ANY, + }, + } + assert json.loads(run_request.content) == EXPECTED_RUN_BODY + assert run_request.headers["authorization"] == "Bearer tok" + assert "x-tenant-id" not in run_request.headers + + +@pytest.mark.parametrize("litellm_params", [None, {}]) +async def test_handle_streaming_rejects_empty_litellm_params(litellm_params: WXOLitellmParams | None) -> None: + stream: Final = WatsonxOrchestrateA2AConfig().handle_streaming( + "req-1", _a2a_params(), litellm_params=litellm_params + ) + + with pytest.raises(ValueError, match=MISSING_LITELLM_PARAMS): + await anext(stream) + + +async def test_handle_streaming_requires_litellm_params_when_only_other_keywords_are_given() -> None: + stream: Final = WatsonxOrchestrateA2AConfig().handle_streaming( + "req-1", _a2a_params(), "https://ignored.test", agent_extra_headers={"x-tenant-id": "acme"} + ) + + with pytest.raises(ValueError, match=MISSING_LITELLM_PARAMS): + await anext(stream) + + +async def test_handle_streaming_hands_litellm_params_to_the_wxo_handler() -> None: + stream: Final = WatsonxOrchestrateA2AConfig().handle_streaming( + "req-1", _a2a_params(), litellm_params={"instance_id": "inst-1"} + ) + + with pytest.raises(ValueError, match=MISSING_HOST): + await anext(stream) + + +@pytest.mark.usefixtures("httpx_transport") +async def test_handle_streaming_runs_the_agent_and_ignores_unrelated_keywords(respx_mock: respx.MockRouter) -> None: + stream_route: Final = _mock_wxo_route(respx_mock, f"{RUNS_URL}/stream") + + chunks: Final = [ + chunk + async for chunk in WatsonxOrchestrateA2AConfig().handle_streaming( + request_id="req-1", + params=_a2a_params(), + api_base="https://ignored.test", + litellm_params=LITELLM_PARAMS, + agent_extra_headers={"x-tenant-id": "acme"}, + ) + ] + + stream_request: Final = stream_route.calls.last.request + assert [chunk["id"] for chunk in chunks] == ["req-1", "req-1", "req-1", "req-1"] + assert chunks[2]["result"] == { + "contextId": ANY, + "kind": "artifact-update", + "taskId": ANY, + "artifact": {"artifactId": ANY, "parts": [{"kind": "text", "text": "wxo says hi"}]}, + } + assert chunks[3]["result"] == { + "contextId": ANY, + "final": True, + "kind": "status-update", + "status": {"state": "completed"}, + "taskId": ANY, + } + assert json.loads(stream_request.content) == EXPECTED_RUN_BODY + assert stream_request.headers["authorization"] == "Bearer tok" + assert "x-tenant-id" not in stream_request.headers diff --git a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 52e44ca5448..6cfdb75fd01 100644 --- a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -2,6 +2,8 @@ import asyncio import json import os import unittest.mock as mock +from types import SimpleNamespace +from typing import Final from unittest.mock import patch import pytest @@ -19,7 +21,12 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import ( ) from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER -from litellm.proxy._types import CallInfo, Litellm_EntityType, WebhookEvent +from litellm.proxy._types import ( + CallInfo, + LiteLLM_UserTable, + Litellm_EntityType, + WebhookEvent, +) from litellm.constants import EMAIL_BUDGET_ALERT_TTL @@ -1417,3 +1424,69 @@ async def test_budget_alert_release_failure_does_not_propagate(base_email_logger ) mock_cache.async_delete_cache.assert_awaited_once() + + +class _RecordingEmailLogger(BaseEmailLogger): + def __init__(self): + super().__init__() + self.recipients = [] + + async def send_email(self, from_email, to_email, subject, html_body): + self.recipients.append(to_email) + + +class _UserTable: + def __init__(self, rows): + self._rows = rows + + async def find_unique(self, where): + return self._rows.get(where["user_id"]) + + +def _prisma_client_with_users(*rows): + table: Final = _UserTable({row.user_id: row for row in rows}) + return SimpleNamespace(db=SimpleNamespace(litellm_usertable=table)) + + +def _key_created_event(user_id): + return SendKeyCreatedEmailEvent( + user_id=user_id, + user_email=None, + virtual_key="sk-test", + max_budget=None, + spend=0.0, + event_group=Litellm_EntityType.USER, + event="key_created", + event_message="Key Created", + ) + + +@pytest.mark.asyncio +async def test_key_created_email_goes_to_the_address_stored_for_the_user(monkeypatch): + stored_user: Final = LiteLLM_UserTable( + user_id="user-1", user_email="stored@example.com" + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + _prisma_client_with_users(stored_user), + ) + email_logger: Final = _RecordingEmailLogger() + + await email_logger.send_key_created_email(_key_created_event("user-1")) + + assert email_logger.recipients == [["stored@example.com"]] + + +@pytest.mark.asyncio +async def test_key_created_email_is_refused_for_a_user_the_database_does_not_know( + monkeypatch, +): + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", _prisma_client_with_users() + ) + email_logger: Final = _RecordingEmailLogger() + + with pytest.raises(ValueError, match="User email not found for user_id: user-2"): + await email_logger.send_key_created_email(_key_created_event("user-2")) + + assert email_logger.recipients == [] diff --git a/tests/unit/files/test_main.py b/tests/unit/files/test_main.py index 2704bdfb6ff..cb70b39f4d5 100644 --- a/tests/unit/files/test_main.py +++ b/tests/unit/files/test_main.py @@ -3,9 +3,12 @@ from urllib.parse import parse_qs, urlparse import httpx import pytest +import respx +from pydantic import ValidationError import litellm from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.llms.openai import OpenAIFileObject NATIVE_VERTEX_ROWS: Final = ( b'{"request": {"contents": [{"role": "user", "parts": [{"text": "Who won the 2024 Tour de France?"}]}],' @@ -69,3 +72,52 @@ def test_create_file_passthrough_kwarg_ships_native_rows_byte_for_byte_under_the assert upload.read() == NATIVE_VERTEX_ROWS assert object_name.startswith("litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/") assert file_object.id == f"gs://my-bucket/{object_name}" + + +OPENAI_FILES_API_BASE: Final = "https://files.test/v1" +PROVIDER_FILE: Final = { + "id": "file-abc123", + "bytes": 120, + "created_at": 1700000000, + "filename": "batch.jsonl", + "object": "file", + "purpose": "batch", +} +UNSET_OPTIONAL_FIELDS: Final = {"status": None, "expires_at": None, "status_details": None} + + +@pytest.mark.parametrize( + "provider_extras, expected", + [ + ({}, {**PROVIDER_FILE, **UNSET_OPTIONAL_FIELDS}), + ( + {"status": "processed", "expires_at": 1800000000, "status_details": "ok", "provider_only": {"a": [1]}}, + {**PROVIDER_FILE, "status": "processed", "expires_at": 1800000000, "status_details": "ok"}, + ), + ], + ids=["required-fields-only", "optional-and-unknown-fields"], +) +@respx.mock +async def test_afile_retrieve_returns_the_provider_file_as_an_openai_file_object(provider_extras, expected): + respx.get(f"{OPENAI_FILES_API_BASE}/files/file-abc123").respond(200, json={**PROVIDER_FILE, **provider_extras}) + + file_object: Final = await litellm.afile_retrieve( + file_id="file-abc123", custom_llm_provider="openai", api_key="sk-test", api_base=OPENAI_FILES_API_BASE + ) + + assert type(file_object) is OpenAIFileObject + assert file_object.model_dump() == expected + + +@respx.mock +async def test_afile_retrieve_rejects_a_provider_file_without_its_size(): + provider_file: Final = {key: value for key, value in PROVIDER_FILE.items() if key != "bytes"} + respx.get(f"{OPENAI_FILES_API_BASE}/files/file-abc123").respond(200, json=provider_file) + + with pytest.raises(ValidationError) as exc_info: + await litellm.afile_retrieve( + file_id="file-abc123", custom_llm_provider="openai", api_key="sk-test", api_base=OPENAI_FILES_API_BASE + ) + + assert exc_info.value.title == "OpenAIFileObject" + assert [error["loc"] for error in exc_info.value.errors()] == [("bytes",)] diff --git a/tests/unit/integrations/arize/test_arize_utils.py b/tests/unit/integrations/arize/test_arize_utils.py index 167b083e147..0d833b90148 100644 --- a/tests/unit/integrations/arize/test_arize_utils.py +++ b/tests/unit/integrations/arize/test_arize_utils.py @@ -5,6 +5,7 @@ from typing import Optional import asyncio +import httpx import pytest import litellm @@ -13,6 +14,7 @@ from litellm.integrations._types.open_inference import ( SpanAttributes, ToolCallAttributes, ) +from litellm.integrations.arize._utils import _coerce_response_obj_for_attrs, _parse_passthrough_response from litellm.integrations.arize.arize import ArizeLogger from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import Choices, StandardCallbackDynamicParams @@ -1513,3 +1515,42 @@ def test_arize_mcp_emitter_is_inert_without_a_standard_logging_object(): written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} assert SpanAttributes.TOOL_NAME not in written + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ('{"id": "msg_1", "usage": {"input_tokens": 3}}', {"id": "msg_1", "usage": {"input_tokens": 3}}), + ("{}", {}), + ], +) +def test_coerce_response_obj_decodes_a_json_object_body(text: str, expected: dict[str, object]): + assert _coerce_response_obj_for_attrs(httpx.Response(200, text=text)) == expected + + +@pytest.mark.parametrize("text", ["[1, 2]", '"text"', "7", "null", "true", "not json"]) +def test_coerce_response_obj_keeps_the_response_when_the_body_is_not_a_json_object(text: str): + response = httpx.Response(200, text=text) + + assert _coerce_response_obj_for_attrs(response) is response + + +@pytest.mark.parametrize( + ("raw", "coerced", "kwargs", "expected"), + [ + (None, {"response": '{"id": "wrapped"}'}, {}, {"id": "wrapped"}), + (None, {"response": "[1, 2]"}, {}, None), + ({"id": "raw"}, {"response": "[1, 2]"}, {}, {"id": "raw"}), + ({"id": "raw"}, {"response": "not json"}, {}, {"id": "raw"}), + (None, {"response": "7"}, {"original_response": '{"id": "original"}'}, {"id": "original"}), + (None, None, {"original_response": '{"id": "original"}'}, {"id": "original"}), + (None, None, {"original_response": "[1, 2]"}, None), + (None, None, {"original_response": '"text"'}, None), + (None, None, {"original_response": "null"}, None), + (None, None, {"original_response": "not json"}, None), + ], +) +def test_parse_passthrough_response_reads_only_json_objects_from_text( + raw: object, coerced: object, kwargs: dict[str, object], expected: dict[str, object] | None +): + assert _parse_passthrough_response(raw, coerced, kwargs) == expected diff --git a/tests/unit/integrations/datadog/test_datadog_llm_obs.py b/tests/unit/integrations/datadog/test_datadog_llm_obs.py index 77d2518696c..a7826e11483 100644 --- a/tests/unit/integrations/datadog/test_datadog_llm_obs.py +++ b/tests/unit/integrations/datadog/test_datadog_llm_obs.py @@ -1139,3 +1139,31 @@ def test_reasoning_content_survives_the_mapping(logger: DataDogLLMObsLogger) -> ) assert payload["meta"]["output"]["messages"][0]["reasoning_content"] == "thinking" + + +@pytest.mark.parametrize( + ("raw_arguments", "shipped"), + [ + ('{"city": "Paris", "days": [1, 2]}', {"city": "Paris", "days": [1, 2]}), + ("{}", {}), + ("[1, 2]", "[1, 2]"), + ("null", "null"), + ("true", "true"), + ('"text"', '"text"'), + ("1.5", "1.5"), + ("", ""), + ], +) +def test_tool_arguments_ship_as_an_object_only_when_they_decode_to_one( + logger: DataDogLLMObsLogger, raw_arguments: str, shipped: object +) -> None: + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": raw_arguments}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == shipped diff --git a/tests/unit/integrations/opik/opik_payload_builder/__init__.py b/tests/unit/integrations/opik/opik_payload_builder/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/opik/opik_payload_builder/test_api.py b/tests/unit/integrations/opik/opik_payload_builder/test_api.py new file mode 100644 index 00000000000..d2e20909d3c --- /dev/null +++ b/tests/unit/integrations/opik/opik_payload_builder/test_api.py @@ -0,0 +1,230 @@ +from collections import OrderedDict +from collections.abc import Mapping +from datetime import datetime, timezone +from types import MappingProxyType +from typing import Final +from uuid import UUID + +import pytest +from pydantic import ValidationError + +from litellm.integrations.opik.opik_payload_builder import build_opik_payload +from litellm.integrations.opik.opik_payload_builder.types import SpanPayload, TracePayload +from litellm.types.utils import ModelResponse, Usage + +_START: Final = datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc) +_END: Final = datetime(2026, 1, 2, 3, 4, 6, tzinfo=timezone.utc) +_MESSAGES: Final = [{"role": "user", "content": "hi"}] +_RESPONSE: Final = {"id": "chatcmpl-1", "choices": []} +_HIDDEN_PARAMS: Final = {"model_id": "deployment-1"} +_LOGGING_METADATA: Final = { + "user_api_key_alias": "team-key", + "requester_metadata": {"opik": {"thread_id": "thread-1"}}, +} +_STANDARD_LOGGING_OBJECT: Final = { + "call_type": "acompletion", + "status": "success", + "model": "gpt-4o", + "metadata": _LOGGING_METADATA, + "messages": _MESSAGES, + "response": _RESPONSE, + "hidden_params": _HIDDEN_PARAMS, + "trace_id": "not-forwarded-to-opik", +} +_RESPONSE_OBJ: Final = ModelResponse( + model="gpt-4o", + created=1767323045, + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), +) +_OPIK_FIELDS_OF_THE_LOGGING_OBJECT: Final = { + "type": "acompletion", + "status": "success", + "model": "gpt-4o", + "hidden_params": _HIDDEN_PARAMS, +} +_EXISTING_TRACE: Final = {"metadata": {"opik": {"current_span_data": {"trace_id": "trace-1", "id": "span-0"}}}} + + +def _payloads_attached_to_trace_1(standard_logging_object: object) -> tuple[TracePayload | None, SpanPayload]: + return build_opik_payload( + kwargs={"standard_logging_object": standard_logging_object, "litellm_params": _EXISTING_TRACE}, + response_obj=_RESPONSE_OBJ, + start_time=_START, + end_time=_END, + project_name="default-project", + ) + + +def _span_attached_to_trace_1(span_id: str, metadata: Mapping[str, object]) -> SpanPayload: + return SpanPayload( + id=span_id, + project_name="default-project", + trace_id="trace-1", + parent_span_id="span-0", + name="gpt-4o_chat.completion_1767323045", + type="llm", + model="gpt-4o", + start_time="2026-01-02T03:04:05Z", + end_time="2026-01-02T03:04:06Z", + input=_MESSAGES, + output=_RESPONSE, + metadata=metadata, + tags=[], + usage={"completion_tokens": 2, "prompt_tokens": 1, "total_tokens": 3}, + ) + + +def test_build_opik_payload_creates_a_trace_and_its_span_from_the_standard_logging_object() -> None: + trace, span = build_opik_payload( + kwargs={ + "standard_logging_object": _STANDARD_LOGGING_OBJECT, + "custom_llm_provider": "openai", + "response_cost": 0.25, + }, + response_obj=_RESPONSE_OBJ, + start_time=_START, + end_time=_END, + project_name="default-project", + ) + metadata: Final = { + "thread_id": "thread-1", + "created_from": "litellm", + **_LOGGING_METADATA, + **_OPIK_FIELDS_OF_THE_LOGGING_OBJECT, + "cost": {"total_tokens": 0.25, "currency": "USD"}, + } + + assert trace is not None + assert (trace, span) == ( + TracePayload( + project_name="default-project", + id=trace.id, + name="chat.completion", + start_time="2026-01-02T03:04:05Z", + end_time="2026-01-02T03:04:06Z", + input=_MESSAGES, + output=_RESPONSE, + metadata=metadata, + tags=["openai"], + thread_id="thread-1", + ), + SpanPayload( + id=span.id, + project_name="default-project", + trace_id=trace.id, + name="gpt-4o_chat.completion_1767323045", + type="llm", + model="gpt-4o", + start_time="2026-01-02T03:04:05Z", + end_time="2026-01-02T03:04:06Z", + input=_MESSAGES, + output=_RESPONSE, + metadata=metadata, + tags=["openai"], + usage={"completion_tokens": 2, "prompt_tokens": 1, "total_tokens": 3}, + provider="openai", + total_cost=0.25, + ), + ) + assert [UUID(trace.id).version, UUID(span.id).version, trace.id != span.id] == [7, 7, True] + assert [span.input is _MESSAGES, span.output is _RESPONSE, span.metadata["hidden_params"] is _HIDDEN_PARAMS] == [ + True, + True, + True, + ] + assert list(span.metadata) == [ + "thread_id", + "created_from", + "user_api_key_alias", + "requester_metadata", + "type", + "status", + "model", + "hidden_params", + "cost", + ] + + +@pytest.mark.parametrize( + "standard_logging_object", + [ + _STANDARD_LOGGING_OBJECT, + OrderedDict(_STANDARD_LOGGING_OBJECT), + MappingProxyType(_STANDARD_LOGGING_OBJECT), + {**_STANDARD_LOGGING_OBJECT, "metadata": MappingProxyType(_LOGGING_METADATA)}, + ], +) +def test_build_opik_payload_reads_any_string_keyed_mapping_as_the_standard_logging_object( + standard_logging_object: object, +) -> None: + trace, span = _payloads_attached_to_trace_1(standard_logging_object) + + assert (trace, span) == ( + None, + _span_attached_to_trace_1( + span.id, + { + "thread_id": "thread-1", + "created_from": "litellm", + **_LOGGING_METADATA, + **_OPIK_FIELDS_OF_THE_LOGGING_OBJECT, + }, + ), + ) + + +@pytest.mark.parametrize( + "standard_logging_object", + [ + {key: value for key, value in _STANDARD_LOGGING_OBJECT.items() if key != "metadata"}, + {**_STANDARD_LOGGING_OBJECT, "metadata": None}, + {**_STANDARD_LOGGING_OBJECT, "metadata": {}}, + {**_STANDARD_LOGGING_OBJECT, "metadata": ""}, + ], +) +def test_build_opik_payload_without_standard_logging_metadata_keeps_only_the_logging_object_fields( + standard_logging_object: object, +) -> None: + trace, span = _payloads_attached_to_trace_1(standard_logging_object) + + assert (trace, span) == ( + None, + _span_attached_to_trace_1(span.id, {"created_from": "litellm", **_OPIK_FIELDS_OF_THE_LOGGING_OBJECT}), + ) + + +def test_build_opik_payload_of_an_empty_standard_logging_object_has_empty_input_and_output() -> None: + _, span = _payloads_attached_to_trace_1({}) + + assert (span.input, span.output, span.metadata) == ({}, {}, {"created_from": "litellm"}) + + +@pytest.mark.parametrize( + "standard_logging_object", + [ + None, + "standard_logging_object", + ["messages", "response"], + list(_STANDARD_LOGGING_OBJECT.items()), + 7, + {**_STANDARD_LOGGING_OBJECT, 7: "keys must be strings"}, + {**_STANDARD_LOGGING_OBJECT, "metadata": "metadata"}, + {**_STANDARD_LOGGING_OBJECT, "metadata": ["user_api_key_alias"]}, + {**_STANDARD_LOGGING_OBJECT, "metadata": 7}, + {**_STANDARD_LOGGING_OBJECT, "metadata": {7: "keys must be strings"}}, + ], +) +def test_build_opik_payload_rejects_a_standard_logging_object_that_is_not_a_string_keyed_mapping( + standard_logging_object: object, +) -> None: + with pytest.raises(ValidationError) as raised: + _payloads_attached_to_trace_1(standard_logging_object) + + assert "input_value" not in str(raised.value) + + +def test_build_opik_payload_without_a_standard_logging_object_raises_a_key_error() -> None: + with pytest.raises(KeyError, match="standard_logging_object"): + build_opik_payload( + kwargs={}, response_obj=_RESPONSE_OBJ, start_time=_START, end_time=_END, project_name="default-project" + ) diff --git a/tests/unit/integrations/opik/test_opik_extractors.py b/tests/unit/integrations/opik/test_opik_extractors.py index 6f85a1c6090..4ef9f500b54 100644 --- a/tests/unit/integrations/opik/test_opik_extractors.py +++ b/tests/unit/integrations/opik/test_opik_extractors.py @@ -1,4 +1,7 @@ +import pytest + from litellm.integrations.opik.opik_payload_builder.extractors import ( + apply_proxy_header_overrides, extract_opik_metadata, ) @@ -82,3 +85,32 @@ def test_extract_opik_metadata_requester_metadata_overrides_all_other_sources(): "workspace": "requester-workspace", "thread_id": "requester-thread", } + + +@pytest.mark.parametrize( + ("opik_tags_header", "expected_tags"), + [ + ('["from-header", "second"]', ["from-request", "from-header", "second"]), + ('["text", 7, null, {"nested": [1]}]', ["from-request", "text", 7, None, {"nested": [1]}]), + ("[]", ["from-request"]), + ('{"not": "a list"}', ["from-request"]), + ('"not-a-list"', ["from-request"]), + ("null", ["from-request"]), + ("not json", ["from-request"]), + ], +) +def test_opik_tags_header_adds_tags_only_when_it_is_a_json_list(opik_tags_header: str, expected_tags: list[object]): + overrides = apply_proxy_header_overrides("project", ["from-request"], None, {"opik_tags": opik_tags_header}) + + assert overrides == ("project", expected_tags, None) + + +def test_opik_headers_override_the_project_name_and_thread_id(): + overrides = apply_proxy_header_overrides( + "project", + ["from-request"], + "thread-from-request", + {"opik_project_name": "header-project", "opik_thread_id": "header-thread", "opik_tags": "", "x-other": "1"}, + ) + + assert overrides == ("header-project", ["from-request"], "header-thread") diff --git a/tests/unit/integrations/test_galileo.py b/tests/unit/integrations/test_galileo.py index d0709b966d4..d067b31a651 100644 --- a/tests/unit/integrations/test_galileo.py +++ b/tests/unit/integrations/test_galileo.py @@ -1,7 +1,11 @@ +from dataclasses import dataclass from datetime import datetime, timezone +from decimal import Decimal +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pydantic import BaseModel, ValidationError from litellm.integrations.galileo import GalileoObserve @@ -905,3 +909,78 @@ async def test_galileo_async_log_success_appends_and_flushes(galileo_v2_env): assert "/ingest/traces/" in flushed_url["url"] assert logger.in_memory_records == [] + + +class RealtimeEvent(BaseModel): + type: str + + +@dataclass(frozen=True) +class ThirdPartyEvent: + model_dump: object + + +@pytest.mark.parametrize( + ("event", "expected_output"), + [ + (RealtimeEvent(type="response.done"), '[{"type": "response.done"}]'), + (ThirdPartyEvent(model_dump=lambda: {"type": "custom"}), '[{"type": "custom"}]'), + (ThirdPartyEvent(model_dump=lambda: [Decimal("2.5")]), '[["2.5"]]'), + (Decimal("1.5"), '["1.5"]'), + ], +) +def test_galileo_realtime_output_serializes_events_that_json_cannot_encode( + galileo_v2_env: None, event: object, expected_output: str +) -> None: + assert GalileoObserve().get_output_str_from_response([event], {"call_type": "_arealtime"}) == expected_output + + +@pytest.mark.parametrize("model_dump", [None, "not callable"]) +def test_galileo_realtime_output_with_a_model_dump_that_cannot_be_called_raises_a_validation_error( + galileo_v2_env: None, model_dump: object +) -> None: + with pytest.raises(ValidationError) as raised: + GalileoObserve().get_output_str_from_response( + [ThirdPartyEvent(model_dump=model_dump)], {"call_type": "_arealtime"} + ) + + assert "input_value" not in str(raised.value) + + +@dataclass(frozen=True) +class ThirdPartyMessage: + json: object + + +def _chat_response_carrying(message: object) -> ModelResponse: + response: Final = ModelResponse(choices=[Choices(message=Message(content="replaced", role="assistant"))]) + response.choices[0].message = message + return response + + +@pytest.mark.parametrize( + ("message", "expected_output"), + [ + ( + ThirdPartyMessage(json=lambda: '{"role":"assistant","content":"hi"}'), + '{"role": "assistant", "content": "hi"}', + ), + ( + ThirdPartyMessage(json=lambda: {"role": "assistant", "content": "hi"}), + '{"role": "assistant", "content": "hi"}', + ), + (ThirdPartyMessage(json=lambda: '"just text"'), "just text"), + (ThirdPartyMessage(json=lambda: None), ""), + ({"role": "assistant", "content": "hi"}, '{"role": "assistant", "content": "hi"}'), + ("plain reply", "plain reply"), + (None, ""), + ], +) +def test_galileo_output_str_of_a_chat_response_whose_message_is_not_a_litellm_message( + galileo_v2_env: None, message: object, expected_output: str +) -> None: + output: Final = GalileoObserve().get_output_str_from_response( + _chat_response_carrying(message), {"call_type": "acompletion"} + ) + + assert output == expected_output diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 175bd95c263..2ce3561d441 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -6762,3 +6762,37 @@ class TestOpenTelemetryNonInferenceUsage(unittest.TestCase): self.assertEqual( self._time_per_output_token_calls("aget_responses", response_obj=self.BACKGROUND_RESPONSE_OBJ), 1 ) + + +def _raw_response_span_attributes(original_response: str) -> dict[str, object]: + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + span = tracer_provider.get_tracer(__name__).start_span("raw_gen_ai_request") + + OpenTelemetry(tracer_provider=tracer_provider).set_raw_request_attributes( + span, + {"litellm_params": {"custom_llm_provider": "vertex_ai"}, "original_response": original_response}, + None, + ) + span.end() + + return dict(span_exporter.get_finished_spans()[0].attributes or {}) + + +@pytest.mark.parametrize( + ("original_response", "expected"), + [ + ('{"id": "r1", "model": "m"}', {"llm.vertex_ai.id": "r1", "llm.vertex_ai.model": "m"}), + ("{}", {}), + ("not json", {"llm.vertex_ai.stringified_raw_response": "not json"}), + ("[1, 2]", {}), + ('"text"', {}), + ("7", {}), + ("null", {}), + ], +) +def test_set_raw_request_attributes_stamps_only_json_object_responses( + original_response: str, expected: dict[str, object] +): + assert _raw_response_span_attributes(original_response) == expected diff --git a/tests/unit/integrations/test_opik_utils.py b/tests/unit/integrations/test_opik_utils.py index a4250acf1dc..376c70a4ed9 100644 --- a/tests/unit/integrations/test_opik_utils.py +++ b/tests/unit/integrations/test_opik_utils.py @@ -4,7 +4,7 @@ import uuid from datetime import datetime, timezone from unittest.mock import patch -from litellm.integrations.opik.utils import create_uuid7 +from litellm.integrations.opik.utils import create_uuid7, get_traces_and_spans_from_payload def _timestamp_ms(uuid_str: str) -> int: @@ -27,3 +27,22 @@ def test_create_uuid7_encodes_timestamp_in_milliseconds(): value = create_uuid7() assert _timestamp_ms(value) == int(fixed.timestamp() * 1000) + + +def test_queued_opik_payloads_are_split_into_traces_and_spans_without_their_null_fields(): + traces, spans = get_traces_and_spans_from_payload( + [ + {"id": "trace-1", "name": "chat.completion", "thread_id": None, "input": {"kept": None}}, + {"id": "span-1", "type": "llm", "parent_span_id": None, "tags": [], "total_cost": 0}, + {"id": "span-2", "type": None}, + ] + ) + + assert (traces, spans) == ( + [{"id": "trace-1", "name": "chat.completion", "input": {"kept": None}}], + [{"id": "span-1", "type": "llm", "tags": [], "total_cost": 0}, {"id": "span-2"}], + ) + + +def test_an_empty_opik_queue_has_no_traces_and_no_spans(): + assert get_traces_and_spans_from_payload([]) == ([], []) diff --git a/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index a0766ac3d58..726ce2f7f1d 100644 --- a/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -1,5 +1,5 @@ import logging -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping from dataclasses import dataclass, field from types import MappingProxyType from typing import Literal, Protocol @@ -33,6 +33,7 @@ from litellm.types.utils import ( ) from litellm.types.vector_stores import ( VectorStoreResultContent, + VectorStoreSearchFailure, VectorStoreSearchResponse, VectorStoreSearchResult, ) @@ -457,6 +458,166 @@ async def test_a_failing_vector_store_is_reported_on_the_streaming_chunk(registr ) +@dataclass +class ThirdPartyMessage: + provider_specific_fields: dict[str, object] | None = None + + +@dataclass +class ThirdPartyChoice: + message: ThirdPartyMessage | None = None + delta: ThirdPartyMessage | None = None + + +@dataclass +class ThirdPartyResponse: + choices: object + + +_SEARCH_FAILURES = ( + VectorStoreSearchFailure(vector_store_id="vs-broken", custom_llm_provider="bedrock", error="search timed out"), +) + + +def _logging_obj_with_search_failures() -> FakeLoggingObj: + logging_obj = FakeLoggingObj({}) + logging_obj.model_call_details["vector_store_search_failures"] = _SEARCH_FAILURES + return logging_obj + + +@pytest.mark.asyncio +async def test_search_failures_join_the_provider_fields_the_message_already_carries() -> None: + existing_fields: dict[str, object] = {"citations": ["doc-1"]} + response = ModelResponse( + choices=[Choices(message=Message(content="an answer", provider_specific_fields=existing_fields))] + ) + + returned = await VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=None) + ).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": _logging_obj_with_search_failures()}, + response=response, + call_type=CallTypes.acompletion, + ) + + assert returned is response + assert _first_message(response).provider_specific_fields is existing_fields + assert existing_fields == {"citations": ["doc-1"], "vector_store_search_failures": _SEARCH_FAILURES} + + +@pytest.mark.asyncio +async def test_search_failures_join_the_provider_fields_the_streaming_delta_already_carries() -> None: + existing_fields: dict[str, object] = {"citations": ["doc-1"]} + chunk = ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="an answer", provider_specific_fields=existing_fields))] + ) + + returned = await VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=None) + ).async_post_call_streaming_deployment_hook( + request_data=_logging_obj_with_search_failures().model_call_details, + response_chunk=chunk, + call_type=CallTypes.acompletion, + ) + + assert returned is chunk + assert chunk.choices[0].delta.provider_specific_fields is existing_fields + assert existing_fields == {"citations": ["doc-1"], "vector_store_search_failures": _SEARCH_FAILURES} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("as_choices", [list, tuple, iter]) +async def test_a_chunk_that_only_looks_like_a_chat_completion_chunk_is_annotated_too( + as_choices: Callable[[list[ThirdPartyChoice]], Iterable[ThirdPartyChoice]], +) -> None: + delta = ThirdPartyMessage() + chunk = ThirdPartyResponse(choices=as_choices([ThirdPartyChoice(delta=None), ThirdPartyChoice(delta=delta)])) + + returned = await VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=None) + ).async_post_call_streaming_deployment_hook( + request_data=_logging_obj_with_search_failures().model_call_details, + response_chunk=chunk, + call_type=CallTypes.acompletion, + ) + + assert returned is chunk + assert delta.provider_specific_fields == {"vector_store_search_failures": _SEARCH_FAILURES} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("choices", [None, [], ()]) +async def test_a_response_without_choices_is_returned_untouched( + choices: object, warnings: list[logging.LogRecord] +) -> None: + response = ThirdPartyResponse(choices=choices) + + returned = await VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=None) + ).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": _logging_obj_with_search_failures()}, + response=response, + call_type=CallTypes.acompletion, + ) + + assert returned is response + assert warnings == [] + + +@pytest.mark.asyncio +async def test_a_chunk_whose_choices_cannot_be_iterated_is_logged_and_passed_through( + warnings: list[logging.LogRecord], +) -> None: + chunk = ThirdPartyResponse(choices=7) + + returned = await VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=None) + ).async_post_call_streaming_deployment_hook( + request_data=_logging_obj_with_search_failures().model_call_details, + response_chunk=chunk, + call_type=CallTypes.acompletion, + ) + + assert returned is chunk + assert [record.levelname for record in warnings] == ["ERROR"] + assert warnings[0].getMessage().startswith("Error adding search results to streaming chunk: ") + assert "input_value" not in warnings[0].getMessage() + + +@pytest.mark.asyncio +async def test_the_search_receives_the_requests_own_metadata_object(registry_with: RegisterStores) -> None: + registry_with("vs-router") + router = RecordingRouter() + metadata = {"user_api_key_team_id": "team-a"} + + await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=router)), + ["vs-router"], + FakeLoggingObj(metadata), + ) + + assert [call["metadata"] is metadata for call in router.calls] == [True] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("litellm_params", [None, "metadata", ["metadata"], {"other": "value"}]) +async def test_a_request_without_metadata_in_its_litellm_params_searches_with_empty_metadata( + registry_with: RegisterStores, litellm_params: object +) -> None: + registry_with("vs-router") + router = RecordingRouter() + logging_obj = FakeLoggingObj({}) + logging_obj.model_call_details["litellm_params"] = litellm_params + + await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=router)), + ["vs-router"], + logging_obj, + ) + + assert [call["metadata"] for call in router.calls] == [{}] + + @pytest.mark.asyncio async def test_error_mode_fails_the_request_instead_of_answering_without_the_knowledge_base( registry_with: RegisterStores, diff --git a/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py b/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py index 6440825e135..a419e73a2e1 100644 --- a/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py +++ b/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock import pytest +from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator from litellm.llms.a2a.chat.transformation import A2AConfig from litellm.types.utils import ModelResponse @@ -85,3 +86,38 @@ def test_transform_request_tags_the_message_with_its_kind(optional_params: dict) ) assert request["params"]["message"]["kind"] == "message" + + +def test_get_model_response_iterator_parses_the_agent_stream(): + iterator = A2AConfig().get_model_response_iterator( + streaming_response=iter( + [ + '{"jsonrpc":"2.0","id":"1","result":{"kind":"task","status":{"state":"completed"},' + '"artifacts":[{"parts":[{"kind":"text","text":"7"}]}]}}' + ] + ), + sync_stream=True, + ) + + chunk = next(iterator) + + assert isinstance(iterator, A2AModelResponseIterator) + assert chunk["text"] == "7" + assert chunk["finish_reason"] == "stop" + + +def test_resolve_agent_config_from_registry_returns_the_explicit_headers_object_untouched(): + headers: dict[str, object] = {"X-Test": "value"} + optional_params: dict[str, object] = {"stream": True} + + resolved = A2AConfig.resolve_agent_config_from_registry( + agent_name="test-agent", + api_base="http://explicit.example", + api_key="explicit-key", + headers=headers, + optional_params=optional_params, + ) + + assert resolved[2] is headers + assert resolved[:2] == ("http://explicit.example", "explicit-key") + assert optional_params == {"stream": True} diff --git a/tests/unit/llms/aiohttp_openai/__init__.py b/tests/unit/llms/aiohttp_openai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/aiohttp_openai/chat/__init__.py b/tests/unit/llms/aiohttp_openai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/aiohttp_openai/chat/test_transformation.py b/tests/unit/llms/aiohttp_openai/chat/test_transformation.py new file mode 100644 index 00000000000..253f9dbe746 --- /dev/null +++ b/tests/unit/llms/aiohttp_openai/chat/test_transformation.py @@ -0,0 +1,101 @@ +from unittest.mock import AsyncMock, Mock + +import pytest +from aiohttp import ClientResponse +from pydantic import ValidationError + +from litellm.llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig +from litellm.types.utils import ModelResponse + + +async def _transform(body: object) -> ModelResponse: + raw_response = Mock(spec=ClientResponse) + raw_response.json = AsyncMock(return_value=body) + return await AiohttpOpenAIChatConfig().transform_response( + model="gpt-4o", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +async def test_transform_response_copies_the_openai_body_onto_the_model_response(): + response = await _transform( + { + "id": "chatcmpl-1", + "created": 1700000000, + "model": "gpt-4o-2024", + "object": "chat.completion", + "system_fingerprint": "fp_1", + "choices": [ + {"index": 0, "finish_reason": "length", "message": {"role": "assistant", "content": "Hi"}}, + { + "index": 1, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + }, + ], + } + ) + + assert response.id == "chatcmpl-1" + assert response.created == 1700000000 + assert response.model == "gpt-4o-2024" + assert response.object == "chat.completion" + assert response.system_fingerprint == "fp_1" + assert [(choice.index, choice.finish_reason, choice.message.content) for choice in response.choices] == [ + (0, "length", "Hi"), + (1, "tool_calls", None), + ] + assert response.choices[1].message.tool_calls[0].function.name == "lookup" + + +async def test_transform_response_fills_choice_defaults_for_an_empty_choice(): + response = await _transform({"choices": [{}]}) + + assert [(choice.index, choice.finish_reason, choice.message.role) for choice in response.choices] == [ + (0, "stop", "assistant") + ] + assert response.id is None + + +async def test_transform_response_returns_no_choices_for_an_empty_choices_list(): + response = await _transform({"id": "chatcmpl-1", "choices": []}) + + assert response.choices == [] + assert response.id == "chatcmpl-1" + + +@pytest.mark.parametrize( + "body", + [ + {"id": "chatcmpl-1"}, + {"choices": None}, + {"choices": 7}, + {"choices": "secret-completion"}, + {"choices": {"message": "secret-completion"}}, + {"choices": ["secret-completion"]}, + {"choices": [{"index": 0}, ["secret-completion"]]}, + ], +) +async def test_transform_response_rejects_choices_that_are_not_a_list_of_objects_without_echoing_them(body: object): + with pytest.raises(ValidationError) as exc_info: + await _transform(body) + + assert "secret-completion" not in str(exc_info.value) + + +async def test_transform_response_rejects_a_choice_whose_message_is_not_an_object(): + with pytest.raises(ValidationError, match="validation error for Choices"): + await _transform({"choices": [{"message": "text"}]}) diff --git a/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py index f385c2f2211..e9736d47537 100644 --- a/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py +++ b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py @@ -12,6 +12,8 @@ import time import httpx import pytest +from openai.types.file_deleted import FileDeleted +from pydantic import ValidationError from unittest.mock import Mock, patch from litellm.llms.anthropic.files.transformation import ( @@ -196,6 +198,45 @@ class TestAnthropicFilesConfig: assert result.purpose == "messages" assert result.status == "uploaded" + def test_create_file_response_maps_the_anthropic_file_onto_an_openai_file(self) -> None: + result = self.config.transform_create_file_response( + model=None, + raw_response=httpx.Response( + 200, + json={ + "id": "file-abc123", + "type": "file", + "filename": "document.pdf", + "mime_type": "application/pdf", + "size_bytes": 12345, + "created_at": "2025-01-15T10:30:00Z", + }, + ), + logging_obj=Mock(), + litellm_params={}, + ) + + assert result == OpenAIFileObject( + id="file-abc123", + bytes=12345, + created_at=1736937000, + filename="document.pdf", + object="file", + purpose="messages", + status="uploaded", + status_details=None, + ) + + @pytest.mark.parametrize("body", [b'["file-abc123"]', b'"file-abc123"', b"null", b"7"]) + def test_create_file_response_rejects_a_body_that_is_not_a_json_object(self, body: bytes) -> None: + with pytest.raises(ValidationError): + self.config.transform_create_file_response( + model=None, + raw_response=httpx.Response(200, content=body), + logging_obj=Mock(), + litellm_params={}, + ) + def test_transform_retrieve_file_request(self): url, params = self.config.transform_retrieve_file_request( file_id="file-abc123", @@ -245,6 +286,43 @@ class TestAnthropicFilesConfig: assert result.id == "file-abc123" assert result.bytes == 5000 + def test_retrieve_file_response_maps_the_anthropic_file_onto_an_openai_file(self) -> None: + result = self.config.transform_retrieve_file_response( + raw_response=httpx.Response( + 200, + json={ + "id": "file-abc123", + "type": "file", + "filename": "document.pdf", + "mime_type": "application/pdf", + "size_bytes": 5000, + "created_at": "2025-06-01T12:00:00Z", + }, + ), + logging_obj=Mock(), + litellm_params={}, + ) + + assert result == OpenAIFileObject( + id="file-abc123", + bytes=5000, + created_at=1748779200, + filename="document.pdf", + object="file", + purpose="messages", + status="uploaded", + status_details=None, + ) + + @pytest.mark.parametrize("body", [b'["file-abc123"]', b'"file-abc123"', b"null", b"7"]) + def test_retrieve_file_response_rejects_a_body_that_is_not_a_json_object(self, body: bytes) -> None: + with pytest.raises(ValidationError): + self.config.transform_retrieve_file_response( + raw_response=httpx.Response(200, content=body), + logging_obj=Mock(), + litellm_params={}, + ) + def test_transform_delete_file_request(self): url, params = self.config.transform_delete_file_request( file_id="file-abc123", @@ -271,6 +349,35 @@ class TestAnthropicFilesConfig: assert result.deleted is True assert result.object == "file" + @pytest.mark.parametrize( + ("payload", "expected_id"), + [ + ({"id": "file-abc123", "type": "file_deleted"}, "file-abc123"), + ({"id": "file-abc123", "unknown": [1, {"nested": None}]}, "file-abc123"), + ({"type": "error", "error": {"type": "not_found_error", "message": "File not found"}}, ""), + ({}, ""), + ], + ) + def test_delete_file_response_reports_the_id_anthropic_returned(self, payload: object, expected_id: str) -> None: + result = self.config.transform_delete_file_response( + raw_response=httpx.Response(200, json=payload), + logging_obj=Mock(), + litellm_params={}, + ) + + assert result == FileDeleted(id=expected_id, deleted=True, object="file") + + @pytest.mark.parametrize( + "body", [b'["file-abc123"]', b'"file-abc123"', b"null", b"7", b'{"id": null}', b'{"id": 7}'] + ) + def test_delete_file_response_rejects_a_body_without_a_string_id(self, body: bytes) -> None: + with pytest.raises(ValidationError): + self.config.transform_delete_file_response( + raw_response=httpx.Response(200, content=body), + logging_obj=Mock(), + litellm_params={}, + ) + def test_transform_list_files_request(self): url, params = self.config.transform_list_files_request( purpose=None, diff --git a/tests/unit/llms/azure/test_azure_speech_audio_transcription.py b/tests/unit/llms/azure/test_azure_speech_audio_transcription.py index b447645bae8..4330f05e79c 100644 --- a/tests/unit/llms/azure/test_azure_speech_audio_transcription.py +++ b/tests/unit/llms/azure/test_azure_speech_audio_transcription.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock import httpx import pytest +from pydantic import ValidationError import litellm from litellm.llms.azure.audio_transcription.transformation import ( @@ -226,3 +227,22 @@ def test_azure_speech_transcription_routes_through_provider_config(monkeypatch): assert audio_handler.call_args.kwargs["custom_llm_provider"] == "azure" +def test_azure_speech_audio_transcription_response_keeps_the_raw_payload_as_hidden_params(): + payload = {"RecognitionStatus": "Success", "NBest": [{"Lexical": "hello world", "Confidence": 0.9}], "Offset": 3} + + response = AzureSpeechAudioTranscriptionConfig().transform_audio_transcription_response( + httpx.Response(200, json=payload) + ) + + assert response.text == "hello world" + assert response._hidden_params == payload + + +@pytest.mark.parametrize("payload", [7, "spoken secret", [{"DisplayText": "spoken secret"}]]) +def test_azure_speech_audio_transcription_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + AzureSpeechAudioTranscriptionConfig().transform_audio_transcription_response( + httpx.Response(200, json=payload) + ) + + assert "spoken secret" not in str(exc_info.value) diff --git a/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py b/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py new file mode 100644 index 00000000000..6ca70324162 --- /dev/null +++ b/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py @@ -0,0 +1,14 @@ +from litellm.llms.bedrock.image_edit.stability_transformation import BedrockStabilityImageEditConfig + + +def test_transform_image_edit_request_returns_the_json_body_and_no_files(): + request = BedrockStabilityImageEditConfig().transform_image_edit_request( + model="stability.stable-image-inpaint-v1:0", + prompt="add a red hat", + image=b"\x89PNG-bytes", + image_edit_optional_request_params={}, + litellm_params={}, + headers={}, + ) + + assert request == ({"output_format": "png", "prompt": "add a red hat", "image": "iVBORy1ieXRlcw=="}, {}) diff --git a/tests/unit/llms/bedrock/image_generation/__init__.py b/tests/unit/llms/bedrock/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/image_generation/test_amazon_titan_transformation.py b/tests/unit/llms/bedrock/image_generation/test_amazon_titan_transformation.py new file mode 100644 index 00000000000..eb8b9de6721 --- /dev/null +++ b/tests/unit/llms/bedrock/image_generation/test_amazon_titan_transformation.py @@ -0,0 +1,57 @@ +import pytest + +from litellm.llms.bedrock.image_generation.amazon_titan_transformation import AmazonTitanImageGenerationConfig + + +@pytest.mark.parametrize( + ("non_default_params", "expected"), + [ + ( + {"size": "1024x512", "n": 2, "quality": "hd"}, + {"imageGenerationConfig": {"width": 1024, "height": 512, "numberOfImages": 2, "quality": "premium"}}, + ), + ({"quality": "low"}, {"imageGenerationConfig": {"quality": "standard"}}), + ({"quality": "auto", "size": None, "n": None}, {}), + ], +) +def test_map_openai_params_builds_the_image_generation_config( + non_default_params: dict[str, object], expected: dict[str, object] +): + optional_params = AmazonTitanImageGenerationConfig.map_openai_params( + non_default_params=non_default_params, optional_params={} + ) + + assert optional_params == expected + + +@pytest.mark.parametrize( + ("optional_params", "expected"), + [ + ( + {"imageGenerationConfig": {"width": 512, "numberOfImages": 2}, "negativeText": "blurry"}, + { + "taskType": "TEXT_IMAGE", + "textToImageParams": {"text": "a cat", "negativeText": "blurry"}, + "imageGenerationConfig": {"width": 512, "numberOfImages": 2}, + }, + ), + ( + {"taskType": "COLOR_GUIDED_GENERATION", "negativeText": ""}, + { + "taskType": "COLOR_GUIDED_GENERATION", + "textToImageParams": {"text": "a cat"}, + "imageGenerationConfig": {}, + }, + ), + ], +) +def test_transform_request_body_builds_the_titan_request( + optional_params: dict[str, object], expected: dict[str, object] +): + request_body = AmazonTitanImageGenerationConfig.transform_request_body( + text="a cat", optional_params=optional_params + ) + + assert request_body == expected + assert list(request_body["textToImageParams"]) == list(expected["textToImageParams"]) + assert optional_params == {} diff --git a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index c40830b238f..495e0f56c84 100644 --- a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -8,7 +8,9 @@ forward_client_headers_to_llm_api were not being passed to Bedrock rerank provid import json from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import pytest +from pydantic import ValidationError import litellm from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo @@ -506,3 +508,60 @@ async def test_bedrock_rerank_records_llm_api_duration(): assert response._hidden_params["litellm_overhead_time_ms"] is not None assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"] + + +RERANK_MODEL_ARN = "arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0" +NON_OBJECT_BODIES = [7, "sensitive-document", [{"index": 0, "relevanceScore": 0.9, "id": "sensitive-document"}]] + + +def _rerank_through_transport(payload: object, *, is_async: bool): + transport = httpx.MockTransport(lambda request: httpx.Response(200, json=payload)) + client = ( + AsyncHTTPHandler(transport=transport) if is_async else HTTPHandler(client=httpx.Client(transport=transport)) + ) + return BedrockRerankHandler().rerank( + model=RERANK_MODEL_ARN, + query=test_query, + documents=test_documents, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "example-secret", + "aws_region_name": "us-west-2", + }, + logging_obj=Mock(), + _is_async=is_async, + client=client, + ) + + +def _assert_is_the_bedrock_ranking(response: litellm.RerankResponse) -> None: + assert response.results == [ + {"index": 2, "relevance_score": 0.95}, + {"index": 0, "relevance_score": 0.1}, + {"index": 1, "relevance_score": 0.05}, + ] + assert response.meta == {"billed_units": {"search_units": 1}, "tokens": {}} + + +def test_bedrock_rerank_maps_the_upstream_ranking(): + _assert_is_the_bedrock_ranking(_rerank_through_transport(bedrock_rerank_response, is_async=False)) + + +async def test_bedrock_arerank_maps_the_upstream_ranking(): + _assert_is_the_bedrock_ranking(await _rerank_through_transport(bedrock_rerank_response, is_async=True)) + + +@pytest.mark.parametrize("payload", NON_OBJECT_BODIES) +def test_bedrock_rerank_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _rerank_through_transport(payload, is_async=False) + + assert "sensitive-document" not in str(exc_info.value) + + +@pytest.mark.parametrize("payload", NON_OBJECT_BODIES) +async def test_bedrock_arerank_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + await _rerank_through_transport(payload, is_async=True) + + assert "sensitive-document" not in str(exc_info.value) diff --git a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py index 20bf65ee385..4526c428f1a 100644 --- a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py @@ -9,10 +9,12 @@ the sharded CI (coverage collection runs against this tree). import json import os +import httpx import pytest from unittest.mock import AsyncMock, patch, MagicMock import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.llms.bedrock.search.transformation import ( AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION, AgentCoreSearchConfig, @@ -635,3 +637,44 @@ class TestAgentCoreSearchEdgeCases: monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) assert search_provider_cost_per_query(model="agentcore/search", custom_llm_provider="agentcore") == (0.0, 0.0) + + +def _transform_text(text: str) -> SearchResponse: + return AgentCoreSearchConfig().transform_search_response( + raw_response=httpx.Response(200, text=text), + logging_obj=MagicMock(), + ) + + +def _as_tuples(response: SearchResponse) -> list[tuple[str, str, str, str | None, str | None]]: + return [(r.title, r.url, r.snippet, r.date, r.last_updated) for r in response.results] + + +EXPECTED_MCP_RESULTS = [ + ("Test Result 1", "https://example.com/1", "Snippet for result 1", "2026-06-16", None), + ("Test Result 2", "https://example.com/2", "Snippet for result 2", None, None), +] + + +@pytest.mark.parametrize( + "text", + [ + json.dumps(_mcp_response_body()), + f"event: message\ndata: {json.dumps(_mcp_response_body())}\n\n", + f"data: 5\n\ndata: [1]\n\ndata: {json.dumps(_mcp_response_body())}\n\n", + json.dumps({"result": {"content": [{"type": "text", "text": json.dumps({"results": MCP_RESULTS})}]}}), + json.dumps({"result": {"content": [{"type": "text", "text": "prose"}], "structuredContent": MCP_RESULTS}}), + ], +) +def test_transform_search_response_reads_results_from_a_real_http_response(text: str): + assert _as_tuples(_transform_text(text)) == EXPECTED_MCP_RESULTS + + +@pytest.mark.parametrize( + "block_text", + ["null", "5", "true", '"text"', "{}", '{"results": null}', '{"results": "text"}', '["scalar", 5, null]'], +) +def test_transform_search_response_ignores_text_blocks_without_result_objects(block_text: str): + body = {"result": {"content": [{"type": "text", "text": block_text}]}} + + assert _transform_text(json.dumps(body)).results == [] diff --git a/tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py b/tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py index d6e2c4a3e06..c947967417d 100644 --- a/tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py +++ b/tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py @@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import ValidationError from litellm.llms.black_forest_labs.image_generation.transformation import ( @@ -355,3 +356,48 @@ class TestBlackForestLabsImageGenerationTransformation: config = get_black_forest_labs_image_generation_config("flux-pro-1.1") assert isinstance(config, BlackForestLabsImageGenerationConfig) + + +def _transform(payload: object) -> ImageResponse: + return BlackForestLabsImageGenerationConfig().transform_image_generation_response( + model="flux-pro-1.1", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + ("result", "expected_urls"), + [ + ({"sample": "https://bfl.example/a.png"}, ["https://bfl.example/a.png"]), + ( + ["https://bfl.example/a.png", {"url": "https://bfl.example/b.png"}, {"seed": 1}, 7], + ["https://bfl.example/a.png", "https://bfl.example/b.png"], + ), + ], +) +def test_transform_image_generation_response_reads_urls_from_the_result(result: object, expected_urls: list[str]): + response = _transform({"status": "Ready", "result": result}) + + assert [image.url for image in response.data] == expected_urls + + +@pytest.mark.parametrize("payload", [{}, {"result": None}, {"result": "https://bfl.example/a.png"}, {"result": []}]) +def test_transform_image_generation_response_without_a_url_is_a_provider_error(payload: dict[str, object]): + with pytest.raises(BlackForestLabsError) as exc_info: + _transform(payload) + + assert exc_info.value.status_code == 500 + + +@pytest.mark.parametrize("payload", [7, "https://bfl.example/a.png", [{"sample": "https://bfl.example/a.png"}]]) +def test_transform_image_generation_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "bfl.example" not in str(exc_info.value) diff --git a/tests/unit/llms/clarifai/__init__.py b/tests/unit/llms/clarifai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/clarifai/chat/__init__.py b/tests/unit/llms/clarifai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/clarifai/chat/test_transformation.py b/tests/unit/llms/clarifai/chat/test_transformation.py new file mode 100644 index 00000000000..e81b2a9558d --- /dev/null +++ b/tests/unit/llms/clarifai/chat/test_transformation.py @@ -0,0 +1,70 @@ +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.clarifai.chat.transformation import ClarifaiConfig +from litellm.llms.openai.common_utils import OpenAIError +from litellm.types.utils import ModelResponse + + +def _transform(raw_response: httpx.Response) -> ModelResponse: + return ClarifaiConfig().transform_response( + model="user.app.model", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_response_builds_the_model_response_from_the_body(): + response = _transform( + httpx.Response( + 200, + json={ + "id": "chatcmpl-1", + "created": 1700000000, + "model": "upstream-model", + "system_fingerprint": "fp_1", + "choices": [{"index": 0, "finish_reason": "length", "message": {"role": "assistant", "content": "Hi"}}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "vendor_field": {"kept": True}, + }, + ) + ) + + assert response.id == "chatcmpl-1" + assert response.created == 1700000000 + assert response.model == "clarifai/user.app.model" + assert response.system_fingerprint == "fp_1" + assert [(choice.finish_reason, choice.message.content) for choice in response.choices] == [("length", "Hi")] + assert response.usage.model_dump()["total_tokens"] == 5 + assert response.vendor_field == {"kept": True} + + +def test_transform_response_keeps_a_missing_model_unset(): + response = _transform(httpx.Response(200, json={"choices": [{"message": {"content": "Hi"}}]})) + + assert response.model is None + assert response.choices[0].message.content == "Hi" + + +@pytest.mark.parametrize("body", [b'["prompt-text"]', b'"prompt-text"', b"7", b"null"]) +def test_transform_response_rejects_a_body_that_is_not_an_object_without_echoing_it(body: bytes): + with pytest.raises(ValidationError) as exc_info: + _transform(httpx.Response(200, content=body)) + + assert "prompt-text" not in str(exc_info.value) + + +def test_transform_response_reports_an_unparseable_body_as_an_openai_error(): + with pytest.raises(OpenAIError, match="Failed to parse Clarifai response") as exc_info: + _transform(httpx.Response(502, content=b"bad gateway")) + + assert exc_info.value.status_code == 502 diff --git a/tests/unit/llms/claude_code/harness/test_transformation.py b/tests/unit/llms/claude_code/harness/test_transformation.py index 6348b4f6a3e..1ae4497ef18 100644 --- a/tests/unit/llms/claude_code/harness/test_transformation.py +++ b/tests/unit/llms/claude_code/harness/test_transformation.py @@ -318,6 +318,29 @@ def test_parse_skips_subagent_messages_and_maps_errors(): ] +def test_user_message_yields_only_its_tool_result_objects(): + line = { + "type": "user", + "message": { + "content": [ + "plain text", + None, + ["tool_result"], + {"type": "text", "text": "not a result"}, + {"type": "tool_result", "tool_use_id": "t1", "content": "ok"}, + {"type": "tool_result", "content": [{"type": "text", "text": "a"}, "b"], "is_error": 1}, + {"type": "tool_result", "tool_use_id": 7, "content": None, "extra": {"kept": "out"}}, + ] + }, + } + + assert ClaudeCodeHarnessConfig().transform_stream_line(line, ClaudeCodeStreamState()) == [ + ToolResult(id="t1", output="ok", is_error=False), + ToolResult(id="", output="a\nb", is_error=True), + ToolResult(id="7", output="", is_error=False), + ] + + def test_parse_thinking_and_mcp_tools(): state = ClaudeCodeStreamState() msg = { diff --git a/tests/unit/llms/codex/harness/test_transformation.py b/tests/unit/llms/codex/harness/test_transformation.py index 6371f74fba6..6df322db1ae 100644 --- a/tests/unit/llms/codex/harness/test_transformation.py +++ b/tests/unit/llms/codex/harness/test_transformation.py @@ -11,7 +11,7 @@ from pathlib import Path from typing import Optional import pytest -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from litellm.harness.context import SessionContext from litellm.harness.errors import HarnessError, HarnessInstallFailed, OptionsMismatch @@ -272,6 +272,36 @@ def test_parse_file_change_and_mcp_and_web_search(): assert events[0].name == "web_search" and events[0].input == {"query": "litellm"} +@pytest.mark.parametrize( + ("changes", "output"), + [ + ([{"path": "a.txt", "kind": "add"}, {"path": "b.txt", "kind": "update", "diff": "@@"}], "add a.txt\nupdate b.txt"), + ([{"path": "only-path.txt"}, {"kind": "delete"}, {}], "only-path.txt\ndelete\n"), + ([], ""), + (None, ""), + ], +) +def test_completed_file_change_lists_each_change(changes: object, output: str): + state = CodexStreamState(started={"i1"}) + item = {"id": "i1", "type": "file_change", "changes": changes, "status": "completed"} + + assert parse_event({"type": "item.completed", "item": item}, state) == [ + ToolResult(id="i1", output=output, is_error=False) + ] + + +@pytest.mark.parametrize( + "changes", + [["a.txt"], [{"path": "a.txt"}, None], "a.txt", {"a.txt": {"kind": "add"}}, 7], +) +def test_completed_file_change_rejects_changes_that_are_not_a_list_of_objects(changes: object): + state = CodexStreamState(started={"i1"}) + item = {"id": "i1", "type": "file_change", "changes": changes, "status": "completed"} + + with pytest.raises(ValidationError): + parse_event({"type": "item.completed", "item": item}, state) + + def test_parse_failed_command_is_error_and_unknown_events_ignored(): state = CodexStreamState() item = { diff --git a/tests/unit/llms/cohere/rerank/test_transformation.py b/tests/unit/llms/cohere/rerank/test_transformation.py new file mode 100644 index 00000000000..ac5acf08d0d --- /dev/null +++ b/tests/unit/llms/cohere/rerank/test_transformation.py @@ -0,0 +1,72 @@ +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.cohere.rerank.transformation import CohereRerankConfig +from litellm.types.rerank import RerankResponse + + +def _transform(payload: object) -> RerankResponse: + return CohereRerankConfig().transform_rerank_response( + model="rerank-english-v3.0", + raw_response=httpx.Response(200, json=payload), + model_response=RerankResponse(), + logging_obj=Mock(), + ) + + +def test_transform_rerank_response_keeps_the_cohere_payload(): + response = _transform( + { + "id": "rerank-1", + "results": [{"index": 1, "relevance_score": 0.9, "document": {"text": "sensitive-document"}}], + "meta": {"billed_units": {"search_units": 1}, "tokens": {"input_tokens": 4}}, + "warnings": ["ignored"], + } + ) + + assert response.model_dump() == { + "id": "rerank-1", + "results": [{"index": 1, "relevance_score": 0.9, "document": {"text": "sensitive-document"}}], + "meta": {"billed_units": {"search_units": 1}, "tokens": {"input_tokens": 4}}, + } + + +def test_transform_rerank_response_of_an_empty_object_has_no_results(): + assert _transform({}).model_dump() == {"id": None, "results": None, "meta": None} + + +@pytest.mark.parametrize("payload", [7, "sensitive-document", [{"id": "sensitive-document"}]]) +def test_transform_rerank_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "sensitive-document" not in str(exc_info.value) + + +@pytest.mark.parametrize("payload", [{"id": 7}, {"results": "not-a-list"}, {"results": [{"index": 0}]}, {"meta": []}]) +def test_transform_rerank_response_rejects_malformed_fields(payload: dict[str, object]): + with pytest.raises(ValidationError): + _transform(payload) + + +def test_map_cohere_rerank_params_returns_every_cohere_param(): + params = CohereRerankConfig().map_cohere_rerank_params( + non_default_params=None, + model="rerank-english-v3.0", + drop_params=False, + query="capital of france", + documents=["Paris", {"text": "Berlin"}], + top_n=1, + ) + + assert params == { + "query": "capital of france", + "documents": ["Paris", {"text": "Berlin"}], + "top_n": 1, + "rank_fields": None, + "return_documents": True, + "max_chunks_per_doc": None, + } diff --git a/tests/unit/llms/cometapi/image_generation/__init__.py b/tests/unit/llms/cometapi/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cometapi/image_generation/test_transformation.py b/tests/unit/llms/cometapi/image_generation/test_transformation.py new file mode 100644 index 00000000000..0a14b955339 --- /dev/null +++ b/tests/unit/llms/cometapi/image_generation/test_transformation.py @@ -0,0 +1,54 @@ +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.cometapi.image_generation.transformation import CometAPIImageGenerationConfig +from litellm.types.utils import ImageResponse + + +def _transform(payload: object) -> ImageResponse: + return CometAPIImageGenerationConfig().transform_image_generation_response( + model="dall-e-3", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=Mock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_each_image_object(): + response = _transform({"created": 1, "data": [{"url": "https://img.cometapi.com/a.png"}, {"b64_json": "QUJD"}, {}]}) + + assert [(image.url, image.b64_json) for image in response.data] == [ + ("https://img.cometapi.com/a.png", None), + (None, "QUJD"), + (None, None), + ] + + +@pytest.mark.parametrize("payload", [{}, {"created": 1}, {"data": []}, {"data": ""}, {"data": {}}, []]) +def test_transform_image_generation_response_without_images_has_no_data(payload: object): + assert _transform(payload).data == [] + + +@pytest.mark.parametrize( + "payload", + [ + ["data", "https://img.cometapi.com/a.png"], + {"data": None}, + {"data": "https://img.cometapi.com/a.png"}, + {"data": {"url": "https://img.cometapi.com/a.png"}}, + {"data": ["https://img.cometapi.com/a.png"]}, + {"data": [{"url": "https://img.cometapi.com/a.png"}, 7]}, + ], +) +def test_transform_image_generation_response_rejects_malformed_payloads(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "img.cometapi.com" not in str(exc_info.value) diff --git a/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py b/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py index 4466e5b8767..a913c79af31 100644 --- a/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock import httpx import pytest +from pydantic import ValidationError from litellm.llms.dashscope.common_utils import DashScopeError @@ -347,3 +348,89 @@ class TestProviderConfigManagerDispatch: present_version_params=[], ) assert isinstance(cfg, DashScopeRerankConfig) + + +def _transform(payload: object, status_code: int = 200) -> RerankResponse: + return DashScopeRerankConfig().transform_rerank_response( + model="qwen3-rerank", + raw_response=httpx.Response(status_code, json=payload), + model_response=RerankResponse(), + logging_obj=MagicMock(), + ) + + +@pytest.mark.parametrize( + ("usage", "expected_total_tokens"), + [ + ({"total_tokens": 79}, 79), + ({}, None), + (None, None), + ], +) +def test_transform_rerank_response_reads_total_tokens_from_usage(usage: object, expected_total_tokens: int | None): + response = _transform({"id": "rerank-1", "results": [{"index": 0, "relevance_score": 0.5}], "usage": usage}) + + assert response.meta == { + "billed_units": {"total_tokens": expected_total_tokens}, + "tokens": {"input_tokens": expected_total_tokens}, + } + + +def test_transform_rerank_response_keeps_provider_id_and_drops_unknown_result_fields(): + response = _transform( + { + "id": "rerank-1", + "results": [{"index": 2, "relevance_score": 0.25, "document": {"text": "doc", "extra": 1}, "extra": 2}], + } + ) + + assert response.id == "rerank-1" + assert response.results == [{"index": 2, "relevance_score": 0.25, "document": {"text": "doc"}}] + + +def test_transform_rerank_response_empty_results_list_yields_no_results(): + assert _transform({"id": "rerank-1", "results": []}).results == [] + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + {"results": 7}, + {"results": ["not an object"]}, + {"results": [{"index": 0, "relevance_score": 0.5}], "usage": "seventy nine"}, + {"results": [{"index": 0, "relevance_score": 0.5}], "usage": {"total_tokens": 1.5}}, + {"results": [{"index": 0, "relevance_score": 0.5}], "id": 7}, + ], +) +def test_transform_rerank_response_rejects_malformed_payloads(payload: object): + with pytest.raises(ValidationError): + _transform(payload) + + +def test_transform_rerank_response_result_without_index_raises_key_error(): + with pytest.raises(KeyError, match="index"): + _transform({"results": [{"relevance_score": 0.5}]}) + + +def test_transform_rerank_response_error_envelope_raises_with_the_provider_status_and_message(): + with pytest.raises(DashScopeError) as exc_info: + _transform({"code": "Throttling", "message": "slow down"}, status_code=429) + + assert exc_info.value.status_code == 429 + assert exc_info.value.message == "slow down" + + +def test_transform_rerank_response_without_results_raises_dashscope_error_naming_the_body(): + with pytest.raises(DashScopeError) as exc_info: + _transform({"id": "rerank-1"}, status_code=502) + + assert exc_info.value.status_code == 502 + assert exc_info.value.message == "No results in DashScope rerank response: {'id': 'rerank-1'}" + + +def test_transform_rerank_response_shape_errors_do_not_echo_the_payload(): + with pytest.raises(ValidationError) as exc_info: + _transform({"results": [["leaked document text"]]}) + + assert "leaked document text" not in str(exc_info.value) diff --git a/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py b/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py index 7c1f5256deb..366413df1c8 100644 --- a/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py +++ b/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py @@ -3,6 +3,7 @@ import os import pathlib from unittest.mock import MagicMock +import httpx import pytest @@ -525,3 +526,82 @@ def test_reconstruct_diarized_transcript_multiple_speaker_changes(): assert "Hello" in result assert "back" in result assert "Thanks" in result + + +def _deepgram_payload(alternative: dict[str, object], channel_fields: dict[str, object]) -> dict[str, object]: + return { + "metadata": {"duration": 2.5}, + "results": {"channels": [{"alternatives": [alternative], **channel_fields}]}, + } + + +def _transform_deepgram_response(payload: object) -> TranscriptionResponse: + return DeepgramAudioTranscriptionConfig().transform_audio_transcription_response( + httpx.Response(200, json=payload) + ) + + +@pytest.mark.parametrize( + ("channel_fields", "expected_language"), + [ + ({}, "en"), + ({"detected_language": None}, "en"), + ({"detected_language": ""}, "en"), + ({"detected_language": "fr"}, "fr"), + ({"detected_language": "de", "language_confidence": 0.98}, "de"), + ({"detected_language": ["es", "en"]}, ["es", "en"]), + ], +) +def test_transform_response_reports_the_detected_language_or_english( + channel_fields: dict[str, object], expected_language: object +): + response = _transform_deepgram_response(_deepgram_payload({"transcript": "bonjour"}, channel_fields)) + + assert response.text == "bonjour" + assert response["language"] == expected_language + assert response["duration"] == 2.5 + assert "words" not in response + + +@pytest.mark.parametrize( + ("words", "expected"), + [ + ([], []), + ("", []), + ({}, []), + ( + [{"word": "hello", "start": 0.0, "end": 0.5, "confidence": 0.9}], + [{"word": "hello", "start": 0.0, "end": 0.5}], + ), + ( + [{"word": None, "start": "0.1", "end": [2]}, {"word": "b", "start": 1, "end": 2}], + [{"word": None, "start": "0.1", "end": [2]}, {"word": "b", "start": 1, "end": 2}], + ), + ], +) +def test_transform_response_maps_words_to_openai_word_timestamps(words: object, expected: list[dict[str, object]]): + payload = _deepgram_payload({"transcript": "hello", "words": words}, {}) + + response = _transform_deepgram_response(payload) + + assert response["words"] == expected + assert response._hidden_params == payload + + +@pytest.mark.parametrize( + "words", + [ + "not a list", + ["not an object"], + [{"word": "hello", "start": 0.0, "end": 0.5}, None], + [{"word": "hello", "start": 0.0, "end": 0.5}, ["nested"]], + [{"word": "hello", "start": 0.0}], + ], +) +def test_transform_response_wraps_malformed_words_with_the_raw_body(words: object): + raw_response = httpx.Response(200, json=_deepgram_payload({"transcript": "hello", "words": words}, {})) + + with pytest.raises(ValueError, match="Error transforming Deepgram response: ") as exc_info: + DeepgramAudioTranscriptionConfig().transform_audio_transcription_response(raw_response) + + assert str(exc_info.value).endswith(f"\nResponse: {raw_response.text}") diff --git a/tests/unit/llms/e2b/__init__.py b/tests/unit/llms/e2b/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/e2b/sandbox/__init__.py b/tests/unit/llms/e2b/sandbox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/e2b/sandbox/test_transformation.py b/tests/unit/llms/e2b/sandbox/test_transformation.py new file mode 100644 index 00000000000..44c0626bc27 --- /dev/null +++ b/tests/unit/llms/e2b/sandbox/test_transformation.py @@ -0,0 +1,121 @@ +import json + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.e2b.sandbox.transformation import E2BSandboxConfig + + +def _client_answering(body: bytes) -> AsyncHTTPHandler: + return AsyncHTTPHandler(transport=httpx.MockTransport(lambda request: httpx.Response(200, content=body))) + + +def _lines(*messages: object) -> list[str]: + return [json.dumps(message) for message in messages] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("created", "expected_domain", "expected_tokens"), + [ + ( + {"sandboxID": "sbx_1", "domain": "eu.e2b.app", "envdAccessToken": "envd", "trafficAccessToken": "traffic"}, + "eu.e2b.app", + ("envd", "traffic"), + ), + ({"sandboxID": "sbx_1"}, "e2b.app", (None, None)), + ({"sandboxID": "sbx_1", "domain": None, "envdAccessToken": "envd"}, "e2b.app", ("envd", None)), + ({"sandboxID": "sbx_1", "domain": ""}, "e2b.app", (None, None)), + ], +) +async def test_acreate_sandbox_builds_the_handle_from_the_create_response( + created: dict[str, object], expected_domain: str, expected_tokens: tuple[str | None, str | None] +): + handle = await E2BSandboxConfig().acreate_sandbox( + api_key="e2b_key", client=_client_answering(json.dumps(created).encode()) + ) + + assert (handle.id, handle.provider, handle.domain) == ("sbx_1", "e2b", expected_domain) + assert handle._hidden_params == { + "envd_access_token": expected_tokens[0], + "traffic_access_token": expected_tokens[1], + "api_key": "e2b_key", + "api_base": "https://api.e2b.app", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [b'["secret-token"]', b'"secret-token"', b"7", b"null"]) +async def test_acreate_sandbox_rejects_a_create_response_that_is_not_an_object_without_echoing_it(body: bytes): + with pytest.raises(ValidationError) as exc_info: + await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=_client_answering(body)) + + assert "secret-token" not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_acreate_sandbox_requires_a_sandbox_id(): + with pytest.raises(KeyError, match="sandboxID"): + await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=_client_answering(b'{"domain": "e2b.app"}')) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("created", [{"sandboxID": 5}, {"sandboxID": "sbx_1", "domain": 5}]) +async def test_acreate_sandbox_rejects_non_string_handle_fields(created: dict[str, object]): + with pytest.raises(ValidationError, match="ContainerHandle"): + await E2BSandboxConfig().acreate_sandbox( + api_key="e2b_key", client=_client_answering(json.dumps(created).encode()) + ) + + +def test_parse_lines_maps_every_message_type_onto_the_result(): + result = E2BSandboxConfig._parse_lines( + [ + *_lines( + {"type": "stdout", "text": "1\n", "timestamp": 1}, + {"type": "stderr", "text": "warn\n"}, + {"type": "stdout"}, + {"type": "result", "png": "BASE64", "is_main_result": True}, + {"type": "error", "name": "ValueError", "value": "bad", "traceback": "tb", "ignored": 1}, + {"type": "number_of_executions", "execution_count": 3}, + {"type": "stdout", "text": "2\n"}, + None, + ), + "", + "not json", + ] + ) + + assert result.model_dump() == { + "stdout": "1\n2\n", + "stderr": "warn\n", + "results": [{"png": "BASE64", "is_main_result": True}], + "error": {"name": "ValueError", "value": "bad", "traceback": "tb"}, + "execution_count": 3, + "object": "code_execution", + } + + +def test_parse_lines_rejects_an_execution_count_that_is_not_a_number(): + with pytest.raises(ValidationError, match="execution_count"): + E2BSandboxConfig._parse_lines(_lines({"type": "number_of_executions", "execution_count": "many"})) + + +@pytest.mark.parametrize( + "message", + [ + ["secret-output"], + "secret-output", + 7, + False, + {"type": "stdout", "text": ["secret-output"]}, + {"type": "stderr", "text": {"secret-output": 1}}, + ], +) +def test_parse_lines_rejects_malformed_messages_without_echoing_them(message: object): + with pytest.raises(ValidationError) as exc_info: + E2BSandboxConfig._parse_lines(_lines({"type": "stdout", "text": "ok"}, message)) + + assert "secret-output" not in str(exc_info.value) diff --git a/tests/unit/llms/elevenlabs/audio_transcription/__init__.py b/tests/unit/llms/elevenlabs/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py b/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py new file mode 100644 index 00000000000..52a10f0a327 --- /dev/null +++ b/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py @@ -0,0 +1,100 @@ +import httpx +import pytest + +from litellm.llms.elevenlabs.audio_transcription.transformation import ElevenLabsAudioTranscriptionConfig +from litellm.types.utils import TranscriptionResponse + + +def _transform(payload: object) -> TranscriptionResponse: + return ElevenLabsAudioTranscriptionConfig().transform_audio_transcription_response( + raw_response=httpx.Response(200, json=payload) + ) + + +def test_transform_audio_transcription_response_keeps_only_spoken_words(): + payload = { + "language_code": "en", + "text": "Hello world", + "words": [ + {"type": "word", "text": "Hello", "start": 0.0, "end": 0.4, "speaker_id": "speaker_0"}, + {"type": "spacing", "text": " ", "start": 0.4, "end": 0.5}, + {"type": "audio_event", "text": "(laughter)", "start": 0.5, "end": 0.9}, + {"type": "word", "text": "world", "start": 0.9, "end": 1.3}, + ], + } + + response = _transform(payload) + + assert response.text == "Hello world" + assert response["task"] == "transcribe" + assert response["language"] == "en" + assert response["words"] == [ + {"word": "Hello", "start": 0.0, "end": 0.4}, + {"word": "world", "start": 0.9, "end": 1.3}, + ] + assert response._hidden_params == payload + + +@pytest.mark.parametrize( + ("word", "expected"), + [ + ({"type": "word"}, [{"word": "", "start": 0, "end": 0}]), + ({"type": "word", "text": None, "start": None, "end": None}, [{"word": None, "start": None, "end": None}]), + ({"type": "word", "text": 7, "start": "0.1", "end": [2]}, [{"word": 7, "start": "0.1", "end": [2]}]), + ({"text": "untyped"}, []), + ({"type": None, "text": "untyped"}, []), + ({}, []), + ], +) +def test_transform_audio_transcription_response_maps_one_word( + word: dict[str, object], expected: list[dict[str, object]] +): + assert _transform({"text": "t", "words": [word]})["words"] == expected + + +@pytest.mark.parametrize("words", [[], "", {}]) +def test_transform_audio_transcription_response_with_empty_words_has_empty_word_list(words: object): + assert _transform({"text": "t", "words": words})["words"] == [] + + +@pytest.mark.parametrize( + ("payload", "expected_text", "expected_language"), + [ + ({}, "", "unknown"), + ({"text": None, "language_code": None}, None, None), + ({"text": "bonjour", "language_code": "fr"}, "bonjour", "fr"), + ({"text": "hola", "language_code": ["es"]}, "hola", ["es"]), + ], +) +def test_transform_audio_transcription_response_without_words_key_has_no_word_list( + payload: dict[str, object], expected_text: str | None, expected_language: object +): + response = _transform(payload) + + assert response.text == expected_text + assert response["language"] == expected_language + assert "words" not in response + assert response._hidden_params == payload + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + "plain text", + 7, + {"text": ["not", "text"]}, + {"text": "t", "words": None}, + {"text": "t", "words": 7}, + {"text": "t", "words": "not a list"}, + {"text": "t", "words": ["not an object"]}, + {"text": "t", "words": [{"type": "word", "text": "ok"}, None]}, + ], +) +def test_transform_audio_transcription_response_wraps_malformed_payloads_with_the_raw_body(payload: object): + raw_response = httpx.Response(200, json=payload) + + with pytest.raises(ValueError, match="Error transforming ElevenLabs response: ") as exc_info: + ElevenLabsAudioTranscriptionConfig().transform_audio_transcription_response(raw_response=raw_response) + + assert str(exc_info.value).endswith(f"\nResponse: {raw_response.text}") diff --git a/tests/unit/llms/fal_ai/image_generation/test_flux_pro_v11_ultra_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_flux_pro_v11_ultra_transformation.py new file mode 100644 index 00000000000..960d927851e --- /dev/null +++ b/tests/unit/llms/fal_ai/image_generation/test_flux_pro_v11_ultra_transformation.py @@ -0,0 +1,56 @@ +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.fal_ai.image_generation.flux_pro_v11_ultra_transformation import FalAIFluxProV11UltraConfig +from litellm.types.utils import ImageResponse + + +def _transform(payload: object) -> ImageResponse: + return FalAIFluxProV11UltraConfig().transform_image_generation_response( + model="fal-ai/flux-pro/v1.1-ultra", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=Mock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_images_and_metadata(): + response = _transform( + { + "images": [{"url": "https://fal.media/a.png", "width": 2752, "height": 1536}, "https://fal.media/b.png"], + "seed": 42, + "timings": {"inference": 2.5}, + "has_nsfw_concepts": [False, False], + } + ) + + assert [image.url for image in response.data] == ["https://fal.media/a.png", "https://fal.media/b.png"] + assert response.data[0].provider_specific_fields == {"width": 2752, "height": 1536} + assert response._hidden_params["seed"] == 42 + assert response._hidden_params["timings"] == {"inference": 2.5} + assert response._hidden_params["has_nsfw_concepts"] == [False, False] + + +@pytest.mark.parametrize("payload", [{}, {"images": None}, {"images": []}]) +def test_transform_image_generation_response_without_images_is_empty(payload: dict[str, object]): + response = _transform(payload) + + assert response.data == [] + assert "seed" not in response._hidden_params + assert "timings" not in response._hidden_params + assert "has_nsfw_concepts" not in response._hidden_params + + +@pytest.mark.parametrize("payload", [7, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}]]) +def test_transform_image_generation_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "fal.media" not in str(exc_info.value) diff --git a/tests/unit/llms/fal_ai/image_generation/test_ideogram_v3_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_ideogram_v3_transformation.py new file mode 100644 index 00000000000..716440004cd --- /dev/null +++ b/tests/unit/llms/fal_ai/image_generation/test_ideogram_v3_transformation.py @@ -0,0 +1,71 @@ +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.fal_ai.image_generation.ideogram_v3_transformation import FalAIIdeogramV3Config +from litellm.types.utils import ImageResponse + + +def _transform(payload: object) -> ImageResponse: + return FalAIIdeogramV3Config().transform_image_generation_response( + model="fal-ai/ideogram/v3", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=Mock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_files_and_seed(): + response = _transform( + {"images": [{"url": "https://fal.media/a.png", "file_name": "a.png"}, "https://fal.media/b.png"], "seed": 42} + ) + + assert [image.url for image in response.data] == ["https://fal.media/a.png", "https://fal.media/b.png"] + assert [image.b64_json for image in response.data] == [None, None] + assert response._hidden_params["seed"] == 42 + + +@pytest.mark.parametrize("payload", [{}, {"images": None}, {"images": "https://fal.media/a.png"}]) +def test_transform_image_generation_response_without_an_image_list_is_empty(payload: dict[str, object]): + response = _transform(payload) + + assert response.data == [] + assert "seed" not in response._hidden_params + + +@pytest.mark.parametrize("payload", [7, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}]]) +def test_transform_image_generation_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "fal.media" not in str(exc_info.value) + + +@pytest.mark.parametrize( + ("size", "expected"), + [ + ("1024x1024", "square_hd"), + (" 1536x1024 ", "landscape_16_9"), + ("640x480", {"width": 640, "height": 480}), + ("wide", "square_hd"), + ("axb", "square_hd"), + ({"width": 640, "height": 480, "unit": "px"}, {"width": 640, "height": 480}), + ({"width": "640"}, {"width": "640"}), + (512, 512), + ], +) +def test_map_openai_params_translates_size_to_image_size(size: object, expected: object): + optional_params = FalAIIdeogramV3Config().map_openai_params( + non_default_params={"size": size, "n": 2, "response_format": "url"}, + optional_params={}, + model="fal-ai/ideogram/v3", + drop_params=False, + ) + + assert optional_params == {"image_size": expected, "num_images": 2} diff --git a/tests/unit/llms/fal_ai/image_generation/test_recraft_v3_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_recraft_v3_transformation.py new file mode 100644 index 00000000000..09591fe24d5 --- /dev/null +++ b/tests/unit/llms/fal_ai/image_generation/test_recraft_v3_transformation.py @@ -0,0 +1,64 @@ +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.fal_ai.image_generation.recraft_v3_transformation import FalAIRecraftV3Config +from litellm.types.utils import ImageResponse + + +def _transform(payload: object) -> ImageResponse: + return FalAIRecraftV3Config().transform_image_generation_response( + model="fal-ai/recraft/v3/text-to-image", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=Mock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_file_objects_and_bare_urls(): + response = _transform( + {"images": [{"url": "https://fal.media/a.png", "content_type": "image/png"}, "https://fal.media/b.png", 7]} + ) + + assert [image.url for image in response.data] == ["https://fal.media/a.png", "https://fal.media/b.png"] + assert [image.b64_json for image in response.data] == [None, None] + + +@pytest.mark.parametrize("payload", [{}, {"images": None}, {"images": "https://fal.media/a.png"}]) +def test_transform_image_generation_response_without_an_image_list_is_empty(payload: dict[str, object]): + assert _transform(payload).data == [] + + +@pytest.mark.parametrize("payload", [7, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}]]) +def test_transform_image_generation_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "fal.media" not in str(exc_info.value) + + +@pytest.mark.parametrize( + ("size", "expected"), + [ + ("1024x1024", "square_hd"), + ("576x1024", "portrait_16_9"), + ("640x480", {"width": 640, "height": 480}), + ("wide", "square_hd"), + ("axb", "square_hd"), + ], +) +def test_map_openai_params_translates_size_to_image_size(size: str, expected: object): + optional_params = FalAIRecraftV3Config().map_openai_params( + non_default_params={"size": size, "n": 2, "response_format": "url"}, + optional_params={}, + model="fal-ai/recraft/v3/text-to-image", + drop_params=False, + ) + + assert optional_params == {"image_size": expected} diff --git a/tests/unit/llms/fal_ai/image_generation/test_stable_diffusion_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_stable_diffusion_transformation.py new file mode 100644 index 00000000000..8f112374efd --- /dev/null +++ b/tests/unit/llms/fal_ai/image_generation/test_stable_diffusion_transformation.py @@ -0,0 +1,76 @@ +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.fal_ai.image_generation.stable_diffusion_transformation import FalAIStableDiffusionConfig +from litellm.types.utils import ImageResponse + + +def _transform(payload: object) -> ImageResponse: + return FalAIStableDiffusionConfig().transform_image_generation_response( + model="fal-ai/stable-diffusion-v35-medium", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=Mock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_images_and_metadata(): + response = _transform( + { + "images": [{"url": "https://fal.media/a.png", "width": 1024}, "https://fal.media/b.png", 7], + "seed": 42, + "timings": {"inference": 2.5}, + "has_nsfw_concepts": [False, False], + } + ) + + assert [image.url for image in response.data] == ["https://fal.media/a.png", "https://fal.media/b.png"] + assert response._hidden_params["seed"] == 42 + assert response._hidden_params["timings"] == {"inference": 2.5} + assert response._hidden_params["has_nsfw_concepts"] == [False, False] + + +@pytest.mark.parametrize("payload", [{}, {"images": None}, {"images": "https://fal.media/a.png"}]) +def test_transform_image_generation_response_without_an_image_list_is_empty(payload: dict[str, object]): + response = _transform(payload) + + assert response.data == [] + assert "seed" not in response._hidden_params + assert "timings" not in response._hidden_params + assert "has_nsfw_concepts" not in response._hidden_params + + +@pytest.mark.parametrize("payload", [7, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}]]) +def test_transform_image_generation_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "fal.media" not in str(exc_info.value) + + +@pytest.mark.parametrize( + ("size", "expected"), + [ + ("1024x1024", "square_hd"), + ("1024x576", "landscape_16_9"), + ("640x480", {"width": 640, "height": 480}), + ("wide", "landscape_4_3"), + ("axb", "landscape_4_3"), + ], +) +def test_map_openai_params_translates_size_to_image_size(size: str, expected: object): + optional_params = FalAIStableDiffusionConfig().map_openai_params( + non_default_params={"size": size, "n": 2, "response_format": "b64_json"}, + optional_params={}, + model="fal-ai/stable-diffusion-v35-medium", + drop_params=False, + ) + + assert optional_params == {"image_size": expected, "num_images": 2, "output_format": "jpeg"} diff --git a/tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py b/tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py index a03b7708238..a3d1879f97c 100644 --- a/tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py +++ b/tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock import httpx import pytest +from pydantic import ValidationError from litellm.llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig from litellm.types.rerank import RerankResponse @@ -341,3 +342,68 @@ class TestFireworksAIRerankTransform: assert headers["Authorization"] == "Bearer test-api-key" assert headers["Content-Type"] == "application/json" + + +def _transform(payload: object) -> RerankResponse: + return FireworksAIRerankConfig().transform_rerank_response( + model="fireworks_ai/fireworks/qwen3-reranker-8b", + raw_response=httpx.Response(200, json=payload), + model_response=RerankResponse(), + logging_obj=MagicMock(), + ) + + +@pytest.mark.parametrize( + ("result", "expected"), + [ + ({"index": "1", "relevance_score": "0.5"}, {"index": 1, "relevance_score": 0.5}), + ({"index": 2, "relevance_score": 1}, {"index": 2, "relevance_score": 1.0}), + ], +) +def test_transform_rerank_response_converts_index_and_score_with_int_and_float( + result: dict[str, object], expected: dict[str, object] +): + assert _transform({"id": "rerank-1", "data": [result]}).results == [expected] + + +@pytest.mark.parametrize("usage_fields", [{}, {"usage": {}}]) +def test_transform_rerank_response_without_usage_counters_reports_zero_tokens(usage_fields: dict[str, object]): + response = _transform({"id": "rerank-1", "data": [{"index": 0, "relevance_score": 0.5}], **usage_fields}) + + assert response.meta == {"billed_units": {"search_units": 0}, "tokens": {"input_tokens": 0, "output_tokens": 0}} + + +def test_transform_rerank_response_keeps_provider_id_and_falls_back_to_results_key(): + response = _transform({"id": "rerank-1", "data": [], "results": [{"index": 3, "relevance_score": 0.25}]}) + + assert response.id == "rerank-1" + assert response.results == [{"index": 3, "relevance_score": 0.25}] + + +def test_transform_rerank_response_empty_results_list_yields_no_results(): + assert _transform({"id": "rerank-1", "results": []}).results == [] + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + {"data": [{"index": 0, "relevance_score": 0.5}], "usage": None}, + {"data": [{"index": 0, "relevance_score": 0.5}], "usage": {"total_tokens": 1.5}}, + {"data": 7}, + {"data": ["not an object"]}, + {"data": [{"index": None, "relevance_score": 0.5}]}, + {"data": [{"index": 0, "relevance_score": [0.5]}]}, + {"id": 7, "data": [{"index": 0, "relevance_score": 0.5}]}, + ], +) +def test_transform_rerank_response_rejects_malformed_payloads(payload: object): + with pytest.raises(ValidationError): + _transform(payload) + + +def test_transform_rerank_response_shape_errors_do_not_echo_the_payload(): + with pytest.raises(ValidationError) as exc_info: + _transform({"data": [{"index": {"leaked": "document text"}, "relevance_score": 0.5}]}) + + assert "document text" not in str(exc_info.value) diff --git a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py index d9509856759..2a69926414a 100644 --- a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py @@ -1,4 +1,8 @@ +from unittest.mock import Mock + import httpx +import pytest +from pydantic import ValidationError from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.gemini.image_generation.transformation import GoogleImageGenConfig @@ -419,3 +423,47 @@ def test_gemini_image_generation_response_without_grounding_has_no_web_search_re ) assert getattr(result.usage, "web_search_requests", None) is None + + +def _transform_response(model: str, payload: object) -> ImageResponse: + return GoogleImageGenConfig().transform_image_generation_response( + model=model, + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(data=[]), + logging_obj=Mock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize("predictions", [[], "", {}]) +def test_imagen_generation_response_with_empty_predictions_has_no_images(predictions: object): + assert _transform_response("gemini/imagen-4.0-generate-001", {"predictions": predictions}).data == [] + + +def test_imagen_generation_response_keeps_predictions_without_image_bytes(): + result = _transform_response( + "gemini/imagen-4.0-generate-001", + {"predictions": [{"bytesBase64Encoded": "first-image"}, {"mimeType": "image/png"}]}, + ) + + assert [image.b64_json for image in result.data or []] == ["first-image", None] + + +@pytest.mark.parametrize( + ("model", "payload"), + [ + ("gemini/imagen-4.0-generate-001", {"predictions": None}), + ("gemini/imagen-4.0-generate-001", {"predictions": "not a list"}), + ("gemini/imagen-4.0-generate-001", {"predictions": [{"bytesBase64Encoded": "a"}, "not an object"]}), + ("gemini-3.1-flash-image-preview", {"usageMetadata": "not an object"}), + ("gemini-3.1-flash-image-preview", {"usageMetadata": ["not an object"]}), + ], +) +def test_image_generation_response_rejects_malformed_payloads_without_echoing_them(model: str, payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform_response(model, payload) + + assert "input_value" not in str(exc_info.value) diff --git a/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py b/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py index 01fe66ca4c7..e9c65cbb591 100644 --- a/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py +++ b/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py @@ -12,6 +12,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest +from pydantic import ValidationError from litellm import LlmProviders from litellm.llms.gigachat.embedding.transformation import ( @@ -339,4 +340,64 @@ class TestGetErrorClass: ) assert isinstance(error, GigaChatEmbeddingError) assert error.status_code == 400 - assert error.message == "embedding failed" \ No newline at end of file + assert error.message == "embedding failed" + + +def _transform(raw_response: httpx.Response) -> EmbeddingResponse: + return GigaChatEmbeddingConfig().transform_embedding_response( + model="Embeddings", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="test-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def test_transform_embedding_response_moves_per_item_usage_into_the_response_usage(): + response = _transform( + httpx.Response( + 200, + json={ + "object": "list", + "model": "Embeddings", + "data": [ + {"object": "embedding", "index": 0, "embedding": [0.1, 2], "usage": {"prompt_tokens": 3}}, + {"object": "embedding", "index": 1, "embedding": [0.5], "usage": {"prompt_tokens": 4}}, + ], + "unknown": "ignored", + }, + ) + ) + + assert response.model_dump() == { + "model": "Embeddings", + "data": [ + {"object": "embedding", "index": 0, "embedding": [0.1, 2]}, + {"object": "embedding", "index": 1, "embedding": [0.5]}, + ], + "object": "list", + "usage": { + "completion_tokens": 0, + "prompt_tokens": 7, + "total_tokens": 7, + "completion_tokens_details": None, + "prompt_tokens_details": None, + }, + } + + +@pytest.mark.parametrize("body", [b"null", b"[]", b"7"]) +def test_transform_embedding_response_body_that_is_not_an_object_raises_type_error(body: bytes): + with pytest.raises(TypeError): + _transform(httpx.Response(200, content=body)) + + +def test_transform_embedding_response_invalid_envelope_field_is_reported_by_name(): + with pytest.raises(ValidationError) as exc_info: + _transform(httpx.Response(200, json={"data": [], "model": 5})) + + assert exc_info.value.title == "EmbeddingResponse" + assert [error["loc"] for error in exc_info.value.errors()] == [("model",)] diff --git a/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py b/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py index f43e2e4d1cb..83c0cb2fec2 100644 --- a/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py +++ b/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py @@ -1,6 +1,8 @@ from unittest.mock import MagicMock, patch +import httpx import pytest +from pydantic import ValidationError from litellm.exceptions import AuthenticationError @@ -8,6 +10,7 @@ from litellm.llms.github_copilot.embedding.transformation import ( GithubCopilotEmbeddingConfig, ) from litellm.llms.github_copilot.common_utils import GetAPIKeyError +from litellm.types.utils import EmbeddingResponse def test_github_copilot_embedding_config_validate_environment(): @@ -201,3 +204,57 @@ def test_github_copilot_embedding_config_transform_response(): assert len(response.data) == 1 assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] assert response.model == "text-embedding-3-small" + + +def _transform(raw_response: httpx.Response) -> EmbeddingResponse: + return GithubCopilotEmbeddingConfig().transform_embedding_response( + model="github_copilot/text-embedding-3-small", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="test-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def test_transform_embedding_response_keeps_the_openai_envelope(): + response = _transform( + httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 5, "total_tokens": 5}, + "unknown": "ignored", + }, + ) + ) + + assert response.model_dump() == { + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "object": "list", + "usage": { + "completion_tokens": 0, + "prompt_tokens": 5, + "total_tokens": 5, + "completion_tokens_details": None, + "prompt_tokens_details": None, + }, + } + + +@pytest.mark.parametrize("body", [b"null", b"[]", b"7", b'"leaked payload text"', b'[{"data": "leaked payload text"}]']) +def test_transform_embedding_response_rejects_a_body_that_is_not_an_object(body: bytes): + with pytest.raises(ValidationError) as exc_info: + _transform(httpx.Response(200, content=body)) + + assert "leaked payload text" not in str(exc_info.value) + + +def test_transform_embedding_response_object_without_data_is_an_invalid_response_object(): + with pytest.raises(Exception, match="Invalid response object"): + _transform(httpx.Response(200, json={"model": "text-embedding-3-small"})) diff --git a/tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py b/tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py index 38355f32da1..44e8145ff33 100644 --- a/tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py +++ b/tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py @@ -1,9 +1,12 @@ from unittest.mock import MagicMock +import httpx import pytest +from pydantic import ValidationError from litellm.llms.jina_ai.embedding.transformation import JinaAIEmbeddingConfig +from litellm.types.utils import EmbeddingResponse JINA_KEY_ENV_NAMES = ("JINA_AI_API_KEY", "JINA_API_KEY", "JINA_AI_TOKEN") @@ -131,3 +134,61 @@ class TestJinaAIEmbeddingTransform: "input": expected_input, } assert result == expected_result + + +def _transform(raw_response: httpx.Response) -> EmbeddingResponse: + return JinaAIEmbeddingConfig().transform_embedding_response( + model="jina-embeddings-v3", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="test-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def test_transform_embedding_response_builds_the_response_from_the_body(): + response = _transform( + httpx.Response( + 200, + json={ + "model": "jina-embeddings-v3", + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 2]}], + "usage": {"prompt_tokens": 3, "total_tokens": 5}, + "unknown": "ignored", + }, + ) + ) + + assert response.model_dump() == { + "model": "jina-embeddings-v3", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 2]}], + "object": "list", + "usage": { + "completion_tokens": 0, + "prompt_tokens": 3, + "total_tokens": 5, + "completion_tokens_details": None, + "prompt_tokens_details": None, + }, + } + + +@pytest.mark.parametrize("body", [b"null", b"7", b'["leaked payload text"]', b'"leaked payload text"']) +def test_transform_embedding_response_rejects_a_body_that_is_not_an_object(body: bytes): + with pytest.raises(ValidationError) as exc_info: + _transform(httpx.Response(200, content=body)) + + assert "leaked payload text" not in str(exc_info.value) + + +@pytest.mark.parametrize("field", ["model", "data", "usage"]) +def test_transform_embedding_response_invalid_field_is_reported_by_name(field: str): + with pytest.raises(ValidationError) as exc_info: + _transform(httpx.Response(200, json={field: 5})) + + assert exc_info.value.title == "EmbeddingResponse" + assert [error["loc"] for error in exc_info.value.errors()] == [(field,)] diff --git a/tests/unit/llms/jina_ai/rerank/__init__.py b/tests/unit/llms/jina_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/jina_ai/rerank/test_transformation.py b/tests/unit/llms/jina_ai/rerank/test_transformation.py new file mode 100644 index 00000000000..9ec02a4548b --- /dev/null +++ b/tests/unit/llms/jina_ai/rerank/test_transformation.py @@ -0,0 +1,111 @@ +from unittest.mock import MagicMock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.jina_ai.rerank.transformation import JinaAIRerankConfig +from litellm.types.rerank import RerankResponse + + +def _transform(payload: object, status_code: int = 200) -> RerankResponse: + return JinaAIRerankConfig().transform_rerank_response( + model="jina-reranker-v2-base-multilingual", + raw_response=httpx.Response(status_code, json=payload), + model_response=RerankResponse(), + logging_obj=MagicMock(), + ) + + +def test_transform_rerank_response_maps_results_and_usage(): + response = _transform( + { + "id": "rerank-1", + "results": [ + {"index": 1, "relevance_score": 0.72, "document": "hello"}, + {"index": 0, "relevance_score": 0.25, "document": {"text": "world", "extra": 1}}, + {"index": 2, "relevance_score": 0, "extra": True}, + ], + "usage": {"total_tokens": 21, "prompt_tokens": 21}, + } + ) + + assert response.id == "rerank-1" + assert response.results == [ + {"index": 1, "relevance_score": 0.72, "document": {"text": "hello"}}, + {"index": 0, "relevance_score": 0.25, "document": {"text": "world"}}, + {"index": 2, "relevance_score": 0.0}, + ] + assert response.meta == {"billed_units": {"total_tokens": 21}, "tokens": {}} + + +@pytest.mark.parametrize( + ("payload", "expected_meta"), + [ + ({"results": []}, {"billed_units": {}, "tokens": {}}), + ({"results": [], "usage": {}}, {"billed_units": {}, "tokens": {}}), + ( + {"results": [], "usage": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3, "unknown": 1}}, + {"billed_units": {"total_tokens": 12}, "tokens": {"input_tokens": 7, "output_tokens": 3}}, + ), + ], +) +def test_transform_rerank_response_keeps_only_known_usage_counters( + payload: dict[str, object], expected_meta: dict[str, object] +): + assert _transform(payload).meta == expected_meta + + +def test_transform_rerank_response_generates_an_id_when_the_provider_sends_none(): + response = _transform({"id": None, "results": [{"index": 0, "relevance_score": 0.5}]}) + + assert isinstance(response.id, str) + assert response.id != "" + + +def test_transform_rerank_response_empty_results_list_yields_no_results(): + assert _transform({"id": "rerank-1", "results": []}).results == [] + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + {"results": [], "usage": None}, + {"results": [], "usage": {"total_tokens": 1.5}}, + {"results": [], "usage": {"input_tokens": "many"}}, + {"results": 7}, + {"results": ["not an object"]}, + {"results": [{"index": 0, "relevance_score": 0.5}], "id": 7}, + {"results": [{"index": None, "relevance_score": 0.5}]}, + ], +) +def test_transform_rerank_response_rejects_malformed_payloads(payload: object): + with pytest.raises(ValidationError): + _transform(payload) + + +def test_transform_rerank_response_without_results_raises_value_error_naming_the_body(): + with pytest.raises(ValueError, match="No results found") as exc_info: + _transform({"id": "rerank-1"}) + + assert str(exc_info.value) == "No results found in the response={'id': 'rerank-1'}" + + +def test_transform_rerank_response_result_without_score_raises_key_error(): + with pytest.raises(KeyError, match="relevance_score"): + _transform({"results": [{"index": 0}]}) + + +def test_transform_rerank_response_non_200_raises_with_the_response_text(): + with pytest.raises(Exception, match="quota exceeded") as exc_info: + _transform({"detail": "quota exceeded"}, status_code=429) + + assert type(exc_info.value) is Exception + + +def test_transform_rerank_response_shape_errors_do_not_echo_the_payload(): + with pytest.raises(ValidationError) as exc_info: + _transform({"results": [["leaked document text"]]}) + + assert "leaked document text" not in str(exc_info.value) diff --git a/tests/unit/llms/manus/files/__init__.py b/tests/unit/llms/manus/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/manus/files/test_transformation.py b/tests/unit/llms/manus/files/test_transformation.py new file mode 100644 index 00000000000..a47e1c8cb9b --- /dev/null +++ b/tests/unit/llms/manus/files/test_transformation.py @@ -0,0 +1,128 @@ +import time +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.manus.files.transformation import ManusFilesConfig +from litellm.types.llms.openai import OpenAIFileObject + + +def _list_files(body: object) -> list[OpenAIFileObject]: + return ManusFilesConfig().transform_list_files_response( + raw_response=httpx.Response(200, json=body), logging_obj=Mock(), litellm_params={} + ) + + +def test_delete_file_response_is_read_into_a_file_deleted_object(): + deleted = ManusFilesConfig().transform_delete_file_response( + raw_response=httpx.Response(200, json={"id": "file-1", "deleted": True, "object": "file", "region": "eu"}), + logging_obj=Mock(), + litellm_params={}, + ) + + assert deleted.model_dump() == {"id": "file-1", "deleted": True, "object": "file", "region": "eu"} + + +@pytest.mark.parametrize("body", [b'["secret-file"]', b'"secret-file"', b"7", b"null"]) +def test_delete_file_response_rejects_a_body_that_is_not_an_object_without_echoing_it(body: bytes): + with pytest.raises(ValidationError) as exc_info: + ManusFilesConfig().transform_delete_file_response( + raw_response=httpx.Response(200, content=body), logging_obj=Mock(), litellm_params={} + ) + + assert "secret-file" not in str(exc_info.value) + + +def test_delete_file_response_requires_the_file_deleted_fields(): + with pytest.raises(ValidationError, match="FileDeleted"): + ManusFilesConfig().transform_delete_file_response( + raw_response=httpx.Response(200, json={"id": "file-1"}), logging_obj=Mock(), litellm_params={} + ) + + +def test_list_files_response_maps_every_listed_file(): + files = _list_files( + { + "object": "list", + "data": [ + { + "id": "file-1", + "bytes": 12, + "filename": "a.pdf", + "purpose": "batch", + "status": "processed", + "status_details": "done", + "created_at": "2024-01-02T03:04:05Z", + }, + {"id": "file-2", "created_at": "2024-01-02T03:04:05.123456+00:00"}, + ], + } + ) + + created_at = int(time.mktime(time.strptime("2024-01-02T03:04:05", "%Y-%m-%dT%H:%M:%S"))) + assert [file.model_dump() for file in files] == [ + { + "id": "file-1", + "bytes": 12, + "created_at": created_at, + "filename": "a.pdf", + "object": "file", + "purpose": "batch", + "status": "processed", + "expires_at": None, + "status_details": "done", + }, + { + "id": "file-2", + "bytes": 0, + "created_at": created_at, + "filename": "", + "object": "file", + "purpose": "assistants", + "status": "uploaded", + "expires_at": None, + "status_details": None, + }, + ] + + +@pytest.mark.parametrize("created_at", ["not-a-date", "", None, 0]) +def test_list_files_response_falls_back_to_now_for_an_unusable_created_at(created_at: object): + before = int(time.time()) + + (file,) = _list_files({"data": [{"id": "file-1", "created_at": created_at}]}) + + assert before <= file.created_at <= int(time.time()) + + +@pytest.mark.parametrize("body", [{}, {"data": []}]) +def test_list_files_response_is_empty_without_listed_files(body: dict[str, object]): + assert _list_files(body) == [] + + +@pytest.mark.parametrize( + "body", + [ + ["secret-file"], + "secret-file", + {"data": None}, + {"data": 7}, + {"data": "secret-file"}, + {"data": ["secret-file"]}, + {"data": [{"id": "file-1"}, ["secret-file"]]}, + {"data": [{"id": "file-1", "created_at": ["secret-file"]}]}, + {"data": [{"id": "file-1", "created_at": 1700000000}]}, + ], +) +def test_list_files_response_rejects_malformed_listings_without_echoing_them(body: object): + with pytest.raises(ValidationError) as exc_info: + _list_files(body) + + assert "secret-file" not in str(exc_info.value) + + +def test_list_files_response_rejects_a_file_the_file_object_cannot_hold(): + with pytest.raises(ValidationError, match="OpenAIFileObject"): + _list_files({"data": [{"id": "file-1", "purpose": "not-a-purpose"}]}) diff --git a/tests/unit/llms/minimax/text_to_speech/__init__.py b/tests/unit/llms/minimax/text_to_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/minimax/text_to_speech/test_transformation.py b/tests/unit/llms/minimax/text_to_speech/test_transformation.py new file mode 100644 index 00000000000..a183c314a61 --- /dev/null +++ b/tests/unit/llms/minimax/text_to_speech/test_transformation.py @@ -0,0 +1,122 @@ +import base64 +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest + +from litellm.llms.minimax.text_to_speech.transformation import MinimaxException, MinimaxTextToSpeechConfig +from litellm.types.llms.openai import HttpxBinaryResponseContent + +_AUDIO: Final = b"ID3\x04minimax-audio" +_REQUEST: Final = httpx.Request("POST", "https://api.minimax.io/v1/t2a_v2") + + +def _transform(payload: object, status_code: int = 200) -> HttpxBinaryResponseContent: + return MinimaxTextToSpeechConfig().transform_text_to_speech_response( + model="speech-02-hd", + raw_response=httpx.Response( + status_code, + json=payload, + request=_REQUEST, + headers={"content-encoding": "identity", "x-trace": "abc"}, + ), + logging_obj=Mock(), + ) + + +@pytest.mark.parametrize( + "payload", + [ + {"data": {"audio": _AUDIO.hex()}, "status": 0, "extra_info": {"audio_length": 5}}, + {"data": {"audio": _AUDIO.hex(), "audio_url": ""}}, + {"data": {"audio": ""}, "audio_file": _AUDIO.hex()}, + {"data": {"audio": None}, "audio_file": base64.b64encode(_AUDIO).decode()}, + {"base_resp": {"status_code": 0}, "audio_file": base64.b64encode(_AUDIO).decode()}, + ], +) +def test_transform_text_to_speech_response_decodes_hex_or_base64_audio(payload: dict[str, object]): + response = _transform(payload).response + + assert response.status_code == 200 + assert response.content == _AUDIO + assert response.headers["content-length"] == str(len(_AUDIO)) + assert response.headers["x-trace"] == "abc" + assert "content-encoding" not in response.headers + + +@pytest.mark.parametrize( + ("payload", "detail"), + [ + ({"status": 2, "ced": "invalid api key"}, "invalid api key"), + ({"status": 2}, "Unknown error"), + ({"status": 1004, "ced": ""}, "API returned status 1004"), + ({"status": "failed", "ced": None, "data": {"audio": _AUDIO.hex()}}, "API returned status failed"), + ], +) +def test_transform_text_to_speech_response_reports_api_status_errors(payload: dict[str, object], detail: str): + with pytest.raises(MinimaxException) as exc_info: + _transform(payload, status_code=401) + + assert exc_info.value.message == f"MiniMax TTS error: {detail}" + assert exc_info.value.status_code == 401 + + +def test_transform_text_to_speech_response_refuses_url_output(): + with pytest.raises(MinimaxException) as exc_info: + _transform({"data": {"audio_url": "https://cdn.example/a.mp3", "audio": _AUDIO.hex()}}) + + assert exc_info.value.message == ( + "URL output format is not yet supported. Use 'hex' format or fetch from URL: https://cdn.example/a.mp3" + ) + assert exc_info.value.status_code == 500 + + +@pytest.mark.parametrize( + ("payload", "keys"), + [ + ({}, []), + ({"data": {}, "status": 0}, ["data", "status"]), + ({"data": {"audio": ""}, "audio_file": ""}, ["data", "audio_file"]), + ({"data": {"audio": None}, "audio_file": None}, ["data", "audio_file"]), + ({"audio_file": 0, "data": {"audio": []}}, ["audio_file", "data"]), + ], +) +def test_transform_text_to_speech_response_without_audio_lists_the_response_keys( + payload: dict[str, object], keys: list[str] +): + with pytest.raises(MinimaxException) as exc_info: + _transform(payload) + + assert exc_info.value.message == f"No audio data in MiniMax response. Response keys: {keys}" + assert exc_info.value.status_code == 500 + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + "not an object", + {"data": None}, + {"data": "not an object"}, + {"data": ["not", "an", "object"]}, + {"data": {"audio": 7}}, + {"data": {"audio": ["49", "44"]}}, + {"audio_file": {"hex": "4944"}}, + ], +) +def test_transform_text_to_speech_response_wraps_malformed_payloads_without_echoing_them(payload: object): + with pytest.raises(MinimaxException) as exc_info: + _transform(payload) + + assert exc_info.value.message.startswith("Error processing MiniMax response: ") + assert "input_value" not in exc_info.value.message + assert exc_info.value.status_code == 500 + + +def test_transform_text_to_speech_response_reports_undecodable_audio(): + with pytest.raises(MinimaxException) as exc_info: + _transform({"data": {"audio": "zzz"}}) + + assert exc_info.value.message.startswith("Failed to decode audio data: ") + assert exc_info.value.status_code == 500 diff --git a/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py b/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py index 2ffe7c3e686..5b0ca09ed13 100644 --- a/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py +++ b/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py @@ -7,9 +7,12 @@ transformation between OpenAI-compatible format and ModelScope API format. from unittest.mock import MagicMock, patch +import httpx import pytest +from pydantic import ValidationError +import litellm from litellm.llms.modelscope.image_generation.transformation import ( ModelScopeImageGenerationConfig, ) @@ -449,3 +452,83 @@ class TestModelScopeImageGenerationTransformation: ) assert isinstance(error, BadRequestError) + + +def _transform_generation_response(payload: object, status_code: int = 200) -> ImageResponse: + return ModelScopeImageGenerationConfig().transform_image_generation_response( + model="Qwen/Qwen-Image", + raw_response=httpx.Response(status_code, json=payload), + model_response=ImageResponse(data=[]), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + ("item", "expected"), + [ + ({"url": "https://a.example/i.png"}, ("https://a.example/i.png", None, None)), + ({"b64_json": "aGVsbG8="}, (None, "aGVsbG8=", None)), + ( + {"url": "https://a.example/i.png", "revised_prompt": "a calmer cat", "seed": 7}, + ("https://a.example/i.png", None, "a calmer cat"), + ), + ({}, (None, None, None)), + ], +) +def test_transform_image_generation_response_maps_one_image( + item: dict[str, object], expected: tuple[str | None, str | None, str | None] +) -> None: + response = _transform_generation_response({"created": 1, "data": [item]}) + + assert [(image.url, image.b64_json, image.revised_prompt) for image in response.data] == [expected] + + +@pytest.mark.parametrize("payload", [{}, {"data": []}, {"data": ""}, {"data": {}}, {"created": 1}]) +def test_transform_image_generation_response_without_images_is_empty(payload: dict[str, object]) -> None: + assert _transform_generation_response(payload).data == [] + + +@pytest.mark.parametrize( + ("error", "status_code", "expected_class", "expected_message"), + [ + ({"message": "Invalid prompt"}, 400, litellm.BadRequestError, "ModelScope error: Invalid prompt"), + ({"message": "Bad key"}, 401, litellm.AuthenticationError, "ModelScope error: Bad key"), + ({"code": "overloaded"}, 503, litellm.InternalServerError, "ModelScope error: {'code': 'overloaded'}"), + ({}, 200, litellm.BadRequestError, "ModelScope error: {}"), + ], +) +def test_transform_image_generation_response_reports_api_error_bodies( + error: dict[str, object], status_code: int, expected_class: type[Exception], expected_message: str +) -> None: + with pytest.raises(expected_class) as exc_info: + _transform_generation_response({"error": error, "data": [{"url": "ignored"}]}, status_code) + + assert str(exc_info.value).endswith(expected_message) + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + "error", + 7, + {"error": "plain text error"}, + {"error": None}, + {"error": ["not", "an", "object"]}, + {"data": 7}, + {"data": None}, + {"data": ["not an object"]}, + {"data": [{"url": "https://a.example/i.png"}, None]}, + ], +) +def test_transform_image_generation_response_rejects_malformed_payloads_without_echoing_them( + payload: object, +) -> None: + with pytest.raises(ValidationError) as exc_info: + _transform_generation_response(payload) + + assert "input_value" not in str(exc_info.value) diff --git a/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py b/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py index 2b03b2d807b..39ba396d8d2 100644 --- a/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py +++ b/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py @@ -10,7 +10,9 @@ truncate. Two defects are covered here: import json from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +from pydantic import ValidationError import litellm from litellm.llms.nvidia_nim.rerank.ranking_transformation import ( @@ -239,3 +241,87 @@ class TestNvidiaNimRetrievalRerankRequestTransform: doc = {"title": "no supported fields here"} request_data = self._build_request([doc]) assert request_data["passages"] == [{"text": json.dumps(doc)}] + + +def _transform_retrieval_response(payload: object) -> RerankResponse: + return NvidiaNimRerankConfig().transform_rerank_response( + model="nvidia/llama-3_2-nv-rerankqa-1b-v2", + raw_response=httpx.Response(200, json=payload), + model_response=RerankResponse(), + logging_obj=MagicMock(), + request_data={"passages": [{"text": "first"}, {"text": "second"}]}, + ) + + +@pytest.mark.parametrize( + ("usage", "expected_total_tokens"), + [ + ({"total_tokens": 42}, 42), + ({"total_tokens": 0}, 2), + ({}, 2), + ], +) +def test_transform_rerank_response_bills_reported_tokens_or_falls_back_to_result_count( + usage: dict[str, object], expected_total_tokens: int +): + response = _transform_retrieval_response( + {"rankings": [{"index": 1, "logit": 2.5}, {"index": 0, "logit": -1}], "usage": usage} + ) + + assert response.meta == {"billed_units": {"total_tokens": expected_total_tokens}} + assert response.results == [ + {"index": 1, "relevance_score": 2.5, "document": {"text": "second"}}, + {"index": 0, "relevance_score": -1.0, "document": {"text": "first"}}, + ] + + +def test_transform_rerank_response_keeps_provider_id_and_bills_one_token_per_result_without_usage(): + response = _transform_retrieval_response({"id": "rank-1", "rankings": [{"index": 0, "logit": 0.5}]}) + + assert response.id == "rank-1" + assert response.meta == {"billed_units": {"total_tokens": 1}} + + +@pytest.mark.parametrize( + "payload", + [ + {"rankings": [], "usage": None}, + {"rankings": [], "usage": {"total_tokens": None}}, + {"rankings": [], "usage": {"total_tokens": "42"}}, + {"rankings": [], "usage": {"total_tokens": 1.5}}, + {"rankings": [], "id": 7}, + ], +) +def test_transform_rerank_response_rejects_malformed_usage_and_id(payload: dict[str, object]): + with pytest.raises(ValidationError): + _transform_retrieval_response(payload) + + +def test_transform_rerank_response_non_object_body_raises_attribute_error(): + with pytest.raises(AttributeError): + _transform_retrieval_response(["not", "an", "object"]) + + +def test_transform_rerank_response_usage_shape_errors_do_not_echo_the_payload(): + with pytest.raises(ValidationError) as exc_info: + _transform_retrieval_response({"rankings": [], "usage": {"total_tokens": "leaked usage text"}}) + + assert "leaked usage text" not in str(exc_info.value) + + +def test_map_cohere_rerank_params_passes_provider_params_through_and_maps_top_n(): + params = NvidiaNimRerankConfig().map_cohere_rerank_params( + non_default_params={"truncate": "END"}, + model="nvidia/llama-3_2-nv-rerankqa-1b-v2", + drop_params=False, + query="which passage shows a cat?", + documents=["a", {"text": "b"}], + top_n=1, + ) + + assert params == { + "query": "which passage shows a cat?", + "documents": ["a", {"text": "b"}], + "top_k": 1, + "truncate": "END", + } diff --git a/tests/unit/llms/oci/embed/test_oci_embed_transformation.py b/tests/unit/llms/oci/embed/test_oci_embed_transformation.py index 4ffd79ff147..01be7c7b904 100644 --- a/tests/unit/llms/oci/embed/test_oci_embed_transformation.py +++ b/tests/unit/llms/oci/embed/test_oci_embed_transformation.py @@ -379,3 +379,78 @@ class TestOCIEmbedConfig: litellm_params={}, ) assert "eu-frankfurt-1" in url + + +def _transform(raw_response: httpx.Response) -> EmbeddingResponse: + return OCIEmbedConfig().transform_embedding_response( + model="cohere.embed-v3.0", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key=None, + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +@pytest.mark.parametrize( + ("token_fields", "expected_prompt_tokens", "expected_total_tokens"), + [ + ({"inputTextTokenCounts": [5, 6]}, 11, 11), + ({"usage": {"promptTokens": 3, "totalTokens": 4}}, 3, 4), + ({"inputTextTokenCounts": [5, 6], "usage": {"promptTokens": 3, "totalTokens": 4}}, 11, 11), + ({}, 0, 0), + ], +) +def test_transform_embedding_response_reads_usage_from_whichever_token_field_is_present( + token_fields: dict[str, object], expected_prompt_tokens: int, expected_total_tokens: int +): + response = _transform( + httpx.Response( + 200, + json={ + "embeddings": [[0.1, 1], [0.5]], + "modelId": "cohere.embed-v3.0", + "modelVersion": "3.0", + "unknown": "ignored", + **token_fields, + }, + ) + ) + + assert response.model_dump() == { + "model": "cohere.embed-v3.0", + "data": [ + {"object": "embedding", "index": 0, "embedding": [0.1, 1.0]}, + {"object": "embedding", "index": 1, "embedding": [0.5]}, + ], + "object": "list", + "usage": { + "completion_tokens": 0, + "prompt_tokens": expected_prompt_tokens, + "total_tokens": expected_total_tokens, + "completion_tokens_details": None, + "prompt_tokens_details": None, + }, + } + + +@pytest.mark.parametrize("body", [b"null", b"7", b'["leaked payload text"]', b'"leaked payload text"']) +def test_transform_embedding_response_body_that_is_not_an_object_is_a_schema_error(body: bytes): + with pytest.raises(OCIError) as exc_info: + _transform(httpx.Response(200, content=body)) + + assert exc_info.value.status_code == 500 + assert exc_info.value.message.startswith("OCI embed response does not match expected schema: ") + assert "leaked payload text" not in exc_info.value.message + + +def test_transform_embedding_response_object_missing_required_fields_names_them(): + with pytest.raises(OCIError) as exc_info: + _transform(httpx.Response(200, json={"embeddings": [[0.1]]})) + + assert exc_info.value.status_code == 500 + assert exc_info.value.message.startswith("OCI embed response does not match expected schema: ") + assert "modelId" in exc_info.value.message + assert "modelVersion" in exc_info.value.message diff --git a/tests/unit/llms/openai/responses/count_tokens/__init__.py b/tests/unit/llms/openai/responses/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/responses/count_tokens/test_handler.py b/tests/unit/llms/openai/responses/count_tokens/test_handler.py new file mode 100644 index 00000000000..2427b8d272f --- /dev/null +++ b/tests/unit/llms/openai/responses/count_tokens/test_handler.py @@ -0,0 +1,26 @@ +import pytest + +from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.responses.count_tokens.handler import OpenAICountTokensHandler + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "input_value", "expected_message"), + [ + ("", "hello", "CountTokens processing error: model parameter is required"), + ("gpt-4o", "", "CountTokens processing error: input parameter is required"), + ("gpt-4o", [], "CountTokens processing error: input parameter is required"), + ], +) +async def test_request_without_model_or_input_is_rejected_before_calling_openai( + model: str, input_value: str | list[object], expected_message: str +) -> None: + with pytest.raises(OpenAIError) as rejected: + await OpenAICountTokensHandler().handle_count_tokens_request( + model=model, + input=input_value, + api_key="sk-test", + ) + + assert (rejected.value.status_code, rejected.value.message) == (500, expected_message) diff --git a/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py b/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py index e45270fb5e3..edd6168c6c3 100644 --- a/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py +++ b/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py @@ -580,3 +580,86 @@ class TestOpenRouterImageGenerationTransformation: assert isinstance(error, OpenRouterException) assert "Test error" in str(error) assert error.status_code == 400 + + +def _transform_generation_response(payload: object) -> ImageResponse: + return OpenRouterImageGenerationConfig().transform_image_generation_response( + model="google/gemini-2.5-flash-image", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(data=[]), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"choices": []}, + {"choices": ""}, + {"choices": {}}, + {"choices": [{}]}, + {"choices": [{"message": {}}]}, + {"choices": [{"message": {"images": ""}}]}, + {"choices": [{"message": {"images": [{}]}}]}, + {"choices": [{"message": {"images": [{"image_url": {}}]}}]}, + {"choices": [{"message": {"images": [{"image_url": {"url": None}}]}}]}, + {"choices": [{"message": {"images": [{"image_url": {"url": ""}}]}}]}, + {"choices": [{"message": {"images": [{"image_url": {"url": 0}}]}}]}, + ], +) +def test_transform_image_generation_response_without_usable_image_url_has_no_images( + payload: dict[str, object], +): + assert _transform_generation_response(payload).data == [] + + +@pytest.mark.parametrize( + ("url", "expected"), + [ + ("data:image/png;base64,aGVsbG8=", ("aGVsbG8=", None)), + ("data:image/png;base64,a,b", ("a,b", None)), + ("data:no-comma", (None, None)), + ("https://example.com/a.png", (None, "https://example.com/a.png")), + ], +) +def test_transform_image_generation_response_maps_one_image_url( + url: str, expected: tuple[str | None, str | None] +): + response = _transform_generation_response( + {"choices": [{"message": {"images": [{"image_url": {"url": url}}]}}]} + ) + + assert [(image.b64_json, image.url) for image in response.data] == [expected] + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + {"choices": 7}, + {"choices": ["not an object"]}, + {"choices": [{"message": "not an object"}]}, + {"choices": [{"message": {"images": 7}}]}, + {"choices": [{"message": {"images": ["not an object"]}}]}, + {"choices": [{"message": {"images": [{"image_url": "not an object"}]}}]}, + {"choices": [{"message": {"images": [{"image_url": {"url": ["a"]}}]}}]}, + {"choices": [{"message": {"images": [{"image_url": {"url": 7}}]}}]}, + ], +) +def test_transform_image_generation_response_wraps_malformed_payloads_without_echoing_them( + payload: object, +): + with pytest.raises(OpenRouterException) as exc_info: + _transform_generation_response(payload) + + message = str(exc_info.value) + assert message.startswith( + "Error transforming OpenRouter image generation response: " + ) + assert "input_value" not in message + assert exc_info.value.status_code == 500 diff --git a/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py index 87e54dfba9b..422658932e6 100644 --- a/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py +++ b/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py @@ -1,4 +1,9 @@ +import httpx +import pytest +from pydantic import ValidationError +from litellm.llms.ovhcloud.audio_transcription.transformation import OVHCloudAudioTranscriptionConfig +from litellm.types.utils import TranscriptionResponse @@ -56,3 +61,40 @@ class TestOVHCloudDurationFieldMigration: mock_response.json.return_value = {"text": "silence", "seconds": 0.0} result = config.transform_audio_transcription_response(mock_response) assert result._hidden_params["duration"] == 0.0 + + +def _transform(payload: object) -> TranscriptionResponse: + return OVHCloudAudioTranscriptionConfig().transform_audio_transcription_response(httpx.Response(200, json=payload)) + + +@pytest.mark.parametrize( + ("payload", "expected_text", "expected_hidden_params"), + [ + ( + {"text": "hello", "seconds": 3.5, "duration": 9}, + "hello", + {"text": "hello", "seconds": 3.5, "duration": 3.5}, + ), + ( + {"transcript": "from transcript", "seconds": None, "duration": 4}, + "from transcript", + {"transcript": "from transcript", "seconds": None, "duration": 4}, + ), + ({"language": "en"}, "", {"language": "en"}), + ], +) +def test_transform_audio_transcription_response_normalizes_text_and_duration( + payload: dict[str, object], expected_text: str, expected_hidden_params: dict[str, object] +): + response = _transform(payload) + + assert response.text == expected_text + assert response._hidden_params == expected_hidden_params + + +@pytest.mark.parametrize("payload", [7, "spoken secret", [{"text": "spoken secret"}]]) +def test_transform_audio_transcription_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "spoken secret" not in str(exc_info.value) diff --git a/tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py b/tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py index 2dfe33b828c..4b677675694 100644 --- a/tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py +++ b/tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest +from pydantic import ValidationError from litellm.llms.recraft.image_generation.transformation import ( @@ -256,3 +257,56 @@ class TestRecraftImageGenerationTransformation: ) assert "Error transforming image generation response" in str(exc_info.value) + + +def _transform(payload: object) -> ImageResponse: + return RecraftImageGenerationConfig().transform_image_generation_response( + model="recraftv3", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_each_image_object(): + response = _transform({"data": [{"url": "https://img.recraft.ai/a.png"}, {"b64_json": "QUJD"}, {}]}) + + assert [(image.url, image.b64_json) for image in response.data] == [ + ("https://img.recraft.ai/a.png", None), + (None, "QUJD"), + (None, None), + ] + + +@pytest.mark.parametrize("data", [[], "", {}]) +def test_transform_image_generation_response_with_empty_data_has_no_images(data: object): + assert _transform({"data": data}).data == [] + + +@pytest.mark.parametrize( + "payload", + [ + 7, + "https://img.recraft.ai/a.png", + [{"url": "https://img.recraft.ai/a.png"}], + {"data": None}, + {"data": "https://img.recraft.ai/a.png"}, + {"data": {"url": "https://img.recraft.ai/a.png"}}, + {"data": ["https://img.recraft.ai/a.png"]}, + {"data": [{"url": "https://img.recraft.ai/a.png"}, 7]}, + ], +) +def test_transform_image_generation_response_rejects_malformed_payloads(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "img.recraft.ai" not in str(exc_info.value) + + +def test_transform_image_generation_response_requires_data(): + with pytest.raises(KeyError): + _transform({"created": 1}) diff --git a/tests/unit/llms/replicate/__init__.py b/tests/unit/llms/replicate/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/replicate/chat/__init__.py b/tests/unit/llms/replicate/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/replicate/chat/test_transformation.py b/tests/unit/llms/replicate/chat/test_transformation.py new file mode 100644 index 00000000000..4c2b1840664 --- /dev/null +++ b/tests/unit/llms/replicate/chat/test_transformation.py @@ -0,0 +1,54 @@ +from unittest.mock import Mock + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.replicate.chat.transformation import ReplicateConfig +from litellm.types.utils import ModelResponse + + +def _transform(raw_response: httpx.Response) -> ModelResponse: + return ReplicateConfig().transform_response( + model="acme/echo-model", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={"input": {"prompt": "Hello"}}, + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + ("output", "content"), + [ + (["Hello", ", ", "world"], "Hello, world"), + ("Hello", "Hello"), + ([], " "), + ("", " "), + ([""], " "), + ], +) +def test_transform_response_joins_the_prediction_output_into_the_message_content(output: object, content: str): + response = _transform(httpx.Response(200, json={"status": "succeeded", "output": output})) + + assert response.choices[0].message.content == content + assert response.model == "replicate/acme/echo-model" + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +def test_transform_response_uses_a_blank_message_when_the_prediction_has_no_output(): + response = _transform(httpx.Response(200, json={"status": "succeeded"})) + + assert response.choices[0].message.content == " " + + +@pytest.mark.parametrize("output", [None, 7, True, ["secret-output", 7], ["secret-output", None], [["secret-output"]]]) +def test_transform_response_rejects_an_output_that_is_not_made_of_strings_without_echoing_it(output: object): + with pytest.raises(ValidationError) as exc_info: + _transform(httpx.Response(200, json={"status": "succeeded", "output": output})) + + assert "secret-output" not in str(exc_info.value) diff --git a/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py index e2dd3bca74f..51a2ccb5640 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py @@ -70,6 +70,19 @@ def test_sagemaker_response_stream_shape_load_failure_returns_none(): assert shape is None +@pytest.mark.parametrize("service_model", [["shapes"], None]) +def test_sagemaker_response_stream_shape_is_none_for_a_service_model_that_is_not_a_mapping( + service_model: object, +): + pytest.importorskip("botocore") + from unittest.mock import patch + + import litellm.llms.sagemaker.common_utils as mod + + with patch("botocore.loaders.Loader.load_service_model", return_value=service_model): + assert mod._load_sagemaker_response_stream_shape() is None + + def test_sagemaker_response_stream_shape_is_structure_shape(): """ The loaded shape should be the botocore StructureShape for diff --git a/tests/unit/llms/snowflake/embedding/test_snowflake_embedding.py b/tests/unit/llms/snowflake/embedding/test_snowflake_embedding.py index 8f4551bc79b..ff2af8a8d7b 100644 --- a/tests/unit/llms/snowflake/embedding/test_snowflake_embedding.py +++ b/tests/unit/llms/snowflake/embedding/test_snowflake_embedding.py @@ -2,9 +2,15 @@ import os import json import copy -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from pydantic import ValidationError import litellm +from litellm.llms.snowflake.embedding.transformation import SnowflakeEmbeddingConfig +from litellm.types.utils import EmbeddingResponse model_name = "snowflake-arctic-embed" @@ -94,3 +100,59 @@ def test_snowflake_env(mock_post): os.environ.pop("SNOWFLAKE_ACCOUNT_ID", None) os.environ.pop("SNOWFLAKE_JWT", None) + + +def _transform(raw_response: httpx.Response) -> EmbeddingResponse: + return SnowflakeEmbeddingConfig().transform_embedding_response( + model="snowflake-arctic-embed-m", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="test-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def test_transform_embedding_response_flattens_each_vector_and_prefixes_the_model(): + response = _transform( + httpx.Response( + 200, + json={ + "object": "list", + "model": "snowflake-arctic-embed-m", + "data": [{"object": "embedding", "index": 0, "embedding": [[0.1, 2]]}], + "usage": {"total_tokens": 5}, + "unknown": "ignored", + }, + ) + ) + + assert response.model_dump() == { + "model": "snowflake/snowflake-arctic-embed-m", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 2]}], + "object": "list", + "usage": { + "completion_tokens": 0, + "prompt_tokens": 0, + "total_tokens": 5, + "completion_tokens_details": None, + "prompt_tokens_details": None, + }, + } + assert response._hidden_params["model"] == "snowflake-arctic-embed-m" + + +@pytest.mark.parametrize("body", [b"null", b"[]", b"7"]) +def test_transform_embedding_response_body_that_is_not_an_object_raises_type_error(body: bytes): + with pytest.raises(TypeError): + _transform(httpx.Response(200, content=body)) + + +def test_transform_embedding_response_invalid_envelope_field_is_reported_by_name(): + with pytest.raises(ValidationError) as exc_info: + _transform(httpx.Response(200, json={"data": [], "model": 5})) + + assert exc_info.value.title == "EmbeddingResponse" + assert [error["loc"] for error in exc_info.value.errors()] == [("model",)] diff --git a/tests/unit/llms/stability/image_generation/test_stability_image_generation.py b/tests/unit/llms/stability/image_generation/test_stability_image_generation.py index c5a78603f9c..42eb1a2365b 100644 --- a/tests/unit/llms/stability/image_generation/test_stability_image_generation.py +++ b/tests/unit/llms/stability/image_generation/test_stability_image_generation.py @@ -9,6 +9,7 @@ from unittest.mock import MagicMock import httpx import pytest +from pydantic import ValidationError from litellm.llms.stability.image_generation import StabilityImageGenerationConfig from litellm.types.llms.stability import ( @@ -304,3 +305,35 @@ class TestStabilityGenerationModels: STABILITY_GENERATION_MODELS["stable-image-core"] == "/v2beta/stable-image/generate/core" ) + + +def _transform(payload: object) -> ImageResponse: + return StabilityImageGenerationConfig().transform_image_generation_response( + model="sd3", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_returns_the_base64_image(): + response = _transform({"image": "QUJD", "finish_reason": "SUCCESS", "seed": 7}) + + assert [(image.b64_json, image.url) for image in response.data] == [("QUJD", None)] + + +@pytest.mark.parametrize("payload", [{}, {"image": ""}, {"image": None, "finish_reason": None}]) +def test_transform_image_generation_response_without_an_image_has_no_data(payload: dict[str, object]): + assert _transform(payload).data == [] + + +@pytest.mark.parametrize("payload", ["base64encodedimage==", ["base64encodedimage=="]]) +def test_transform_image_generation_response_rejects_non_object_bodies(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "base64encodedimage" not in str(exc_info.value) diff --git a/tests/unit/llms/tinyfish/test_tinyfish_search.py b/tests/unit/llms/tinyfish/test_tinyfish_search.py index 69afbb416aa..e7020238ba6 100644 --- a/tests/unit/llms/tinyfish/test_tinyfish_search.py +++ b/tests/unit/llms/tinyfish/test_tinyfish_search.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.tinyfish.search.transformation import ( TinyfishSearchConfig, _append_domain_filters, @@ -821,3 +822,56 @@ class TestDefaultMissingResultFields: assert raw_json["results"][0] == "string item" assert raw_json["results"][1] == 42 assert raw_json["results"][2] == {"title": "ok", "url": "", "snippet": ""} + + +def test_transform_search_response_defaults_missing_fields_of_a_decoded_http_body(): + payload = { + "query": "q", + "results": [ + {"title": "T", "url": "https://a.example", "snippet": "S", "position": 1}, + {"title": None, "position": 2}, + ], + } + + response = TinyfishSearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload, headers={"x-request-id": "req-1"}), + logging_obj=None, + ) + + assert [result.model_dump(exclude_none=True) for result in response.results] == [ + {"title": "T", "url": "https://a.example", "snippet": "S", "position": 1}, + {"title": "", "url": "", "snippet": "", "position": 2}, + ] + assert response.model_dump()["query"] == "q" + assert response._hidden_params["headers"]["x-request-id"] == "req-1" + + +@pytest.mark.parametrize("payload", [[], "text", 5, {}, {"results": "text"}, {"results": [5]}]) +def test_transform_search_response_wraps_a_decoded_body_of_the_wrong_shape(payload: object): + with pytest.raises(BaseLLMException, match="TinyFish Search: Response shape does not match") as exc_info: + TinyfishSearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=None + ) + + assert exc_info.value.status_code == 200 + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + ('{"error": {"code": "INVALID_INPUT", "message": "query is required"}}', "query is required"), + ('{"error": {"message": ""}}', '{"error": {"message": ""}}'), + ('{"error": {"message": 5}}', '{"error": {"message": 5}}'), + ('{"error": "flat"}', '{"error": "flat"}'), + ('["not", "an", "envelope"]', '["not", "an", "envelope"]'), + ("Bad Gateway", "Bad Gateway"), + ], +) +def test_transform_search_response_unwraps_only_the_tinyfish_error_envelope(body: str, expected: str): + with pytest.raises(BaseLLMException) as exc_info: + TinyfishSearchConfig().transform_search_response(raw_response=httpx.Response(400, text=body), logging_obj=None) + + assert ( + exc_info.value.message == f"TinyFish Search: {expected}. See https://docs.tinyfish.ai/search-api for details." + ) + assert exc_info.value.status_code == 400 diff --git a/tests/unit/llms/together_ai/rerank/__init__.py b/tests/unit/llms/together_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/together_ai/rerank/test_handler.py b/tests/unit/llms/together_ai/rerank/test_handler.py new file mode 100644 index 00000000000..bc2f30364d2 --- /dev/null +++ b/tests/unit/llms/together_ai/rerank/test_handler.py @@ -0,0 +1,96 @@ +import json +from collections.abc import Iterator + +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.llms.together_ai.rerank.handler import TogetherAIRerank + +API_BASE = "https://api.together.example/v1" +RERANKED = { + "id": "rr-1", + "results": [ + {"index": 1, "relevance_score": 0.9, "document": {"text": "Paris"}}, + {"index": 0, "relevance_score": 0.1}, + ], + "usage": {"total_tokens": 7}, +} +NON_OBJECT_BODIES = [7, "sensitive-document", [{"id": "sensitive-document"}]] + + +@pytest.fixture +def rerank_route(monkeypatch: pytest.MonkeyPatch) -> Iterator[respx.Route]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock as router: + yield router.post(f"{API_BASE}/rerank") + + +def _rerank(*, is_async: bool): + return TogetherAIRerank().rerank( + model="Salesforce/Llama-Rank-V1", + api_key="sk-test", + api_base=API_BASE, + query="capital of france", + documents=["Berlin", {"text": "Paris"}], + top_n=1, + _is_async=is_async, + ) + + +def _assert_is_the_reranking(response: litellm.RerankResponse, route: respx.Route) -> None: + assert response.id == "rr-1" + assert response.results == [ + {"index": 1, "relevance_score": 0.9, "document": {"text": "Paris"}}, + {"index": 0, "relevance_score": 0.1}, + ] + assert response.meta == {"billed_units": {"total_tokens": 7}, "tokens": {}} + assert route.calls.last.request.headers["authorization"] == "Bearer sk-test" + assert json.loads(route.calls.last.request.content) == { + "model": "Salesforce/Llama-Rank-V1", + "query": "capital of france", + "top_n": 1, + "documents": ["Berlin", {"text": "Paris"}], + "return_documents": True, + } + + +def test_rerank_returns_the_upstream_ranking(rerank_route: respx.Route): + rerank_route.mock(return_value=httpx.Response(200, json=RERANKED)) + + _assert_is_the_reranking(_rerank(is_async=False), rerank_route) + + +async def test_async_rerank_returns_the_upstream_ranking(rerank_route: respx.Route): + rerank_route.mock(return_value=httpx.Response(200, json=RERANKED)) + + _assert_is_the_reranking(await _rerank(is_async=True), rerank_route) + + +@pytest.mark.parametrize("payload", NON_OBJECT_BODIES) +def test_rerank_rejects_non_object_bodies(rerank_route: respx.Route, payload: object): + rerank_route.mock(return_value=httpx.Response(200, json=payload)) + + with pytest.raises(ValidationError) as exc_info: + _rerank(is_async=False) + + assert "sensitive-document" not in str(exc_info.value) + + +@pytest.mark.parametrize("payload", NON_OBJECT_BODIES) +async def test_async_rerank_rejects_non_object_bodies(rerank_route: respx.Route, payload: object): + rerank_route.mock(return_value=httpx.Response(200, json=payload)) + + with pytest.raises(ValidationError) as exc_info: + await _rerank(is_async=True) + + assert "sensitive-document" not in str(exc_info.value) + + +def test_rerank_without_results_is_a_value_error(rerank_route: respx.Route): + rerank_route.mock(return_value=httpx.Response(200, json={"id": "rr-1"})) + + with pytest.raises(ValueError, match="No results found"): + _rerank(is_async=False) diff --git a/tests/unit/llms/vertex_ai/image_generation/test_image_generation_handler.py b/tests/unit/llms/vertex_ai/image_generation/test_image_generation_handler.py new file mode 100644 index 00000000000..ac5e14c0388 --- /dev/null +++ b/tests/unit/llms/vertex_ai/image_generation/test_image_generation_handler.py @@ -0,0 +1,43 @@ +import pytest +from pydantic import ValidationError + +from litellm.llms.vertex_ai.image_generation.image_generation_handler import VertexImageGeneration +from litellm.types.utils import ImageResponse + + +def _process(json_response: dict[str, object]) -> ImageResponse: + return VertexImageGeneration().process_image_generation_response( + json_response, ImageResponse(), "imagegeneration@006" + ) + + +@pytest.mark.parametrize( + ("predictions", "expected"), + [ + ([{"bytesBase64Encoded": "QUJD", "mimeType": "image/png"}, {"bytesBase64Encoded": "REVG"}], ["QUJD", "REVG"]), + ([{"bytesBase64Encoded": None}], [None]), + ([], []), + ], +) +def test_process_image_generation_response_maps_each_prediction( + predictions: list[dict[str, object]], expected: list[str | None] +): + response = _process({"predictions": predictions, "deployedModelId": "1"}) + + assert [image.b64_json for image in response.data] == expected + + +@pytest.mark.parametrize( + "predictions", + [None, 7, ["QUJD"], [{"bytesBase64Encoded": "QUJD"}, None], [{"bytesBase64Encoded": 7}]], +) +def test_process_image_generation_response_rejects_malformed_predictions(predictions: object): + with pytest.raises(ValidationError) as exc_info: + _process({"predictions": predictions}) + + assert "QUJD" not in str(exc_info.value) + + +def test_process_image_generation_response_requires_the_encoded_bytes_key(): + with pytest.raises(KeyError): + _process({"predictions": [{"mimeType": "image/png"}]}) diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index a72a570c2a2..0461c739395 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -1,6 +1,8 @@ from unittest.mock import MagicMock, patch import httpx +import pytest +from pydantic import ValidationError from litellm.llms.vertex_ai.image_generation import ( @@ -12,6 +14,7 @@ from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import ( VertexAIImagenImageGenerationConfig, ) +from litellm.types.utils import ImageResponse class TestVertexAIGeminiImageGenerationConfig: @@ -635,3 +638,64 @@ class TestVertexAIImageGenerationIntegration: assert "us-central1" in url assert "imagegeneration@006" in url assert "predict" in url + + +def _transform_gemini_response(payload: object) -> ImageResponse: + return VertexAIGeminiImageGenerationConfig().transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_gemini_image_generation_response_maps_usage_by_modality(): + response = _transform_gemini_response( + { + "candidates": [{"content": {"parts": [{"inlineData": {"data": "aGVsbG8="}}]}}], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 7, + "totalTokenCount": 12, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 3}, + {"modality": "IMAGE", "tokenCount": 2}, + ], + }, + } + ) + + assert [image.b64_json for image in response.data] == ["aGVsbG8="] + assert response.usage.input_tokens == 5 + assert response.usage.output_tokens == 7 + assert response.usage.total_tokens == 12 + assert response.usage.input_tokens_details.text_tokens == 3 + assert response.usage.input_tokens_details.image_tokens == 2 + + +@pytest.mark.parametrize("usage_metadata", [None, {}, [], "", 0]) +def test_gemini_image_generation_response_with_falsy_usage_metadata_keeps_zeroed_usage(usage_metadata: object): + response = _transform_gemini_response( + { + "candidates": [{"content": {"parts": []}, "groundingMetadata": {"webSearchQueries": ["a"]}}], + "usageMetadata": usage_metadata, + } + ) + + assert response.data == [] + assert (response.usage.input_tokens, response.usage.output_tokens, response.usage.total_tokens) == (0, 0, 0) + assert response.usage.web_search_requests == 1 + + +@pytest.mark.parametrize("usage_metadata", ["not an object", ["not", "an", "object"], 7, True]) +def test_gemini_image_generation_response_rejects_non_object_usage_metadata_without_echoing_it( + usage_metadata: object, +): + with pytest.raises(ValidationError) as exc_info: + _transform_gemini_response({"candidates": [], "usageMetadata": usage_metadata}) + + assert "input_value" not in str(exc_info.value) diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py index fd667e8f425..8c73b72a65a 100644 --- a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock, Mock, patch import httpx import pytest +from pydantic import ValidationError import litellm from litellm.llms.vertex_ai.text_to_speech.transformation import ( @@ -633,3 +634,32 @@ def test_litellm_speech_vertex_ai_chirp(mock_get_token, mock_ensure_token, mock_ assert "headers" in call_kwargs assert "Authorization" in call_kwargs["headers"] assert call_kwargs["headers"]["Authorization"] == "Bearer mock-token" + + +@pytest.mark.parametrize("payload", [{}, {"audioContent": ""}, {"audioContent": None}, {"audioContent": []}]) +def test_transform_text_to_speech_response_without_audio_content_reports_it_missing(payload: dict[str, object]): + with pytest.raises(ValueError, match="No audioContent in Vertex AI TTS response"): + VertexAITextToSpeechConfig().transform_text_to_speech_response( + model="vertex_ai/chirp", + raw_response=httpx.Response(200, json=payload), + logging_obj=MagicMock(), + ) + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + {"audioContent": 7}, + {"audioContent": ["UklGRiQAAABXQVZFZm10IA=="]}, + ], +) +def test_transform_text_to_speech_response_rejects_malformed_payloads_without_echoing_them(payload: object): + with pytest.raises(ValidationError) as exc_info: + VertexAITextToSpeechConfig().transform_text_to_speech_response( + model="vertex_ai/chirp", + raw_response=httpx.Response(200, json=payload), + logging_obj=MagicMock(), + ) + + assert "input_value" not in str(exc_info.value) diff --git a/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py b/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py index 5eb4bf31845..777e52f4e2a 100644 --- a/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py +++ b/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest +from pydantic import ValidationError from litellm.llms.voyage.rerank.transformation import VoyageRerankConfig from litellm.types.rerank import RerankResponse @@ -343,3 +344,81 @@ class TestVoyageRerankTransform: assert prompt_cost == 0.0 assert completion_cost == 0.0 + + +def _transform(payload: object) -> RerankResponse: + return VoyageRerankConfig().transform_rerank_response( + model="rerank-2.5", + raw_response=httpx.Response(200, json=payload), + model_response=RerankResponse(), + logging_obj=MagicMock(), + ) + + +def test_transform_rerank_response_keeps_provider_id_and_drops_unknown_result_fields(): + response = _transform( + { + "id": "rerank-1", + "data": [ + {"index": 2, "relevance_score": 0.25, "document": {"text": "doc", "extra": 1}, "extra": 2}, + {"index": 0, "relevance_score": 1, "document": "plain"}, + ], + "usage": {"total_tokens": 9}, + } + ) + + assert response.id == "rerank-1" + assert response.results == [ + {"index": 2, "relevance_score": 0.25, "document": {"text": "doc"}}, + {"index": 0, "relevance_score": 1.0, "document": {"text": "plain"}}, + ] + + +@pytest.mark.parametrize( + ("payload", "expected_total_tokens"), + [ + ({"data": []}, 0), + ({"data": [], "usage": {}}, 0), + ({"data": [], "usage": {"total_tokens": 79}}, 79), + ({"data": [], "usage": {"total_tokens": None}}, None), + ], +) +def test_transform_rerank_response_reads_total_tokens_from_usage( + payload: dict[str, object], expected_total_tokens: int | None +): + assert _transform(payload).meta == { + "billed_units": {"total_tokens": expected_total_tokens}, + "tokens": {"input_tokens": expected_total_tokens, "output_tokens": 0}, + } + + +def test_transform_rerank_response_empty_data_list_yields_no_results(): + assert _transform({"id": "rerank-1", "data": []}).results == [] + + +@pytest.mark.parametrize( + "payload", + [ + ["not", "an", "object"], + {"data": 7}, + {"data": ["not an object"]}, + {"data": [{"index": 0, "relevance_score": 0.5}], "usage": None}, + {"data": [{"index": 0, "relevance_score": 0.5}], "usage": {"total_tokens": 1.5}}, + {"data": [{"index": 0, "relevance_score": 0.5}], "id": 7}, + ], +) +def test_transform_rerank_response_rejects_malformed_payloads(payload: object): + with pytest.raises(ValidationError): + _transform(payload) + + +def test_transform_rerank_response_result_without_index_raises_key_error(): + with pytest.raises(KeyError, match="index"): + _transform({"data": [{"relevance_score": 0.5}]}) + + +def test_transform_rerank_response_shape_errors_do_not_echo_the_payload(): + with pytest.raises(ValidationError) as exc_info: + _transform({"data": [["leaked document text"]]}) + + assert "leaked document text" not in str(exc_info.value) diff --git a/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py b/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py index efe592f515e..b6e5007c456 100644 --- a/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py +++ b/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py @@ -6,6 +6,10 @@ Validates the WatsonX transcription response transformation. from unittest.mock import MagicMock +import httpx +import pytest +from pydantic import ValidationError + from litellm.llms.watsonx.audio_transcription.transformation import ( IBMWatsonXAudioTranscriptionConfig, ) @@ -83,3 +87,57 @@ class TestWatsonXAudioTranscription: # Verify duration is set via dictionary assignment assert result["duration"] == 5.5 + + +def _transform_transcription_response(payload: object) -> TranscriptionResponse: + return IBMWatsonXAudioTranscriptionConfig().transform_audio_transcription_response( + httpx.Response(200, json=payload) + ) + + +@pytest.mark.parametrize( + ("payload", "expected_extras"), + [ + ({"text": "hello"}, {}), + ({"text": "hello", "model": "whisper-large-v3-turbo"}, {}), + ( + {"text": "hello", "duration": 1.5, "language": "en", "task": "transcribe"}, + {"duration": 1.5, "language": "en", "task": "transcribe"}, + ), + ( + {"text": "hello", "segments": [{"id": 0, "text": "hello"}], "words": None}, + {"segments": [{"id": 0, "text": "hello"}], "words": None}, + ), + ], +) +def test_transform_audio_transcription_response_copies_every_field_except_model( + payload: dict[str, object], expected_extras: dict[str, object] +) -> None: + response = _transform_transcription_response(payload) + + assert response.text == "hello" + assert not hasattr(response, "model") + assert {key: response[key] for key in expected_extras} == expected_extras + + +@pytest.mark.parametrize("payload", [{}, {"model": "whisper"}, {"duration": 2.0}, {"text": None, "usage": None}]) +def test_transform_audio_transcription_response_without_text_or_usage_reports_the_body( + payload: dict[str, object], +) -> None: + with pytest.raises(ValueError, match="Invalid response format") as exc_info: + _transform_transcription_response(payload) + + assert exc_info.value.args == ( + "Invalid response format. Received response does not match the expected format. Got: ", + payload, + ) + + +@pytest.mark.parametrize("payload", [["not", "an", "object"], "plain text", 7, True]) +def test_transform_audio_transcription_response_rejects_non_object_bodies_without_echoing_them( + payload: object, +) -> None: + with pytest.raises(ValidationError) as exc_info: + _transform_transcription_response(payload) + + assert "input_value" not in str(exc_info.value) diff --git a/tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py b/tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py index 58f6bb23498..afcebcf0990 100644 --- a/tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py +++ b/tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py @@ -1,8 +1,13 @@ +from unittest.mock import MagicMock + +import httpx import pytest +from pydantic import ValidationError from litellm.llms.watsonx.embed.transformation import IBMWatsonXEmbeddingConfig +from litellm.types.utils import EmbeddingResponse class TestIBMWatsonXEmbeddingConfig: @@ -42,3 +47,88 @@ class TestIBMWatsonXEmbeddingConfig: optional_params={"project_id": "test-project-id"}, headers={}, ) + + +def _transform(payload: object) -> EmbeddingResponse: + return IBMWatsonXEmbeddingConfig().transform_embedding_response( + model="ibm/slate-125m-english-rtrvr", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="test-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def test_transform_embedding_response_numbers_each_result_and_bills_the_input_tokens(): + response = _transform( + { + "model_id": "ibm/slate-125m-english-rtrvr", + "results": [{"embedding": [0.1, 2], "input": "ignored"}, {"embedding": [0.5]}], + "input_token_count": 7, + } + ) + + assert response.object == "list" + assert response.data == [ + {"object": "embedding", "index": 0, "embedding": [0.1, 2]}, + {"object": "embedding", "index": 1, "embedding": [0.5]}, + ] + assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (7, 0, 7) + + +@pytest.mark.parametrize( + ("payload", "expected_tokens"), + [ + ({"input_token_count": None}, 0), + ({"results": []}, 0), + ({}, 0), + ], +) +def test_transform_embedding_response_defaults_missing_results_and_token_count( + payload: dict[str, object], expected_tokens: int +): + response = _transform(payload) + + assert response.data == [] + assert (response.usage.prompt_tokens, response.usage.total_tokens) == (expected_tokens, expected_tokens) + + +@pytest.mark.parametrize( + "payload", + [ + {"results": None}, + {"results": 7}, + {"results": ["not an object"]}, + {"results": [{"embedding": [0.1]}, None]}, + {"input_token_count": 1.5}, + {"input_token_count": "many"}, + ], +) +def test_transform_embedding_response_rejects_malformed_payloads(payload: object): + with pytest.raises(ValidationError): + _transform(payload) + + +@pytest.mark.parametrize("payload", [["not", "an", "object"], "not an object"]) +def test_transform_embedding_response_non_object_body_raises_attribute_error(payload: object): + with pytest.raises(AttributeError): + _transform(payload) + + +def test_transform_embedding_response_result_without_embedding_raises_key_error(): + with pytest.raises(KeyError, match="embedding"): + _transform({"results": [{"input": "x"}]}) + + +@pytest.mark.parametrize( + "payload", + [{"results": ["leaked payload text"]}, {"results": {"leaked payload text": 1}}], +) +def test_transform_embedding_response_shape_errors_do_not_echo_the_payload(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "leaked payload text" not in str(exc_info.value) diff --git a/tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py b/tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py index ccbd318959f..fbdbddacd5f 100644 --- a/tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py +++ b/tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py @@ -9,6 +9,7 @@ from unittest.mock import MagicMock import httpx import pytest +from pydantic import ValidationError from litellm.llms.watsonx.common_utils import ( WatsonXAIError, @@ -260,3 +261,64 @@ class TestIBMWatsonXRerankTransform: assert "return_documents" in supported_params assert "max_tokens_per_doc" in supported_params assert len(supported_params) == 5 + + +def _transform(payload: object) -> RerankResponse: + return IBMWatsonXRerankConfig().transform_rerank_response( + model="watsonx/cross-encoder/ms-marco-minilm-l-12-v2", + raw_response=httpx.Response(200, json=payload), + model_response=RerankResponse(), + logging_obj=MagicMock(), + ) + + +def test_transform_rerank_response_keeps_the_upstream_id_documents_and_token_count(): + response = _transform( + { + "id": "rerank-1", + "results": [ + {"index": 1, "score": 0.5, "input": "plain text"}, + {"index": 0, "score": 0.25, "input": {"text": "object text"}}, + ], + "input_token_count": 62, + } + ) + + assert response.id == "rerank-1" + assert response.results == [ + {"index": 1, "relevance_score": 0.5, "document": {"text": "plain text"}}, + {"index": 0, "relevance_score": 0.25, "document": {"text": "object text"}}, + ] + assert response.meta == {"tokens": {"input_tokens": 62}} + + +def test_transform_rerank_response_without_a_token_count_reports_zero_tokens(): + response = _transform({"results": []}) + + assert response.results == [] + assert response.meta == {"tokens": {"input_tokens": 0}} + + +@pytest.mark.parametrize( + "payload", + [ + ["sensitive-document"], + "sensitive-document", + {"results": "sensitive-document"}, + {"results": 7}, + {"results": ["sensitive-document"]}, + {"results": [{"index": 0, "score": 0.5}, None]}, + {"results": [{"index": 0, "score": 0.5}], "id": 7}, + {"results": [{"index": 0, "score": 0.5}], "input_token_count": "many"}, + ], +) +def test_transform_rerank_response_rejects_malformed_payloads(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert "sensitive-document" not in str(exc_info.value) + + +def test_transform_rerank_response_requires_index_and_score_on_each_result(): + with pytest.raises(KeyError): + _transform({"results": [{"index": 0}]}) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py index cfcff73b857..2e255ddf853 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py @@ -173,6 +173,49 @@ def test_mcp_oauth_token_identity_detects_change_under_encryption(): assert mcp_oauth_token_identity(unchanged) != mcp_oauth_token_identity(changed) +def test_mcp_oauth_token_identity_reads_client_and_scopes_from_json_string_credentials(): + from litellm.proxy._experimental.mcp_server.db import mcp_oauth_token_identity + + credentials: Final = json.dumps( + {"client_id": "cid", "client_secret": "csec", "scopes": ["a"], "upstream_resource": "api://audience"} + ) + + assert mcp_oauth_token_identity(_identity_server(credentials=credentials)) == ( + "https://up.example.com/mcp", + None, + "oauth2", + "authorization_code", + None, + "https://idp.example.com/authorize", + "https://idp.example.com/token", + "https://idp.example.com/register", + "cid", + "csec", + ["a"], + "api://audience", + ) + + +@pytest.mark.parametrize("credentials", ["[]", '["client_id"]', '"cid"', "5", "null", "not json"]) +def test_mcp_oauth_token_identity_treats_stored_credentials_that_are_not_a_json_object_as_empty(credentials): + from litellm.proxy._experimental.mcp_server.db import mcp_oauth_token_identity + + assert mcp_oauth_token_identity(_identity_server(credentials=credentials)) == ( + "https://up.example.com/mcp", + None, + "oauth2", + "authorization_code", + None, + "https://idp.example.com/authorize", + "https://idp.example.com/token", + "https://idp.example.com/register", + None, + None, + None, + None, + ) + + def _oauth_row(user_id: str, server_id: str = "srv-1"): """A stored per-user OAuth token row (payload tagged type=oauth2, legacy plain-base64 encoding).""" row = _legacy_row(json.dumps({"type": "oauth2", "access_token": "tok-" + user_id})) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_env_vars.py index fff4221f243..6a08e8eeab6 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -1016,6 +1016,34 @@ async def test_get_user_env_vars_returns_empty_for_missing_row(): assert await get_user_env_vars(prisma, "alice", "srv-1") == {} +@pytest.mark.asyncio +async def test_get_user_env_vars_stringifies_non_string_json_values(env_vars_salt_key): + from types import SimpleNamespace + + from litellm.proxy._experimental.mcp_server.db import get_user_env_vars + + row: Final = SimpleNamespace(values_b64=_encrypted_user_env_blob({"PORT": 8080, "DEBUG": True, "EMPTY": None})) + + assert await get_user_env_vars(_mock_env_vars_prisma(row=row), "alice", "srv-1") == { + "PORT": "8080", + "DEBUG": "True", + "EMPTY": "None", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_json", ["[]", '[{"TOKEN": "t"}]', '"TOKEN"', "5", "null", "true", "not json"]) +async def test_get_user_env_vars_treats_a_blob_that_is_not_a_json_object_as_unset(env_vars_salt_key, stored_json): + from types import SimpleNamespace + + from litellm.proxy._experimental.mcp_server.db import get_user_env_vars + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + row: Final = SimpleNamespace(values_b64=encrypt_value_helper(stored_json)) + + assert await get_user_env_vars(_mock_env_vars_prisma(row=row), "alice", "srv-1") == {} + + @pytest.mark.asyncio async def test_decode_user_env_vars_warns_when_undecryptable( env_vars_salt_key, monkeypatch diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index 3a7792db1fe..b4a6aaf706e 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -7,6 +7,7 @@ from unittest.mock import Mock, patch import pytest +import responses from click.testing import CliRunner from litellm.constants import CLI_JWT_EXPIRATION_HOURS @@ -2044,3 +2045,33 @@ class TestGetStoredApiKeyRefresh: assert captured.out == "" assert captured.err == "Could not renew the key: token request failed with 503: temporarily_unavailable\n" save.assert_not_called() + + +@pytest.mark.parametrize( + "body, expected_detail", + [ + ('{"detail": "Too many CLI login attempts.", "retry": {"after": [30, null]}}', ": Too many CLI login attempts."), + ('{"detail": 5}', ""), + ('{"detail": ""}', ""), + ('{"detail": {"message": "nested"}}', ""), + ("{}", ""), + ('["detail"]', ""), + ('"detail"', ""), + ("7", ""), + ("true", ""), + ("null", ""), + ("gateway", ""), + ("", ""), + ], +) +@responses.activate +def test_login_shows_the_error_detail_only_when_the_proxy_answers_with_a_json_object(body, expected_detail): + responses.post("https://test.example.com/sso/cli/start", body=body, status=429) + + result = CliRunner().invoke(login, obj={"base_url": "https://test.example.com"}) + + assert result.exit_code == 0 + assert result.output == ( + "Authentication failed: Starting CLI login failed: HTTP 429 from https://test.example.com/sso/cli/start" + f"{expected_detail}\n" + ) diff --git a/tests/unit/proxy/client/test_chat.py b/tests/unit/proxy/client/test_chat.py index 8fe1bfcbb2f..9ab7150c288 100644 --- a/tests/unit/proxy/client/test_chat.py +++ b/tests/unit/proxy/client/test_chat.py @@ -7,6 +7,7 @@ import sys import pytest import requests +from pydantic import ValidationError from litellm.proxy.client.chat import ChatClient from litellm.proxy.client.exceptions import UnauthorizedError @@ -256,3 +257,53 @@ def test_completions_stream_gives_up_at_the_timeout_instead_of_hanging(hanging_s next(client.completions_stream(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}])) assert time.monotonic() - started < 10 + + +@pytest.mark.parametrize( + "body, expected", + [ + ( + b'data: {"choices": [{"delta": {"content": "Hel"}}]}\n\n' + b'data: {"choices": [{"delta": {"content": "lo"}}]}\n\n' + b"data: [DONE]\n\n", + [{"choices": [{"delta": {"content": "Hel"}}]}, {"choices": [{"delta": {"content": "lo"}}]}], + ), + ( + b': keep-alive\n\nevent: ping\n\ndata: not json\n\ndata: {"id": "caf\xc3\xa9"}\n\ndata: 7\n\n', + [{"id": "café"}, 7], + ), + (b'data: {"id": 1}\n\ndata: [DONE] \n\ndata: {"id": 2}\n\n', [{"id": 1}]), + (b"", []), + ], +) +@responses.activate +def test_completions_stream_yields_parsed_sse_chunks(client, base_url, sample_messages, body, expected): + responses.add(responses.POST, f"{base_url}/chat/completions", body=body, status=200) + + assert list(client.completions_stream(model="gpt-4", messages=sample_messages)) == expected + + +@responses.activate +def test_completions_stream_rejects_bytes_that_are_not_utf8(client, base_url, sample_messages): + responses.add(responses.POST, f"{base_url}/chat/completions", body=b"data: \xff\n\n", status=200) + + with pytest.raises(UnicodeDecodeError): + list(client.completions_stream(model="gpt-4", messages=sample_messages)) + + +@pytest.mark.parametrize("line", ['data: {"id": 1}', 5, memoryview(b'data: {"id": 1}')]) +@responses.activate +def test_completions_stream_rejects_lines_that_are_not_bytes(client, base_url, sample_messages, monkeypatch, line): + responses.add(responses.POST, f"{base_url}/chat/completions", body=b"", status=200) + monkeypatch.setattr(requests.Response, "iter_lines", lambda self: iter([line])) + + with pytest.raises(ValidationError): + list(client.completions_stream(model="gpt-4", messages=sample_messages)) + + +@responses.activate +def test_completions_stream_accepts_bytearray_lines(client, base_url, sample_messages, monkeypatch): + responses.add(responses.POST, f"{base_url}/chat/completions", body=b"", status=200) + monkeypatch.setattr(requests.Response, "iter_lines", lambda self: iter([bytearray(b'data: {"id": 1}')])) + + assert list(client.completions_stream(model="gpt-4", messages=sample_messages)) == [{"id": 1}] diff --git a/tests/unit/proxy/db/test_object_permission_repository.py b/tests/unit/proxy/db/test_object_permission_repository.py new file mode 100644 index 00000000000..e2d79592915 --- /dev/null +++ b/tests/unit/proxy/db/test_object_permission_repository.py @@ -0,0 +1,73 @@ +from collections.abc import Mapping +from types import SimpleNamespace +from typing import Final + +import pytest + +from litellm.repositories.object_permission_repository import ObjectPermissionRepository + + +class _RecordingPermissionTable: + def __init__(self, stored: Mapping[str, object]) -> None: + self.stored: Final = stored + self.created: Final[list[Mapping[str, object]]] = [] + self.updated: Final[list[tuple[Mapping[str, object], Mapping[str, object]]]] = [] + + async def create(self, data: Mapping[str, object]) -> Mapping[str, object]: + self.created.append(data) + return {**self.stored, **data} + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> Mapping[str, object]: + self.updated.append((where, data)) + return {**self.stored, **data} + + +def _repository(table: _RecordingPermissionTable) -> ObjectPermissionRepository: + return ObjectPermissionRepository(SimpleNamespace(db=SimpleNamespace(litellm_objectpermissiontable=table))) + + +@pytest.mark.asyncio +async def test_create_permission_writes_only_the_fields_it_was_given() -> None: + table: Final = _RecordingPermissionTable({"object_permission_id": "perm-1"}) + + permission: Final = await _repository(table).create_permission( + mcp_servers=["server-1"], mcp_tool_permissions={"server-1": ["search"]}, models=[] + ) + + assert table.created == [ + {"mcp_servers": ["server-1"], "mcp_tool_permissions": {"server-1": ["search"]}, "models": []} + ] + assert permission.model_dump(exclude_unset=True) == { + "object_permission_id": "perm-1", + "mcp_servers": ["server-1"], + "mcp_tool_permissions": {"server-1": ["search"]}, + "models": [], + } + + +@pytest.mark.asyncio +async def test_create_permission_without_fields_writes_an_empty_row() -> None: + table: Final = _RecordingPermissionTable({"object_permission_id": "perm-1"}) + + permission: Final = await _repository(table).create_permission() + + assert table.created == [{}] + assert permission.model_dump(exclude_unset=True) == {"object_permission_id": "perm-1"} + + +@pytest.mark.asyncio +async def test_update_permission_changes_only_the_fields_it_was_given() -> None: + table: Final = _RecordingPermissionTable( + {"object_permission_id": "perm-1", "models": ["gpt-4o"], "agents": ["agent-1"]} + ) + + permission: Final = await _repository(table).update_permission("perm-1", models=["gpt-4o-mini"], skills=[]) + + assert table.updated == [({"object_permission_id": "perm-1"}, {"models": ["gpt-4o-mini"], "skills": []})] + assert permission is not None + assert permission.model_dump(exclude_unset=True) == { + "object_permission_id": "perm-1", + "models": ["gpt-4o-mini"], + "agents": ["agent-1"], + "skills": [], + } diff --git a/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py index 7ed1a436cb6..d1b110f288a 100644 --- a/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py +++ b/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py @@ -13,9 +13,12 @@ seam stayed untouched, so a guard that raises after the provider call would stil import base64 from contextlib import ExitStack from dataclasses import dataclass +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi import Response @@ -273,3 +276,68 @@ async def test_cancel__unified_job_id_allowed_when_managed_files_required(seams) await _cancel(_unified_job_id()) assert seams.router.acancel_fine_tuning_job.call_count == 1 + + +FINE_TUNING_API_BASE: Final = "https://fine-tuning.test/v1" +PROVIDER_JOBS_PAGE: Final = { + "object": "list", + "data": [ + { + "id": RAW_JOB_ID, + "created_at": 1234567890, + "fine_tuned_model": None, + "finished_at": None, + "hyperparameters": {"n_epochs": 1}, + "model": "gpt-4o-mini", + "object": "fine_tuning.job", + "organization_id": "org-test", + "result_files": [], + "seed": 0, + "status": "running", + "trained_tokens": None, + "training_file": RAW_FILE_ID, + "validation_file": None, + } + ], + "has_more": False, +} + + +async def _list(custom_llm_provider: str | None): + return await endpoints.list_fine_tuning_jobs( + request=FakeRequest(), + fastapi_response=Response(), + custom_llm_provider=custom_llm_provider, + target_model_names=None, + after=None, + limit=None, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +@pytest.mark.asyncio +async def test_list__rejected_when_neither_a_provider_nor_a_target_model_is_named(seams): + with pytest.raises(ProxyException) as exc: + await _list(custom_llm_provider=None) + + assert exc.value.code == "400" + assert exc.value.message == "Invalid request, No litellm managed file id or custom_llm_provider provided." + + +@pytest.mark.asyncio +@respx.mock +async def test_list__returns_the_jobs_the_configured_provider_lists(seams): + respx.get(f"{FINE_TUNING_API_BASE}/fine_tuning/jobs").mock( + return_value=httpx.Response(200, json=PROVIDER_JOBS_PAGE) + ) + provider_config: Final = { + "custom_llm_provider": "openai", + "api_key": "sk-provider", + "api_base": FINE_TUNING_API_BASE, + } + + with patch.object(endpoints, "fine_tuning_config", [provider_config]): + page: Final = await _list(custom_llm_provider="openai") + + assert [job.id for job in page.data] == [RAW_JOB_ID] + assert page.has_more is False diff --git a/tests/unit/proxy/management_endpoints/policy_endpoints/test_endpoints.py b/tests/unit/proxy/management_endpoints/policy_endpoints/test_endpoints.py index 86f9aeafc08..62863163b8e 100644 --- a/tests/unit/proxy/management_endpoints/policy_endpoints/test_endpoints.py +++ b/tests/unit/proxy/management_endpoints/policy_endpoints/test_endpoints.py @@ -6,11 +6,21 @@ without needing a running proxy. """ import pytest +from fastapi import Request +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.management_endpoints.policy_endpoints.endpoints import ( GuardrailTestResultEntry, _compute_overall_action, _test_guardrail_definitions, + list_policies, +) +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.types.proxy.policy_engine import ( + PolicyGuardrailsResponse, + PolicyListResponse, + PolicyScopeResponse, + PolicySummaryItem, ) @@ -293,3 +303,68 @@ class TestEnrichPolicyTemplateStreamKeepalive: assert b": ping\n\n" not in chunks assert chunks[0] == b'data: {"type": "competitor", "name": "Rival Air"}\n\n' assert chunks[-1].startswith(b'data: {"type": "done"') + + +def _policy_list_request() -> Request: + return Request({"type": "http", "method": "GET", "path": "/policy/list", "headers": []}) + + +def _summary_item(inherit, resolved_guardrails, inheritance_chain) -> PolicySummaryItem: + return PolicySummaryItem( + inherit=inherit, + scope=PolicyScopeResponse(), + guardrails=PolicyGuardrailsResponse(), + resolved_guardrails=resolved_guardrails, + inheritance_chain=inheritance_chain, + ) + + +@pytest.fixture +def policy_registry(): + registry = get_policy_registry() + registry.clear() + yield registry + registry.clear() + + +@pytest.mark.asyncio +async def test_list_policies_is_empty_before_any_policy_is_loaded(policy_registry): + response = await list_policies(request=_policy_list_request(), user_api_key_dict=UserAPIKeyAuth()) + + assert response == PolicyListResponse(policies={}, total_count=0) + + +@pytest.mark.parametrize( + ("policies_config", "expected_policies"), + [ + ({}, {}), + ( + {"solo": {"description": "standalone", "guardrails": {"add": ["pii"]}}}, + {"solo": _summary_item(None, ["pii"], ["solo"])}, + ), + ( + { + "base": {"guardrails": {"add": ["pii"]}}, + "child": {"inherit": "base", "guardrails": {"add": ["audit"], "remove": ["pii"]}}, + "conditional": {"guardrails": {"add": ["toxicity"]}, "condition": {"model": "gpt-4.*"}}, + "empty": {"guardrails": {"remove": ["pii"]}}, + }, + { + "base": _summary_item(None, ["pii"], ["base"]), + "child": _summary_item("base", ["audit"], ["base", "child"]), + "conditional": _summary_item(None, ["toxicity"], ["conditional"]), + "empty": _summary_item(None, [], ["empty"]), + }, + ), + ], +) +@pytest.mark.asyncio +async def test_list_policies_reports_each_loaded_policy_with_its_resolved_guardrails( + policy_registry, policies_config, expected_policies +): + policy_registry.load_policies(policies_config) + + response = await list_policies(request=_policy_list_request(), user_api_key_dict=UserAPIKeyAuth()) + + assert response == PolicyListResponse(policies=expected_policies, total_count=0) + assert list(response.policies) == list(expected_policies) diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 8159890ef16..452fdec648f 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -6167,6 +6167,53 @@ class TestStrategyRouterWriteValidation: async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id) as table: assert hasattr(table, "create") + @pytest.mark.asyncio + async def test_slot_counts_tuned_routers_whose_stored_params_are_json_strings( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from fastapi import HTTPException + + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot + from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + + fake = self._FakeDb([]) + fake.tx_obj.litellm_proxymodeltable.find_many = AsyncMock( + return_value=[ + { + "model_id": "a", + "model_name": "router-a", + "litellm_params": json.dumps( + {"model": "auto_router/complexity_router", "complexity_router_config": self._TUNED_A_EDITED} + ), + "model_info": json.dumps({"id": "a"}), + } + ] + ) + monkeypatch.setattr(proxy_server._license_check, "auto_router_capability_limit", lambda: 1) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr( + proxy_server, + "heuristic_v1_tuning_baselines", + snapshot_tuning_baselines( + [self._db_router_row("a", self._TUNED_A), self._db_router_row("b", self._TUNED_B)] + ), + ) + + with pytest.raises(HTTPException) as refused: + async with _auto_router_capability_slot( + fake, + effective_params={ + "model": "auto_router/complexity_router", + "complexity_router_config": self._TUNED_B_EDITED, + }, + model_id="b", + ): + pass + + assert refused.value.status_code == 403 + assert "changed heuristic scoring rules" in str(refused.value.detail) + @pytest.mark.asyncio async def test_add_new_model_refuses_a_second_tuned_heuristic_v1_router_without_a_model_id(self) -> None: """A create request carries no model_info at all, yet the quota still judges it: Deployment mints the diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 440f93d1387..7d3bd049a03 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1219,6 +1219,34 @@ async def test_legacy_update_without_budget_fields_skips_budget_write(monkeypatc assert prisma.db.litellm_organizationtable.update.await_args.kwargs["data"]["organization_alias"] == "renamed" +@pytest.mark.asyncio +async def test_legacy_update_writes_sent_metadata_to_the_organization_row(monkeypatch): + prisma = await _run_legacy_update_organization( + monkeypatch, + body={"organization_id": "org-1", "metadata": {"team": "search", "limits": {"rpm": 5}}}, + existing_budget_id="budget-1", + ) + + organization_write = prisma.db.litellm_organizationtable.update.await_args + assert organization_write.kwargs["where"] == {"organization_id": "org-1"} + assert json.loads(organization_write.kwargs["data"]["metadata"]) == {"team": "search", "limits": {"rpm": 5}} + prisma.db.litellm_budgettable.update.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [[], ["litellm_budget_table"], "organization_id", 5, 1.5, True, None]) +async def test_legacy_update_rejects_json_body_that_is_not_an_object_before_reading_the_database(monkeypatch, body): + from pydantic import ValidationError + + from litellm.proxy import proxy_server + + with pytest.raises(ValidationError): + await _run_legacy_update_organization(monkeypatch, body=body, existing_budget_id="budget-1") + + proxy_server.prisma_client.db.litellm_organizationtable.find_unique.assert_not_awaited() + proxy_server.prisma_client.db.litellm_organizationtable.update.assert_not_awaited() + + def test_build_budget_write_data_recomputes_reset_at_on_duration(): """A sent budget_duration recomputes budget_reset_at so the reset window follows the new duration.""" from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index 14ee9db6ffd..d5dbd3df6a6 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -2,6 +2,7 @@ import inspect import json from collections.abc import Mapping, Sequence from contextlib import contextmanager +from datetime import datetime from types import MappingProxyType, SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, Mock, patch @@ -1521,3 +1522,67 @@ async def test_add_tag_to_deployment_model_not_found(): assert exc_info.value.status_code == 500 assert "not found in database" in str(exc_info.value.detail) + + +class _StoredTagTable: + def __init__(self, model_info: object) -> None: + self.model_info = model_info + + async def find_many(self, where: object = None, include: object = None) -> list[SimpleNamespace]: + return [ + SimpleNamespace( + tag_name="routed-tag", + description="Routes to one model", + models=["model-1"], + model_info=self.model_info, + budget_id=None, + created_at=datetime(2025, 1, 1), + updated_at=datetime(2025, 1, 2), + created_by="user-123", + litellm_budget_table=None, + ) + ] + + +class _NoDynamicTagSpend: + async def group_by(self, by: object, where: object, min: object, max: object) -> list[object]: + return [] + + +@pytest.mark.parametrize( + ("stored_model_info", "returned_model_info"), + [ + ('{"model-1": "gpt-4o"}', {"model-1": "gpt-4o"}), + ({"model-1": "gpt-4o"}, {"model-1": "gpt-4o"}), + (None, {}), + ], +) +def test_tag_info_and_tag_list_return_the_stored_model_info_decoded( + monkeypatch, stored_model_info, returned_model_info +): + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace( + db=SimpleNamespace( + litellm_tagtable=_StoredTagTable(stored_model_info), + litellm_dailytagspend=_NoDynamicTagSpend(), + ) + ), + ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + info_response = client.post("/tag/info", json={"names": ["routed-tag"]}) + list_response = client.get("/tag/list") + finally: + app.dependency_overrides.clear() + + assert info_response.status_code == 200 + assert info_response.json()["routed-tag"]["model_info"] == returned_model_info + assert list_response.status_code == 200 + assert [tag["model_info"] for tag in list_response.json()] == [returned_model_info] diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 094320225cd..1c9bd470e52 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2533,3 +2533,22 @@ class TestRecordPartialUsageForFailure: assert "combined_usage_object" not in logging_obj.model_call_details assert "response_cost" not in logging_obj.model_call_details + + +@pytest.mark.parametrize( + ("all_chunks", "interrupted"), + [ + (['data: {"type": "content_block_delta"}'], True), + (['data: {"type": "message_delta"}', 'data: {"type": "message_stop"}'], False), + ([b'data: {"type": "message_delta"}\ndata: {"type": "message_stop"}\n'], False), + (['data: {"type": "message_delta"}', "data: [1, 2]", 'data: "text"', "data: 7", "data: null"], False), + (['data: {"type": "content_block_stop"}', "data: [1, 2]", "data: not json"], True), + (['data: {"type": "message_start"}', 'data: {"type": ["message_delta"]}', "data: {}"], True), + (["data: [1, 2]", "data: null", "event: message_delta"], True), + ([], True), + ], +) +def test_stream_was_interrupted_skips_data_lines_that_are_not_json_objects( + all_chunks: list[str | bytes], interrupted: bool +): + assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is interrupted diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py new file mode 100644 index 00000000000..f9df5a120f5 --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py @@ -0,0 +1,95 @@ +from datetime import datetime + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( + VertexPassthroughLoggingHandler, +) +from litellm.types.utils import EmbeddingResponse, ModelResponse + +_PREDICT_ROUTE = "/v1/projects/p/locations/us-central1/publishers/google/models/text-embedding-004:predict" +_INTERACTIONS_ROUTE = "https://aiplatform.googleapis.com/v1beta1/projects/p/locations/global/interactions" + + +def _handle(url_route: str, payload: object) -> PassThroughEndpointLoggingTypedDict: + logging_obj = Logging( + model="unknown", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime(2026, 1, 1), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.optional_params = {} + response = httpx.Response(200, json=payload) + return VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=response, + logging_obj=logging_obj, + url_route=url_route, + result=response.text, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + cache_hit=False, + request_body={"model": "gemini-omni-flash-preview"}, + ) + + +def test_predict_response_with_text_embeddings_is_logged_as_an_embedding_response(): + result = _handle( + _PREDICT_ROUTE, + { + "predictions": [ + {"embeddings": {"values": [0.1, 0.2], "statistics": {"token_count": 3}}}, + {"embeddings": {"values": [0.3, 0.4], "statistics": {"token_count": 4}}}, + ], + "metadata": {"billableCharacterCount": 9}, + }, + ) + + response = result["result"] + assert isinstance(response, EmbeddingResponse) + assert response.data == [ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, + {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]}, + ] + assert response.usage.prompt_tokens == 7 + assert result["kwargs"]["model"] == "text-embedding-004" + assert result["kwargs"]["custom_llm_provider"] == "vertex_ai" + + +@pytest.mark.parametrize("payload", [["not", "an", "object"], "text", 7, 1.5, True]) +def test_predict_response_that_is_not_a_json_object_is_rejected_without_echoing_it(payload: object): + with pytest.raises(ValidationError) as exc_info: + _handle(_PREDICT_ROUTE, payload) + + assert "input_value" not in str(exc_info.value) + + +def test_interactions_usage_object_is_read_into_prompt_and_completion_tokens(): + result = _handle( + _INTERACTIONS_ROUTE, + { + "id": "interactions/abc", + "model": "gemini-omni-flash-preview", + "usage": { + "total_tokens": 41, + "total_input_tokens": 12, + "input_tokens_by_modality": [{"modality": "text", "tokens": 12}], + "total_output_tokens": 9, + "output_tokens_by_modality": [{"modality": "text", "tokens": 9}], + "total_thought_tokens": 20, + }, + }, + ) + + response = result["result"] + assert isinstance(response, ModelResponse) + assert response.usage.prompt_tokens == 12 + assert response.usage.completion_tokens == 29 + assert response.usage.completion_tokens_details.text_tokens == 9 + assert result["kwargs"]["custom_llm_provider"] == "vertex_ai" diff --git a/tests/unit/proxy/policy_engine/test_policy_registry.py b/tests/unit/proxy/policy_engine/test_policy_registry.py new file mode 100644 index 00000000000..db7775b6238 --- /dev/null +++ b/tests/unit/proxy/policy_engine/test_policy_registry.py @@ -0,0 +1,154 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timezone +from types import SimpleNamespace + +import pytest + +from litellm.proxy.policy_engine.policy_registry import PolicyRegistry +from litellm.types.proxy.policy_engine import PolicyCreateRequest, PolicyUpdateRequest + +_NOW = datetime(2026, 1, 1, tzinfo=timezone.utc) + +_REQUESTED_PIPELINE = {"mode": "pre_call", "steps": [{"guardrail": "pii-guard", "on_fail": "next"}]} + +_STORED_PIPELINE = { + "mode": "pre_call", + "steps": [ + { + "guardrail": "pii-guard", + "on_fail": "next", + "on_pass": "allow", + "on_error": None, + "pass_data": False, + "modify_response_message": None, + } + ], +} + +_MALFORMED_PIPELINES = [ + pytest.param({"mode": "pre_call", "steps": []}, "steps", id="no-steps"), + pytest.param({"mode": "pre_call"}, "steps", id="steps-missing"), + pytest.param({"steps": [{"guardrail": "pii-guard"}]}, "mode", id="mode-missing"), + pytest.param({"mode": "during_call", "steps": [{"guardrail": "pii-guard"}]}, "mode", id="unknown-mode"), + pytest.param( + {"mode": "pre_call", "steps": [{"guardrail": "pii-guard"}], "order": "strict"}, "order", id="unknown-key" + ), + pytest.param({"mode": "pre_call", "steps": "pii-guard"}, "steps", id="steps-not-a-list"), +] + + +class _PolicyTable: + def __init__(self) -> None: + self.writes: list[Mapping[str, object]] = [] + + def _row(self, data: Mapping[str, object]) -> SimpleNamespace: + pipeline = data.get("pipeline") + return SimpleNamespace( + policy_id="policy-1", + policy_name=data.get("policy_name", "pii-policy"), + version_number=1, + version_status=data.get("version_status", "draft"), + parent_version_id=None, + is_latest=True, + published_at=None, + production_at=None, + inherit=None, + description=None, + guardrails_add=data.get("guardrails_add", []), + guardrails_remove=data.get("guardrails_remove", []), + condition=None, + pipeline=json.loads(pipeline) if isinstance(pipeline, str) else None, + created_at=_NOW, + updated_at=_NOW, + created_by=None, + updated_by=None, + ) + + async def create(self, data: Mapping[str, object]) -> SimpleNamespace: + self.writes.append(data) + return self._row(data) + + async def find_unique(self, where: Mapping[str, object]) -> SimpleNamespace: + return self._row({"policy_name": "pii-policy", "version_status": "draft"}) + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> SimpleNamespace: + self.writes.append(data) + return self._row(data) + + +def _prisma_client(table: _PolicyTable) -> SimpleNamespace: + return SimpleNamespace(db=SimpleNamespace(litellm_policytable=table)) + + +@pytest.mark.asyncio +async def test_add_policy_to_db_stores_the_pipeline_with_step_defaults_filled_in(): + table = _PolicyTable() + registry = PolicyRegistry() + + response = await registry.add_policy_to_db( + PolicyCreateRequest(policy_name="pii-policy", guardrails_add=["pii-guard"], pipeline=_REQUESTED_PIPELINE), + _prisma_client(table), + ) + + assert [json.loads(str(write["pipeline"])) for write in table.writes] == [_STORED_PIPELINE] + assert response.pipeline == _STORED_PIPELINE + stored_policy = registry.get_policy("pii-policy") + assert stored_policy is not None + assert stored_policy.pipeline is not None + assert stored_policy.pipeline.model_dump() == _STORED_PIPELINE + + +@pytest.mark.asyncio +async def test_update_policy_in_db_stores_the_pipeline_with_step_defaults_filled_in(): + table = _PolicyTable() + + response = await PolicyRegistry().update_policy_in_db( + "policy-1", + PolicyUpdateRequest(pipeline=_REQUESTED_PIPELINE), + _prisma_client(table), + ) + + assert [json.loads(str(write["pipeline"])) for write in table.writes] == [_STORED_PIPELINE] + assert response.pipeline == _STORED_PIPELINE + + +@pytest.mark.parametrize(("pipeline", "rejected_field"), _MALFORMED_PIPELINES) +@pytest.mark.asyncio +async def test_add_policy_to_db_rejects_a_malformed_pipeline_before_writing( + pipeline: dict[str, object], rejected_field: str +): + table = _PolicyTable() + registry = PolicyRegistry() + + with pytest.raises(Exception, match="Error adding policy to DB: 1 validation error") as raised: + await registry.add_policy_to_db( + PolicyCreateRequest(policy_name="pii-policy", pipeline=pipeline), + _prisma_client(table), + ) + + assert str(raised.value).startswith( + f"Error adding policy to DB: 1 validation error for GuardrailPipeline\n{rejected_field}\n" + ) + assert table.writes == [] + assert registry.get_policy("pii-policy") is None + + +@pytest.mark.parametrize(("pipeline", "rejected_field"), _MALFORMED_PIPELINES) +@pytest.mark.asyncio +async def test_update_policy_in_db_rejects_a_malformed_pipeline_before_writing( + pipeline: dict[str, object], rejected_field: str +): + table = _PolicyTable() + + with pytest.raises(Exception, match="Error updating policy in DB: 1 validation error") as raised: + await PolicyRegistry().update_policy_in_db( + "policy-1", + PolicyUpdateRequest(pipeline=pipeline), + _prisma_client(table), + ) + + assert str(raised.value).startswith( + f"Error updating policy in DB: 1 validation error for GuardrailPipeline\n{rejected_field}\n" + ) + assert table.writes == [] diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 9df8e6f4d67..024fe5229be 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -11,6 +11,8 @@ from litellm.caching import DualCache from litellm.models.budget import LiteLLM_BudgetTable from litellm.proxy import proxy_server from litellm.proxy._types import ( + LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, @@ -18,11 +20,13 @@ from litellm.proxy._types import ( ) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + project_cache_key, team_membership_reservation_cache_key, ) from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, estimate_request_max_cost, + get_budget_window_start, release_unbound_budget_reservation, reserve_budget_for_request, ) @@ -342,3 +346,83 @@ async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_call assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_reservation_holds_cost_against_cached_org_and_project_budgets(spend_counter_cache: DualCache): + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key="org_id:org-budgeted:with_budget", + value=LiteLLM_OrganizationTable( + organization_id="org-budgeted", + budget_id="org-budget", + spend=1.5, + models=[], + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=10.0), + ), + model_type=LiteLLM_OrganizationTable, + ) + await cache.async_set_cache( + key=project_cache_key("project-budgeted"), + value=LiteLLM_ProjectTableCachedObj( + project_id="project-budgeted", spend=2.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=20.0) + ), + model_type=LiteLLM_ProjectTableCachedObj, + ) + + reservation: Final = await reserve_budget_for_request( + request_body={"model": "gpt-4o", "input": "hello"}, + route="/v1/responses", + llm_router=None, + valid_token=UserAPIKeyAuth(token="hashed-org-project", org_id="org-budgeted", project_id="project-budgeted"), + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + assert reservation is not None + reserved_cost: Final = reservation["reserved_cost"] + assert reserved_cost > 0 + assert reservation["entries"] == [ + { + "counter_key": "spend:org:org-budgeted", + "entity_type": "Organization", + "entity_id": "org-budgeted", + "reserved_cost": reserved_cost, + "applied_adjustment": 0.0, + }, + { + "counter_key": "spend:project:project-budgeted", + "entity_type": "Project", + "entity_id": "project-budgeted", + "reserved_cost": reserved_cost, + "applied_adjustment": 0.0, + }, + ] + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-budgeted") == pytest.approx(reserved_cost) + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:project:project-budgeted") == pytest.approx( + reserved_cost + ) + + +@pytest.mark.parametrize( + ("window", "expected"), + [ + ( + '{"budget_duration": "1h", "reset_at": "2030-01-01T01:00:00Z"}', + datetime(2030, 1, 1, 0, 0, tzinfo=timezone.utc), + ), + ('{"reset_at": "2030-01-01T01:00:00Z"}', None), + ("{}", None), + ('["budget_duration", "1h"]', None), + ('"1h"', None), + ("null", None), + ("not json", None), + ], +) +def test_budget_window_start_reads_json_encoded_windows(window: str, expected: datetime | None) -> None: + assert get_budget_window_start(window) == expected diff --git a/tests/unit/rag/ingestion/test_base_ingestion.py b/tests/unit/rag/ingestion/test_base_ingestion.py new file mode 100644 index 00000000000..d20b46d356b --- /dev/null +++ b/tests/unit/rag/ingestion/test_base_ingestion.py @@ -0,0 +1,101 @@ +from collections import UserString +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion +from litellm.types.utils import CredentialItem + +_FILE_URL: Final = "https://files.example/docs/report.pdf" +_STORED_CREDENTIAL: Final = CredentialItem( + credential_name="gemini-prod", + credential_info={}, + credential_values={"api_key": "stored-key"}, +) + + +def _vector_store_options(litellm_credential_name: object) -> dict[str, object]: + return { + "custom_llm_provider": "gemini", + "litellm_credential_name": litellm_credential_name, + "api_key": "caller-key", + "api_base": "https://caller.example", + } + + +def test_a_stored_credential_named_by_the_vector_store_replaces_the_caller_supplied_values( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "credential_list", [_STORED_CREDENTIAL]) + vector_store: Final = _vector_store_options("gemini-prod") + + ingestion: Final = GeminiRAGIngestion(ingest_options={"vector_store": vector_store}) + + assert ingestion.vector_store_config is vector_store + assert vector_store == { + "custom_llm_provider": "gemini", + "litellm_credential_name": "gemini-prod", + "api_key": "stored-key", + } + + +@pytest.mark.parametrize( + "litellm_credential_name", + ["gemini-staging", "", None, 7, True, ["gemini-prod"], {"credential_name": "gemini-prod"}], +) +def test_a_credential_name_that_matches_no_stored_credential_leaves_the_vector_store_config_alone( + monkeypatch: pytest.MonkeyPatch, litellm_credential_name: object +): + monkeypatch.setattr(litellm, "credential_list", [_STORED_CREDENTIAL]) + + ingestion: Final = GeminiRAGIngestion( + ingest_options={"vector_store": _vector_store_options(litellm_credential_name)} + ) + + assert ingestion.vector_store_config == _vector_store_options(litellm_credential_name) + + +def test_a_credential_name_that_is_not_a_string_is_not_resolved_even_when_it_equals_a_stored_name( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "credential_list", [_STORED_CREDENTIAL]) + litellm_credential_name: Final = UserString("gemini-prod") + + ingestion: Final = GeminiRAGIngestion( + ingest_options={"vector_store": _vector_store_options(litellm_credential_name)} + ) + + assert litellm_credential_name == "gemini-prod" + assert ingestion.vector_store_config == _vector_store_options(litellm_credential_name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("headers", "expected_content_type"), + [ + ([("content-type", "application/pdf")], "application/pdf"), + ([("Content-Type", "text/plain; charset=utf-8")], "text/plain; charset=utf-8"), + ([("content-type", "")], ""), + ([("content-type", "text/plain"), ("content-type", "text/html")], "text/plain, text/html"), + ([], "application/octet-stream"), + ([("content-length", "8")], "application/octet-stream"), + ], +) +async def test_upload_from_a_url_takes_the_content_type_from_the_response_header( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + headers: list[tuple[str, str]], + expected_content_type: str, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "user_url_validation", False) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.get(_FILE_URL).mock(return_value=httpx.Response(200, content=b"%PDF-1.7", headers=headers)) + ingestion: Final = GeminiRAGIngestion(ingest_options={"vector_store": {"custom_llm_provider": "gemini"}}) + + uploaded: Final = await ingestion.upload(file_url=_FILE_URL) + + assert uploaded == ("report.pdf", b"%PDF-1.7", expected_content_type, None) diff --git a/tests/unit/rag/ingestion/test_gemini_ingestion.py b/tests/unit/rag/ingestion/test_gemini_ingestion.py new file mode 100644 index 00000000000..8122a9d334b --- /dev/null +++ b/tests/unit/rag/ingestion/test_gemini_ingestion.py @@ -0,0 +1,111 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion + +_API_BASE: Final = "https://gemini.example" +_STORE: Final = "fileSearchStores/docs-1" +_START_UPLOAD_URL: Final = f"{_API_BASE}/upload/v1beta/{_STORE}:uploadToFileSearchStore" +_UPLOAD_SESSION_URL: Final = "https://gemini.example/upload/session-1" +_DOCUMENT: Final = "fileSearchStores/docs-1/documents/notes-1" + + +def _ingestion(chunking_strategy: object) -> GeminiRAGIngestion: + return GeminiRAGIngestion( + ingest_options={ + "chunking_strategy": chunking_strategy, + "vector_store": { + "custom_llm_provider": "gemini", + "vector_store_id": _STORE, + "api_key": "test-key", + "api_base": _API_BASE, + }, + } + ) + + +async def _store_notes(ingestion: GeminiRAGIngestion) -> tuple[str | None, str | None]: + return await ingestion.store( + file_content=b"first note", + filename="notes.txt", + content_type="text/plain", + chunks=[], + embeddings=None, + ) + + +@pytest.fixture +def start_upload(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.put(_UPLOAD_SESSION_URL).mock(return_value=httpx.Response(200, json={"name": _DOCUMENT})) + return respx_mock.post(_START_UPLOAD_URL).mock( + return_value=httpx.Response(200, headers={"x-goog-upload-url": _UPLOAD_SESSION_URL}) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("white_space_config", "expected"), + [ + ( + {"max_tokens_per_chunk": 200, "max_overlap_tokens": 20}, + {"maxTokensPerChunk": 200, "maxOverlapTokens": 20}, + ), + ({"max_tokens_per_chunk": 200}, {"maxTokensPerChunk": 200, "maxOverlapTokens": 400}), + ({"unrelated": True}, {"maxTokensPerChunk": 800, "maxOverlapTokens": 400}), + ( + {"max_tokens_per_chunk": "200", "max_overlap_tokens": None}, + {"maxTokensPerChunk": "200", "maxOverlapTokens": None}, + ), + ({7: "not a field", "max_overlap_tokens": 1.5}, {"maxTokensPerChunk": 800, "maxOverlapTokens": 1.5}), + (MappingProxyType({"max_overlap_tokens": 0}), {"maxTokensPerChunk": 800, "maxOverlapTokens": 0}), + ], +) +async def test_white_space_config_is_sent_as_the_chunking_config_of_the_upload( + start_upload: respx.Route, white_space_config: Mapping[object, object], expected: Mapping[str, object] +): + stored: Final = await _store_notes(_ingestion({"white_space_config": white_space_config})) + + assert stored == (_STORE, _DOCUMENT) + assert json.loads(start_upload.calls.last.request.content) == { + "displayName": "notes.txt", + "chunkingConfig": {"whiteSpaceConfig": expected}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunking_strategy", + [None, {"type": "auto"}, {"white_space_config": None}, {"white_space_config": {}}, {"white_space_config": 0}], +) +async def test_upload_without_a_white_space_config_sends_no_chunking_config( + start_upload: respx.Route, chunking_strategy: object +): + stored: Final = await _store_notes(_ingestion(chunking_strategy)) + + assert stored == (_STORE, _DOCUMENT) + assert json.loads(start_upload.calls.last.request.content) == {"displayName": "notes.txt"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "white_space_config", + ["private-setting", [800, 400], [("max_tokens_per_chunk", 200)], 800, True, 1.5], +) +async def test_a_white_space_config_that_is_not_a_mapping_is_rejected_before_any_upload( + start_upload: respx.Route, white_space_config: object +): + with pytest.raises(ValidationError) as raised: + await _store_notes(_ingestion({"white_space_config": white_space_config})) + + assert "private-setting" not in str(raised.value) + assert not start_upload.called diff --git a/tests/unit/rag/ingestion/test_openai_ingestion.py b/tests/unit/rag/ingestion/test_openai_ingestion.py new file mode 100644 index 00000000000..74b68c9a517 --- /dev/null +++ b/tests/unit/rag/ingestion/test_openai_ingestion.py @@ -0,0 +1,101 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion + +_API_BASE: Final = "https://openai.example/v1" +_STATIC_CHUNKING: Final = {"type": "static", "static": {"max_chunk_size_tokens": 800, "chunk_overlap_tokens": 400}} +_ATTACHED_FILE: Final = { + "id": "file-1", + "object": "vector_store.file", + "created_at": 1767323045, + "vector_store_id": "vs_1", + "status": "completed", + "usage_bytes": 10, +} +_UPLOADED_FILE: Final = { + "id": "file-1", + "object": "file", + "bytes": 10, + "created_at": 1767323045, + "filename": "notes.txt", + "purpose": "assistants", + "status": "processed", +} + + +def _ingestion(chunking_strategy: object) -> OpenAIRAGIngestion: + return OpenAIRAGIngestion( + ingest_options={ + "chunking_strategy": chunking_strategy, + "vector_store": { + "custom_llm_provider": "openai", + "vector_store_id": "vs_1", + "api_key": "sk-test", + "api_base": _API_BASE, + }, + } + ) + + +@pytest.fixture +def attach_file(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + return respx_mock.post(f"{_API_BASE}/vector_stores/vs_1/files").mock( + return_value=httpx.Response(200, json=_ATTACHED_FILE) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("chunking_strategy", "sent_chunking_strategy"), + [(None, {"type": "auto"}), ({"type": "auto"}, {"type": "auto"}), (_STATIC_CHUNKING, _STATIC_CHUNKING)], +) +async def test_an_uploaded_file_is_attached_to_the_vector_store_with_the_chunking_strategy( + attach_file: respx.Route, respx_mock: respx.MockRouter, chunking_strategy: object, sent_chunking_strategy: object +): + respx_mock.post(f"{_API_BASE}/files").mock(return_value=httpx.Response(200, json=_UPLOADED_FILE)) + + stored: Final = await _ingestion(chunking_strategy).store( + file_content=b"first note", + filename="notes.txt", + content_type="text/plain", + chunks=[], + embeddings=None, + ) + + assert stored == ("vs_1", "file-1") + assert json.loads(attach_file.calls.last.request.content) == { + "file_id": "file-1", + "chunking_strategy": sent_chunking_strategy, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("chunking_strategy", "sent_chunking_strategy"), + [(None, {"type": "auto"}), (_STATIC_CHUNKING, _STATIC_CHUNKING)], +) +async def test_an_existing_file_is_attached_to_the_vector_store_with_the_chunking_strategy( + attach_file: respx.Route, chunking_strategy: object, sent_chunking_strategy: object +): + stored: Final = await _ingestion(chunking_strategy).store( + file_content=None, + filename=None, + content_type=None, + chunks=[], + embeddings=None, + existing_file_id="file-9", + ) + + assert stored == ("vs_1", "file-9") + assert json.loads(attach_file.calls.last.request.content) == { + "file_id": "file-9", + "chunking_strategy": sent_chunking_strategy, + } diff --git a/tests/unit/rag/test_main.py b/tests/unit/rag/test_main.py index 54748efb480..81318a1e113 100644 --- a/tests/unit/rag/test_main.py +++ b/tests/unit/rag/test_main.py @@ -12,12 +12,14 @@ aquery carries the completion response with real usage and cost. import asyncio import json +from types import MappingProxyType from typing import Final from unittest.mock import patch import httpx import pytest import respx +from pydantic import ValidationError import litellm from litellm._internal_context import is_internal_call @@ -538,6 +540,91 @@ async def test_aquery_forwards_vector_store_params_to_search_but_not_completion( assert not ({"milvus_text_field", "outputFields"} & set(completion_kwargs)) +_UNDECODABLE_FILE: Final = {"filename": "notes.txt", "content": "x"} + + +@pytest.mark.parametrize( + ("ingest_options", "expected_provider"), + [ + ({"vector_store": {"custom_llm_provider": "bedrock"}}, "bedrock"), + ({"vector_store": MappingProxyType({"custom_llm_provider": "bedrock"})}, "bedrock"), + ({"vector_store": {"vector_store_id": "vs_1", 7: "ignored"}}, None), + ({}, None), + ], +) +def test_ingest_failure_is_attributed_to_the_vector_store_provider( + ingest_options: dict[str, object], expected_provider: str | None +) -> None: + with pytest.raises(litellm.APIConnectionError, match="Invalid base64-encoded string") as raised: + litellm.ingest(ingest_options=ingest_options, file=_UNDECODABLE_FILE) + + assert raised.value.llm_provider == expected_provider + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("ingest_options", "expected_provider"), + [ + ({"vector_store": {"custom_llm_provider": "bedrock"}}, "bedrock"), + ({"vector_store": MappingProxyType({"custom_llm_provider": "bedrock"})}, "bedrock"), + ({"vector_store": {"vector_store_id": "vs_1", 7: "ignored"}}, None), + ({}, None), + ], +) +async def test_aingest_failure_is_attributed_to_the_vector_store_provider( + ingest_options: dict[str, object], expected_provider: str | None +) -> None: + with pytest.raises(litellm.APIConnectionError, match="Invalid base64-encoded string") as raised: + await litellm.aingest(ingest_options=ingest_options, file=_UNDECODABLE_FILE) + + assert raised.value.llm_provider == expected_provider + + +@pytest.mark.parametrize("vector_store", [None, "openai", ["openai"], [{"api_key": "sk-test"}]]) +def test_ingest_failure_with_a_vector_store_that_is_not_a_mapping_raises_a_validation_error( + vector_store: object, +) -> None: + with pytest.raises(ValidationError) as raised: + litellm.ingest(ingest_options={"vector_store": vector_store}, file=_UNDECODABLE_FILE) + + assert "sk-test" not in str(raised.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("vector_store", [None, "openai", ["openai"], [{"api_key": "sk-test"}]]) +async def test_aingest_failure_with_a_vector_store_that_is_not_a_mapping_raises_a_validation_error( + vector_store: object, +) -> None: + with pytest.raises(ValidationError) as raised: + await litellm.aingest(ingest_options={"vector_store": vector_store}, file=_UNDECODABLE_FILE) + + assert "sk-test" not in str(raised.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", [None, 5]) +async def test_aquery_with_a_provider_that_is_not_a_string_bills_only_the_completion(provider: object) -> None: + messages: Final = [{"role": "user", "content": "hello"}] + default_provider_response: Final = await litellm.aquery( + model="gpt-4o-mini", + messages=messages, + retrieval_config={"vector_store_id": "vs_test_123"}, + mock_response="hi there", + ) + + response: Final = await litellm.aquery( + model="gpt-4o-mini", + messages=messages, + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": provider}, + mock_response="hi there", + ) + + await _drain_logging_worker() + + assert response._hidden_params["response_cost"] > 0 + assert response._hidden_params["response_cost"] == default_provider_response._hidden_params["response_cost"] + + def test_rag_call_types_are_registered(): """ query/aquery/ingest/aingest are @client-decorated entry points, so their diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index bae6db9ee88..3aed934f40c 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -11,11 +11,13 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from prisma import models as prisma_models from prisma.builder import QueryBuilder +from pydantic import ValidationError from litellm.models.base import DomainModel from litellm.models.budget import LiteLLM_BudgetTable from litellm.models.credentials import CredentialItem from litellm.models.team import LiteLLM_TeamTable +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.repositories.base_repository import BaseRepository from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE @@ -308,6 +310,29 @@ class TestBudgetRepository: assert budget.budget_id == "budget-1" +_ROW_TIMESTAMP: Final = datetime(2026, 1, 1, 12, 0, 0) + + +def _stored_proxy_model_row(*, litellm_params: str, model_info: str) -> prisma_models.LiteLLM_ProxyModelTable: + return prisma_models.LiteLLM_ProxyModelTable( + model_id="other-model", + model_name="gpt-4o", + litellm_params=litellm_params, + model_info=model_info, + blocked=False, + created_at=_ROW_TIMESTAMP, + created_by="admin", + updated_at=_ROW_TIMESTAMP, + updated_by="admin", + ) + + +def _proxy_model_client(rows: list[prisma_models.LiteLLM_ProxyModelTable]) -> SimpleNamespace: + return SimpleNamespace( + db=SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=AsyncMock(return_value=rows))) + ) + + class TestModelRepository: @pytest.fixture def repo(self): @@ -331,6 +356,56 @@ class TestModelRepository: ).build_query() assert 'where: { model_id: { not: "current-model" } }' in " ".join(query.split()) + @pytest.mark.parametrize("double_encoded", [False, True]) + @pytest.mark.asyncio + async def test_find_all_except_returns_stored_rows_with_decrypted_params( + self, monkeypatch: pytest.MonkeyPatch, double_encoded: bool + ) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt") + stored_params: Final = json.dumps( + {"model": "openai/gpt-4o", "api_key": encrypt_value_helper("sk-secret"), "rpm": 5, "tags": ["prod"]} + ) + stored_info: Final = json.dumps({"id": "other-model", "team_id": "team-1"}) + row: Final = _stored_proxy_model_row( + litellm_params=json.dumps(stored_params) if double_encoded else stored_params, + model_info=json.dumps(stored_info) if double_encoded else stored_info, + ) + + models: Final = await ModelRepository(_proxy_model_client([row])).find_all_except("current-model") + + assert [model.model_dump() for model in models] == [ + { + "model_id": "other-model", + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-secret", "rpm": 5, "tags": ["prod"]}, + "model_info": {"id": "other-model", "team_id": "team-1"}, + "blocked": False, + "created_at": _ROW_TIMESTAMP, + "created_by": "admin", + "updated_at": _ROW_TIMESTAMP, + "updated_by": "admin", + } + ] + + @pytest.mark.asyncio + async def test_find_all_except_keeps_empty_params_and_missing_model_info(self) -> None: + row: Final = _stored_proxy_model_row(litellm_params="{}", model_info="null") + + models: Final = await ModelRepository(_proxy_model_client([row])).find_all_except("current-model") + + assert [(model.litellm_params, model.model_info) for model in models] == [({}, None)] + + @pytest.mark.parametrize("stored_params", ['["sk-secret"]', "1", "true", json.dumps('["sk-secret"]'), '"1"']) + @pytest.mark.asyncio + async def test_find_all_except_rejects_rows_whose_params_are_not_a_json_object(self, stored_params: str) -> None: + row: Final = _stored_proxy_model_row(litellm_params=stored_params, model_info="null") + + with pytest.raises(ValidationError) as rejected: + await ModelRepository(_proxy_model_client([row])).find_all_except("current-model") + + assert [error["type"] for error in rejected.value.errors()] == ["dict_type"] + assert "sk-secret" not in str(rejected.value) + def test_table_is_wrapped_for_config_sync(self, repo): from litellm.proxy.common_utils.config_sync_pubsub import ( _PublishOnWriteActions, diff --git a/tests/unit/router_strategy/adaptive_router/test_adaptive_router.py b/tests/unit/router_strategy/adaptive_router/test_adaptive_router.py index d717c4e8c89..9144bbc1788 100644 --- a/tests/unit/router_strategy/adaptive_router/test_adaptive_router.py +++ b/tests/unit/router_strategy/adaptive_router/test_adaptive_router.py @@ -7,6 +7,7 @@ from litellm.router_strategy.adaptive_router import adaptive_router as ar_module import pytest from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter +from litellm.router_strategy.adaptive_router.bandit import initial_cell from litellm.router_strategy.adaptive_router.signals import Turn from litellm.types.router import ( AdaptiveRouterConfig, @@ -416,3 +417,11 @@ def test_session_state_expiry_is_refreshed_on_access(): second_exp = r._session_states_expiry[("sess-A", "fast")] assert second_exp > first_exp + + +@pytest.mark.parametrize("model", ["fast", "smart"]) +def test_cell_returns_the_cold_start_prior_for_an_available_model(model): + r = _make_router() + cell = r.cell(RequestType.CODE_GENERATION, model) + assert cell == initial_cell(r.model_to_prefs[model], RequestType.CODE_GENERATION) + assert cell.total_samples == 0 diff --git a/tests/unit/secret_managers/test_aws_secret_manager.py b/tests/unit/secret_managers/test_aws_secret_manager.py new file mode 100644 index 00000000000..d9922b43ec8 --- /dev/null +++ b/tests/unit/secret_managers/test_aws_secret_manager.py @@ -0,0 +1,36 @@ +import base64 + +import pytest +from botocore.stub import Stubber + +from litellm.secret_managers.aws_secret_manager import AWSKeyManagementService_V2 + +_CIPHERTEXT = b"ciphertext-bytes" +_KEY_ID = "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab" + + +@pytest.mark.parametrize( + ("plaintext", "expected"), + [ + pytest.param(b"True", True, id="true literal"), + pytest.param(b"False", False, id="false literal"), + pytest.param(b" True\n", True, id="padded true literal"), + pytest.param(b"1", "1", id="integer literal"), + pytest.param(b"[True]", "[True]", id="list literal"), + pytest.param(b"'True'", "'True'", id="quoted string literal"), + pytest.param(b"sk-not-a-literal", "sk-not-a-literal", id="plain secret"), + pytest.param(b"", "", id="empty secret"), + ], +) +def test_decrypt_value_turns_only_boolean_literals_into_bools(monkeypatch, plaintext, expected): + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + monkeypatch.setenv("LITELLM_LICENSE", "license-for-test") + monkeypatch.setenv("ENCRYPTED_SETTING", "aws_kms/" + base64.b64encode(_CIPHERTEXT).decode()) + kms = AWSKeyManagementService_V2() + + with Stubber(kms.kms_client) as stubber: + stubber.add_response("decrypt", {"KeyId": _KEY_ID, "Plaintext": plaintext}, {"CiphertextBlob": _CIPHERTEXT}) + decrypted = kms.decrypt_value(secret_name="ENCRYPTED_SETTING") + + assert decrypted == expected + assert type(decrypted) is type(expected) diff --git a/tests/unit/secret_managers/test_secret_managers_main.py b/tests/unit/secret_managers/test_secret_managers_main.py index acc91b691d4..32040251795 100644 --- a/tests/unit/secret_managers/test_secret_managers_main.py +++ b/tests/unit/secret_managers/test_secret_managers_main.py @@ -431,3 +431,38 @@ def test_secret_manager_would_be_consulted_is_false_without_a_client(monkeypatch monkeypatch.setattr(litellm, "secret_manager_client", None) assert secret_manager_would_be_consulted("os.environ/ANY_NAME") is False + + +class _FixedValueSecretManager(CustomSecretManager): + def __init__(self, value): + self.value = value + + def sync_read_secret(self, secret_name, optional_params=None, timeout=None): + return self.value + + async def async_read_secret(self, secret_name, optional_params=None, timeout=None): + return self.value + + +@pytest.mark.parametrize( + ("stored", "expected"), + [ + ("True", True), + ("False", False), + ("1", "1"), + ("[True]", "[True]"), + ("'True'", "'True'"), + ("sk-not-a-literal", "sk-not-a-literal"), + ("", ""), + ], +) +def test_get_secret_turns_only_boolean_literals_from_the_secret_manager_into_bools(monkeypatch, stored, expected): + monkeypatch.setattr(litellm, "secret_manager_client", _FixedValueSecretManager(stored)) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + monkeypatch.delenv("STORED_FLAG", raising=False) + + secret = get_secret("os.environ/STORED_FLAG") + + assert secret == expected + assert type(secret) is type(expected) diff --git a/tests/unit/test_dashscope_image_generation.py b/tests/unit/test_dashscope_image_generation.py index 6f91fe9a0e0..960807a033b 100644 --- a/tests/unit/test_dashscope_image_generation.py +++ b/tests/unit/test_dashscope_image_generation.py @@ -9,6 +9,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest +from pydantic import ValidationError import litellm from litellm.llms.dashscope.image_generation.transformation import ( @@ -455,3 +456,98 @@ def test_litellm_image_generation_dashscope_end_to_end(model: str): assert "input" in body assert "messages" in body["input"] assert body["parameters"]["size"] == "1024*1024" + + +def _transform_response(payload: object) -> ImageResponse: + return DashScopeImageGenerationConfig().transform_image_generation_response( + model="qwen-image-2.0", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"output": {}}, + {"output": {"choices": []}}, + {"output": {"choices": ""}}, + {"output": {"choices": [{}]}}, + {"output": {"choices": [{"message": {}}]}}, + {"output": {"choices": [{"message": {"content": {}}}]}}, + {"output": {"choices": [{"message": {"content": [{}]}}]}}, + {"output": {"choices": [{"message": {"content": [{"image": ""}]}}]}}, + {"output": {"choices": [{"message": {"content": [{"text": "hi"}]}}]}}, + {"code": "Partial", "output": {"choices": []}}, + ], +) +def test_transform_response_without_image_content_has_no_images( + payload: dict[str, object], +): + assert _transform_response(payload).data == [] + + +def test_transform_response_skips_content_items_without_an_image(): + response = _transform_response( + { + "output": { + "choices": [ + {"message": {"content": [{"text": "caption"}, {"image": "https://a.example/1.png"}]}}, + {"message": {"content": [{"image": None}, {"image": "https://a.example/2.png"}]}}, + ] + } + } + ) + + assert [image.url for image in response.data] == [ + "https://a.example/1.png", + "https://a.example/2.png", + ] + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ({"code": "InvalidParameter", "message": "Size not supported"}, "Size not supported"), + ({"code": "Throttled"}, "{'code': 'Throttled'}"), + ({"code": 429, "message": {"detail": "slow down"}}, "{'detail': 'slow down'}"), + ], +) +def test_transform_response_reports_api_error_bodies( + payload: dict[str, object], message: str +): + with pytest.raises(BaseLLMException) as exc_info: + _transform_response(payload) + + assert exc_info.value.message == message + assert exc_info.value.status_code == 200 + + +@pytest.mark.parametrize( + "payload", + [ + ["code"], + "code", + 7, + {"output": None}, + {"output": ["not", "an", "object"]}, + {"output": {"choices": 7}}, + {"output": {"choices": ["not an object"]}}, + {"output": {"choices": [{"message": "not an object"}]}}, + {"output": {"choices": [{"message": {"content": None}}]}}, + {"output": {"choices": [{"message": {"content": ["not an object"]}}]}}, + ], +) +def test_transform_response_rejects_malformed_payloads_without_echoing_them( + payload: object, +): + with pytest.raises(ValidationError) as exc_info: + _transform_response(payload) + + assert "input_value" not in str(exc_info.value) diff --git a/tests/unit/vector_stores/test_vector_store_registry.py b/tests/unit/vector_stores/test_vector_store_registry.py index 762176d6a81..acfaccf8e2d 100644 --- a/tests/unit/vector_stores/test_vector_store_registry.py +++ b/tests/unit/vector_stores/test_vector_store_registry.py @@ -1,10 +1,14 @@ import json +from collections.abc import Mapping, Sequence +from types import SimpleNamespace +from typing import Final from unittest.mock import patch import httpx import pytest import respx from fastapi.testclient import TestClient +from pydantic import BaseModel, ValidationError from datetime import datetime, timezone @@ -13,7 +17,7 @@ from unittest.mock import AsyncMock, MagicMock import litellm from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.vector_stores.main import search -from litellm.vector_stores.vector_store_registry import VectorStoreRegistry +from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry @pytest.fixture(autouse=True) @@ -249,3 +253,93 @@ async def test_config_owned_store_survives_db_liveness_check_while_missing_db_st prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_awaited_once_with( where={"vector_store_id": "vs_from_db"} ) + + +_INDEX_PARAMS: Final = {"vector_store_index": "real-index-name", "vector_store_name": "azure-ai-search-store"} +_INDEX_CREATED_AT: Final = datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc) +_INDEX_ROW: Final = { + "id": "idx-1", + "index_name": "team-docs", + "litellm_params": _INDEX_PARAMS, + "index_info": {"dimensions": 1536}, + "created_at": _INDEX_CREATED_AT, + "created_by": "user-1", + "updated_at": _INDEX_CREATED_AT, + "updated_by": "user-2", +} + + +class GeneratedIndexRow(BaseModel): + id: str + index_name: str + litellm_params: dict[str, str] + index_info: dict[str, int] | None = None + created_at: datetime + created_by: str | None = None + updated_at: datetime + updated_by: str | None = None + + +def _database_listing_index_rows(rows: Sequence[object]) -> SimpleNamespace: + async def find_many(order: Mapping[str, str]) -> Sequence[object]: + return rows if order == {"created_at": "desc"} else [] + + return SimpleNamespace( + db=SimpleNamespace(litellm_managedvectorstoreindextable=SimpleNamespace(find_many=find_many)) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "row", + [ + _INDEX_ROW, + GeneratedIndexRow(**_INDEX_ROW), + list(_INDEX_ROW.items()), + {**_INDEX_ROW, "column_added_by_a_later_migration": "ignored"}, + ], +) +async def test_vector_store_index_rows_from_the_db_are_returned_as_indexes(row: object) -> None: + indexes: Final = await VectorStoreIndexRegistry._get_vector_store_indexes_from_db( + _database_listing_index_rows([row]) + ) + + assert [index.model_dump() for index in indexes] == [_INDEX_ROW] + + +@pytest.mark.asyncio +async def test_vector_store_index_row_without_optional_columns_gets_empty_defaults() -> None: + indexes: Final = await VectorStoreIndexRegistry._get_vector_store_indexes_from_db( + _database_listing_index_rows([{"id": "idx-1", "index_name": "team-docs", "litellm_params": _INDEX_PARAMS}]) + ) + + assert [index.model_dump() for index in indexes] == [ + { + "id": "idx-1", + "index_name": "team-docs", + "litellm_params": _INDEX_PARAMS, + "index_info": None, + "created_at": None, + "created_by": None, + "updated_at": None, + "updated_by": None, + } + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "row", + [ + {"index_name": "team-docs", "litellm_params": _INDEX_PARAMS}, + {**_INDEX_ROW, "id": 7}, + {**_INDEX_ROW, "litellm_params": None}, + {**_INDEX_ROW, "litellm_params": {"vector_store_index": "real-index-name"}}, + {**_INDEX_ROW, "index_info": ["not", "a", "mapping"]}, + {**_INDEX_ROW, 7: "column names must be strings"}, + {**_INDEX_ROW, b"index_name": "column names are not decoded"}, + ], +) +async def test_malformed_vector_store_index_row_from_the_db_raises_a_validation_error(row: object) -> None: + with pytest.raises(ValidationError): + await VectorStoreIndexRegistry._get_vector_store_indexes_from_db(_database_listing_index_rows([row])) From a99bccacead30b02acab705ead77f74d1e9e6109 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 5 Oct 2026 11:10:10 -0700 Subject: [PATCH 007/191] fix(ui): restore inline Lens onboarding and responsive layout (#44604) * fix(ui): improve gateway layouts on mobile * fix(ui): limit mobile cleanup to navigation and header * fix(ui): restore inline Lens onboarding and responsive layout * fix(ui): smooth Lens tab and panel connections * fix(ui): keep trace time controls within narrow panels --- .../src/components/lens/LensModeSwitch.tsx | 10 +- .../lens/LensPage.integration.test.tsx | 3 +- .../lens/LensSetup.integration.test.tsx | 109 +++++------ .../lens/LensWorkspace.integration.test.tsx | 21 ++- .../src/components/lens/LensWorkspace.tsx | 175 +++++++++++------- .../components/lens/hooks/useLensReadiness.ts | 5 +- .../lens/investigations/InvestigationList.tsx | 1 + .../InvestigationsView.integration.test.tsx | 41 +--- .../investigations/InvestigationsView.tsx | 10 +- .../components/lens/model/readiness.test.ts | 22 ++- .../src/components/lens/model/readiness.ts | 8 + .../lens/onboarding/LensGettingStarted.tsx | 9 +- .../lens/onboarding/LensIntroDialog.tsx | 80 -------- .../lens/onboarding/LensIntroduction.tsx | 18 +- .../lens/onboarding/OnboardingSetup.tsx | 5 +- .../lens/onboarding/OnboardingSteps.tsx | 30 ++- .../TracingSetupCard.integration.test.tsx | 30 ++- .../onboarding/tracing/TracingSetupCard.tsx | 2 - .../components/lens/settings/LensSettings.tsx | 22 +-- .../src/components/lens/storage.ts | 5 - .../lens/traces/detail/run/RunHeader.tsx | 24 +-- .../AgentTracesSection.integration.test.tsx | 23 +-- .../lens/traces/list/AgentTracesTable.tsx | 1 + .../lens/traces/list/runSearch/RunSearch.tsx | 4 +- .../traces/list/runSearch/RunsToolbar.tsx | 12 +- .../components/lens/ui/LensPreviewButton.tsx | 24 --- .../src/components/lens/ui/frame.ts | 20 +- .../components/shared/search/SearchBox.tsx | 4 +- .../shared/timeline/TimeRangeControls.tsx | 11 +- .../shared/timeline/Timeline.test.ts | 16 +- .../components/shared/timeline/Timeline.tsx | 12 +- .../tests/lens-test-utils.tsx | 7 - 32 files changed, 345 insertions(+), 419 deletions(-) delete mode 100644 ui/litellm-dashboard/src/components/lens/onboarding/LensIntroDialog.tsx delete mode 100644 ui/litellm-dashboard/src/components/lens/ui/LensPreviewButton.tsx diff --git a/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx b/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx index 0cfa9e47bce..c68a1112c8b 100644 --- a/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx +++ b/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx @@ -30,13 +30,12 @@ function ActivityDot({ activity }: { activity: InvestigationActivity }) { ); } -/** Inverted corner joining the tab's side border to the card's top border; plain CSS borders so both snap to the same pixels. */ function NotchCorner({ side, demo }: { side: "left" | "right"; demo: boolean }) { return (