fix: replace jwt by pyjwt

docs/how_to_guides/telemetry.mdx

@@ -2,7 +2,11 @@ import Dashboard from './assets/dashboard.png';
 import ViewTraces from './assets/view_traces.mp4';
 import GrafanaOtelConfig from './assets/grafana_otel_config.mp4';
 
-# Capture Metrics for your Guards
+# Overview 
+
+In this document, we explain how to set up Guardrails with [OpenTelemetry (OTEL)](https://opentelemetry.io/) using either Grafana or a self-hosted OTEL collector. With this functionality enabled, you can measure latency of Guards, Large Language Models (LLMs), scuccess rates, and other metrics for your Guardrails-protected LLM calls. 
+
+## Metrics you can capture using OTEL
 
 This package is instrumented using the OpenTelemetry Python SDK. By viewing the captured traces and derived metrics, we're able to get useful insights into how our Guards, and our LLM apps in general perform. Among other things, we're able to find:
 
@@ -12,12 +16,11 @@ This package is instrumented using the OpenTelemetry Python SDK. By viewing the
 4. The rate at which validators Pass and Fail, within a Guard and across Guards
 5. Deep dives into singular guard and validator calls
 
-Since we are using OpenTelemetry, traces and metrics can be written to any OpenTelemetry enabled service or OTLP endpoint. This includes all major metrics providers like Grafana, New Relic, Prometheus, and Splunk.
+Since we are using OpenTelemetry, traces and metrics can be written to any OpenTelemetry-enabled service or OTLP endpoint. This includes all major metrics providers like Grafana, New Relic, Prometheus, and Splunk. 
 
-This guide will show how to set up your python project to log traces to Grafana and an OTEL collector.
+This guide will show how to set up your Python project to log traces to Grafana and to a self-hosted OTEL collector. For other OTEL endpoints, consult your metrics provider's documentation on OTEL support.
 
-
-## Setup with Grafana
+## Configure OTEL for Grafana
 
 Grafana Cloud is a free offering by Grafana that's easy to setup and is our preferred location for storing metrics.
 
@@ -41,11 +44,11 @@ OTEL_EXPORTER_OTLP_HEADERS
 
 ### Setup a Guard with Telemetry
 
-We first have to install the ```ValidLength``` guardrail from Guardrails Hub.
+1. First, install the ```ValidLength``` guardrail from Guardrails Hub.
 
 ```guardrails hub install hub://guardrails/valid_length```
 
-Then, set up your Guard the default tracer provided in the guardrails library. You can still use your desired validators
+2. Next, set up your Guard the default tracer provided in the guardrails library. You can still use your desired validators:
 
 <h5 a><strong><code>main.py</code></strong></h5>
 ```python
@@ -69,7 +72,7 @@ guard(
 )
 ```
 
-Before running the file, make sure to set the environment variables you got from Grafana
+3. Before running the file, make sure to set the environment variables you got from Grafana
 
 ```bash
 export GRAFANA_INSTANCE_ID=
@@ -82,34 +85,41 @@ export OTEL_EXPORTER_OTLP_ENDPOINT=
 export OTEL_EXPORTER_OTLP_HEADERS=
 ```
 
-Finally, run the python script
+4. Finally, run the python script
 
 ```bash
 
 python main.py
 
 ```
 
-### Viewing traces
+### View traces
+
+There are two ways to view traces: using the Explore tab or using the Guardrails Grafana dashboard template. 
 
-The simplest way to do this is to go to your grafana stack and click on the "Explore" tab. You should see a list of traces that you can filter by service name, operation name, and more.
+#### Use the Explore tab
+
+The simplest way to do this is to go to your grafana stack and click on the "**Explore** tab. You should see a list of traces that you can filter by service name, operation name, and more.
 
 <video autoPlay={true} style={{width: "100%"}} loop={true} controls={true}>
  <source src={ViewTraces} type="video/mp4"/>
 </video>
 
+#### Use the Guardrails Grafana dashboard template
+
 While this is easy to do, it's not the best way to get a big-picture view of how your guards are doing. For that, we should use the Guardrails Grafana dashboard template.
 
-The template can be found here https://grafana.com/grafana/dashboards/20600-standard-guardrails-dash/
-and instructions to use the template are found here https://grafana.com/docs/grafana/latest/dashboards/build-dashboards/import-dashboards/
+**[Use the template](https://grafana.com/grafana/dashboards/20600-standard-guardrails-dash/)**
+
+**[Template instructions](https://grafana.com/docs/grafana/latest/dashboards/build-dashboards/import-dashboards/)**
 
 <img src={Dashboard}/>
 
 
-## OTEL Collector
+## Configure OTEL for a self-hosted OpenTelemetry Collector
 
 For advanced use cases (like if you have a metrics provider in a VPC), you can use a self-hosted OpenTelemetry Collector to receive traces and metrics from your Guard.
-Standard [open telemetry environment variables](https://opentelemetry-python.readthedocs.io/en/stable/getting-started.html#configure-the-exporter) are used to configure the collector. Use the default_otel_collector_tracer when configuring your guard.
+Standard [open telemetry environment variables](https://opentelemetry.io/docs/languages/python/exporters/) are used to configure the collector. Use the `default_otel_collector_tracer` when configuring your guard.
 
 ```python
 from guardrails import Guard, OnFailAction

guardrails/cli/server/hub_client.py

@@ -4,14 +4,15 @@
 from typing import Any, Dict, Optional
 
 import requests
-from jwt import JWT
-from jwt.exceptions import JWTDecodeError
+import jwt
+from jwt import ExpiredSignatureError, DecodeError
+
 
 from guardrails.classes.credentials import Credentials
 from guardrails.cli.logger import logger
 from guardrails.cli.server.module_manifest import ModuleManifest
 
-FIND_NEW_TOKEN = "You can find a new token at https://hub.guardrailsai.com/tokens"
+FIND_NEW_TOKEN = "You can find a new token at https://hub.guardrailsai.com/keys"
 
 TOKEN_EXPIRED_MESSAGE = f"""Your token has expired. Please run `guardrails configure`\
 to update your token.
@@ -89,13 +90,11 @@ def get_jwt_token(creds: Credentials) -> Optional[str]:
     # check for jwt expiration
     if token:
         try:
-            JWT().decode(token, do_verify=False)
-        except JWTDecodeError as e:
-            # if the error message includes "Expired", then the token is expired
-            if "Expired" in str(e):
-                raise ExpiredTokenError(TOKEN_EXPIRED_MESSAGE)
-            else:
-                raise InvalidTokenError(TOKEN_INVALID_MESSAGE)
+            jwt.decode(token, options={"verify_signature": False, "verify_exp": True})
+        except ExpiredSignatureError:
+            raise ExpiredTokenError(TOKEN_EXPIRED_MESSAGE)
+        except DecodeError:
+            raise InvalidTokenError(TOKEN_INVALID_MESSAGE)
     return token
 
 

guardrails/hub_token/token.py

@@ -1,9 +1,9 @@
 from guardrails.classes.credentials import Credentials
-from jwt import JWT
-from jwt.exceptions import JWTDecodeError
+import jwt
+from jwt import ExpiredSignatureError, DecodeError
 from typing import Optional
 
-FIND_NEW_TOKEN = "You can find a new token at https://hub.guardrailsai.com/tokens"
+FIND_NEW_TOKEN = "You can find a new token at https://hub.guardrailsai.com/keys"
 
 TOKEN_EXPIRED_MESSAGE = f"""Your token has expired. Please run `guardrails configure`\
 to update your token.
@@ -39,11 +39,9 @@ def get_jwt_token(creds: Credentials) -> Optional[str]:
     # check for jwt expiration
     if token:
         try:
-            JWT().decode(token, do_verify=False)
-        except JWTDecodeError as e:
-            # if the error message includes "Expired", then the token is expired
-            if "Expired" in str(e):
-                raise ExpiredTokenError(TOKEN_EXPIRED_MESSAGE)
-            else:
-                raise InvalidTokenError(TOKEN_INVALID_MESSAGE)
+            jwt.decode(token, options={"verify_signature": False, "verify_exp": True})
+        except ExpiredSignatureError:
+            raise ExpiredTokenError(TOKEN_EXPIRED_MESSAGE)
+        except DecodeError:
+            raise InvalidTokenError(TOKEN_INVALID_MESSAGE)
     return token

pyproject.toml

@@ -51,8 +51,8 @@ faker = "^25.2.0"
 jsonref = "^1.1.0"
 jsonformer = {version = "0.12.0", optional = true}
 jsonschema = {version = "^4.22.0", extras = ["format"]}
-jwt = "^1.3.1"
 pip = ">=22"
+pyjwt = "^2.8.0"
 opentelemetry-sdk = "^1.24.0"
 opentelemetry-exporter-otlp-proto-grpc = "^1.24.0"
 opentelemetry-exporter-otlp-proto-http = "^1.24.0"

tests/unit_tests/cli/server/test_hub_client.py

@@ -1,13 +1,16 @@
 import datetime
-import math
 
 import pytest
-from jwt import JWT, jwk_from_dict
+import jwt
+from datetime import timezone
+
 
 from guardrails.classes.credentials import Credentials
 from guardrails.cli.server.hub_client import (
     TOKEN_EXPIRED_MESSAGE,
     TOKEN_INVALID_MESSAGE,
+    InvalidTokenError,
+    ExpiredTokenError,
     get_jwt_token,
 )
 
@@ -43,39 +46,27 @@ def test_get_auth():
 
 
 def test_get_jwt_token():
-    expiration = math.floor(datetime.datetime.now().timestamp() + 1000)
-
-    jwk = jwk_from_dict(
-        {
-            "alg": "HS256",
-            "kty": "oct",
-            "kid": "050bf691-4348-4891-940f-99af8354e82b",
-            "k": "eCE35cBrbRsO1GhrbxLXnGrVATgUFZDrPyyuOar4crw",
-        }
-    )
-
-    valid_jwt = JWT().encode(
-        {
-            "exp": expiration,
-        },
-        jwk,
-        "HS256",
-    )
-    creds = {"token": valid_jwt}
-    assert get_jwt_token(Credentials.from_dict(creds)) == valid_jwt
-
-    with pytest.raises(Exception) as e:
-        expiration = math.floor(datetime.datetime.now().timestamp() - 1000)
-        expired_jwt = JWT().encode(
-            {
-                "exp": expiration,
-            },
-            jwk,
-            "HS256",
-        )
+    # Create a JWT that expires in the future
+    secret_key = "secret"
+    timedelta = datetime.timedelta(seconds=1000)
+    expiration = datetime.datetime.now(tz=timezone.utc) + timedelta
+    valid_jwt = jwt.encode({"exp": expiration}, secret_key, algorithm="HS256")
+    creds = Credentials.from_dict({"token": valid_jwt})
+
+    # Test valid token
+    assert get_jwt_token(creds) == valid_jwt
+
+    # Test with an expired JWT
+    with pytest.raises(ExpiredTokenError) as e:
+        expired = datetime.datetime.now(tz=timezone.utc) - timedelta
+        expired_jwt = jwt.encode({"exp": expired}, secret_key, algorithm="HS256")
         get_jwt_token(Credentials.from_dict({"token": expired_jwt}))
+
     assert str(e.value) == TOKEN_EXPIRED_MESSAGE
 
-    with pytest.raises(Exception) as e:
-        get_jwt_token(Credentials.from_dict({"token": "invalid_token"}))
+    # Test with an invalid token format
+    with pytest.raises(InvalidTokenError) as e:
+        invalid_jwt = "invalid"
+        get_jwt_token(Credentials.from_dict({"token": invalid_jwt}))
+
     assert str(e.value) == TOKEN_INVALID_MESSAGE