use TokenStore interface

Arize-ai · Sep 20, 2024 · a31d6fe · a31d6fe
1 parent 80d3a8f
commit a31d6fe
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/src/phoenix/server/api/routers/auth.py b/src/phoenix/server/api/routers/auth.py
@@ -36,7 +36,6 @@
 from phoenix.server.bearer_auth import PhoenixUser, create_access_and_refresh_tokens
 from phoenix.server.email.templates.types import PasswordResetTemplateBody
 from phoenix.server.email.types import EmailSender
-from phoenix.server.jwt_store import JwtStore
 from phoenix.server.rate_limiters import ServerRateLimiter, fastapi_ip_rate_limiter
 from phoenix.server.types import (
     AccessTokenClaims,
@@ -72,7 +71,7 @@
 async def login(request: Request) -> Response:
     assert isinstance(access_token_expiry := request.app.state.access_token_expiry, timedelta)
     assert isinstance(refresh_token_expiry := request.app.state.refresh_token_expiry, timedelta)
-    token_store: JwtStore = request.app.state.get_token_store()
+    token_store: TokenStore = request.app.state.get_token_store()
     data = await request.json()
     email = data.get("email")
     password = data.get("password")

diff --git a/src/phoenix/server/api/routers/oauth2.py b/src/phoenix/server/api/routers/oauth2.py
@@ -29,13 +29,13 @@
 from phoenix.db import models
 from phoenix.db.enums import UserRole
 from phoenix.server.bearer_auth import create_access_and_refresh_tokens
-from phoenix.server.jwt_store import JwtStore
 from phoenix.server.oauth2 import OAuth2Client
 from phoenix.server.rate_limiters import (
     ServerRateLimiter,
     fastapi_ip_rate_limiter,
     fastapi_route_rate_limiter,
 )
+from phoenix.server.types import TokenStore
 
 _LOWERCASE_ALPHANUMS_AND_UNDERSCORES = r"[a-z0-9_]+"
 
@@ -116,7 +116,7 @@ async def create_tokens(
         return _redirect_to_login(error=_INVALID_OAUTH2_STATE_MESSAGE)
     assert isinstance(access_token_expiry := request.app.state.access_token_expiry, timedelta)
     assert isinstance(refresh_token_expiry := request.app.state.refresh_token_expiry, timedelta)
-    token_store: JwtStore = request.app.state.get_token_store()
+    token_store: TokenStore = request.app.state.get_token_store()
     if not isinstance(
         oauth2_client := request.app.state.oauth2_clients.get_client(idp_name), OAuth2Client
     ):