fix(auth): import PyJWT and cryptography at use time so a base litellm install still imports

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

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"})
+