From f8b31844bb8eb34fbee60481cf7bdc316be09fc5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 15:13:38 -0700 Subject: [PATCH] fix(auth): import PyJWT and cryptography at use time so a base litellm install still imports --- litellm/llms/base_llm/auth/jwt_signing.py | 30 +++++++++----- .../llms/base_llm/auth/test_jwt_signing.py | 39 +++++++++++++++++++ 2 files changed, 60 insertions(+), 9 deletions(-) diff --git a/litellm/llms/base_llm/auth/jwt_signing.py b/litellm/llms/base_llm/auth/jwt_signing.py index a7e7f8272f9..71aba6fac52 100644 --- a/litellm/llms/base_llm/auth/jwt_signing.py +++ b/litellm/llms/base_llm/auth/jwt_signing.py @@ -14,13 +14,16 @@ import hashlib import json from collections.abc import Mapping from types import MappingProxyType -from typing import Final, TypeAlias +from typing import TYPE_CHECKING, Final, TypeAlias -import jwt -from cryptography.hazmat.primitives.asymmetric import ec -from cryptography.hazmat.primitives.serialization import load_pem_private_key +if TYPE_CHECKING: + from cryptography.hazmat.primitives.asymmetric import ec ALG: Final = "ES256" +MISSING_SIGNING_DEPENDENCIES_MESSAGE: Final = ( + "the internal_issuer identity source needs PyJWT and cryptography, which a base litellm install " + "does not include: pip install 'litellm[proxy]'" +) _JWK_CURVE_NAME: Final = "P-256" _JWK_KEY_TYPE: Final = "EC" _COORDINATE_BYTE_LENGTH: Final = 32 # P-256 field element width, RFC 7518 6.2.1.2/6.2.1.3 @@ -29,8 +32,13 @@ Jwk: TypeAlias = Mapping[str, str] Jwks: TypeAlias = Mapping[str, tuple[Jwk, ...]] -def load_es256_private_key(pem: str) -> ec.EllipticCurvePrivateKey: +def load_es256_private_key(pem: str) -> "ec.EllipticCurvePrivateKey": """Parses an unencrypted PEM EC private key. Never echoes the key material in an error.""" + try: + from cryptography.hazmat.primitives.asymmetric import ec + from cryptography.hazmat.primitives.serialization import load_pem_private_key + except ImportError as e: + raise ImportError(MISSING_SIGNING_DEPENDENCIES_MESSAGE) from e try: key: Final = load_pem_private_key(pem.encode(), password=None) except (ValueError, TypeError) as e: @@ -46,7 +54,7 @@ def _b64url_coordinate(value: int) -> str: return base64.urlsafe_b64encode(value.to_bytes(_COORDINATE_BYTE_LENGTH, "big")).rstrip(b"=").decode("ascii") -def _jwk_thumbprint_members(public_key: ec.EllipticCurvePublicKey) -> Jwk: +def _jwk_thumbprint_members(public_key: "ec.EllipticCurvePublicKey") -> Jwk: """RFC 7638 3.2's exact EC member set (crv, kty, x, y) and nothing else: an extra member here would change the thumbprint and desync it from the ``kid`` published in the JWKS.""" numbers: Final = public_key.public_numbers() @@ -60,7 +68,7 @@ def _jwk_thumbprint_members(public_key: ec.EllipticCurvePublicKey) -> Jwk: ) -def rfc7638_thumbprint(public_key: ec.EllipticCurvePublicKey) -> str: +def rfc7638_thumbprint(public_key: "ec.EllipticCurvePublicKey") -> str: """RFC 7638: SHA-256 over the lexicographically member-ordered, whitespace-free JSON rendering of the thumbprint members, base64url-encoded without padding.""" canonical: Final = json.dumps( @@ -70,11 +78,11 @@ def rfc7638_thumbprint(public_key: ec.EllipticCurvePublicKey) -> str: return base64.urlsafe_b64encode(hashlib.sha256(canonical.encode()).digest()).rstrip(b"=").decode("ascii") -def build_jwk(public_key: ec.EllipticCurvePublicKey, kid: str) -> Jwk: +def build_jwk(public_key: "ec.EllipticCurvePublicKey", kid: str) -> Jwk: return MappingProxyType({**_jwk_thumbprint_members(public_key), "use": "sig", "alg": ALG, "kid": kid}) -def build_jwks(public_key: ec.EllipticCurvePublicKey) -> Jwks: +def build_jwks(public_key: "ec.EllipticCurvePublicKey") -> Jwks: kid: Final = rfc7638_thumbprint(public_key) return MappingProxyType({"keys": (build_jwk(public_key, kid),)}) @@ -97,6 +105,10 @@ def jwks_document_json(pem: str) -> str: def sign_es256_jwt(pem: str, claims: Mapping[str, object]) -> str: """Signs ``claims`` with the PEM key, stamping ``kid`` as its RFC 7638 thumbprint so a verifier can look the signing key up in the published JWKS by ``kid`` alone.""" + try: + import jwt + except ImportError as e: + raise ImportError(MISSING_SIGNING_DEPENDENCIES_MESSAGE) from e key: Final = load_es256_private_key(pem) kid: Final = rfc7638_thumbprint(key.public_key()) headers: Final = {"kid": kid} # mutable-ok: PyJWT requires a real dict, not a Mapping diff --git a/tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py b/tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py index 77c303ef4e2..30b400f4adb 100644 --- a/tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py +++ b/tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py @@ -1,6 +1,9 @@ import base64 import hashlib import json +import subprocess +import sys +import textwrap import time from typing import Final @@ -10,6 +13,7 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import ec, rsa from litellm.llms.base_llm.auth.jwt_signing import ( + MISSING_SIGNING_DEPENDENCIES_MESSAGE, build_jwk, build_jwks, jwks_document_json, @@ -172,3 +176,38 @@ class TestSignEs256Jwt: with pytest.raises(jwt.exceptions.InvalidSignatureError): jwt.decode(token, other_key.public_key(), algorithms=["ES256"]) + + +class TestBaseSdkImport: + """A base ``pip install litellm`` has neither PyJWT nor cryptography (both are proxy extras), + and ``litellm/__init__`` reaches this module through the Anthropic provider, so a + module-level import of either would break ``import litellm`` for every base SDK user.""" + + def test_module_imports_with_pyjwt_and_cryptography_absent(self): + script: Final = textwrap.dedent( + """ + import sys + + class Blocker: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] in {"jwt", "cryptography"}: + raise ModuleNotFoundError(f"No module named {name!r}") + + sys.meta_path.insert(0, Blocker()) + import litellm + from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json + try: + jwks_document_json("not a key") + except ImportError as e: + print(e) + """ + ) + result: Final = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr[-2000:] + assert result.stdout.strip() == MISSING_SIGNING_DEPENDENCIES_MESSAGE + + def test_signing_reports_the_missing_extra(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setitem(sys.modules, "jwt", None) + with pytest.raises(ImportError, match="litellm\\[proxy\\]"): + sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"}) +