Fixes auth TokenPair change bug (also fix for null logprob from inference server) (LAION-AI#2155)

Author: yk

CODEOWNERS

@@ -10,3 +10,4 @@
 /safety @SummerSigh @shahules786
 /inference/ @yk @andreaskoepf @olliestanley @AbdBarho
 /backend/ @andreaskoepf @melvinebenezer @yk
+/oasst-shared/ @andreaskoepf @melvinebenezer @yk @olliestanley @AbdBarho

inference/server/oasst_inference_server/routes/auth.py

@@ -86,8 +86,8 @@ async def callback_discord(
     refresh_token = auth.create_refresh_token(user.id)
 
     token_pair = protocol.TokenPair(
-        protocol.Token(access_token=access_token, token_type="bearer"),
-        protocol.Token(access_token=refresh_token, token_type="refresh"),
+        access_token=protocol.Token(access_token=access_token, token_type="bearer"),
+        refresh_token=protocol.Token(access_token=refresh_token, token_type="refresh"),
     )
 
     return token_pair
@@ -154,8 +154,8 @@ async def callback_github(
     refresh_token = auth.create_refresh_token(user.id)
 
     token_pair = protocol.TokenPair(
-        protocol.Token(access_token=access_token, token_type="bearer"),
-        protocol.Token(access_token=refresh_token, token_type="refresh"),
+        access_token=protocol.Token(access_token=access_token, token_type="bearer"),
+        refresh_token=protocol.Token(access_token=refresh_token, token_type="refresh"),
     )
 
     return token_pair
@@ -212,8 +212,8 @@ async def login_debug(username: str, db: database.AsyncSession = Depends(deps.cr
     refresh_token = auth.create_refresh_token(user.id)
 
     token_pair = protocol.TokenPair(
-        protocol.Token(access_token=access_token, token_type="bearer"),
-        protocol.Token(access_token=refresh_token, token_type="refresh"),
+        access_token=protocol.Token(access_token=access_token, token_type="bearer"),
+        refresh_token=protocol.Token(access_token=refresh_token, token_type="refresh"),
     )
 
     return token_pair

inference/text-client/text_client_utils.py

@@ -12,8 +12,8 @@ def __init__(self, backend_url, http_client=requests):
 
     def login(self, username):
         auth_data = self.http_client.get(f"{self.backend_url}/auth/login/debug", params={"username": username}).json()
-        assert auth_data["token_type"] == "bearer"
-        bearer_token = auth_data["access_token"]
+        assert auth_data["access_token"]["token_type"] == "bearer"
+        bearer_token = auth_data["access_token"]["access_token"]
         logger.debug(f"Logged in as {username} with token {bearer_token}")
         self.auth_headers = {"Authorization": f"Bearer {bearer_token}"}
 

inference/worker/interface.py

@@ -41,7 +41,7 @@ class GenerateStreamRequest(pydantic.BaseModel):
 
 class Token(pydantic.BaseModel):
     text: str
-    logprob: float
+    logprob: float | None
     id: int
 
     def __len__(self) -> int:

oasst-shared/oasst_shared/schemas/inference.py

@@ -207,7 +207,7 @@ class PongResponse(WorkerResponseBase):
 class TokenResponse(WorkerResponseBase):
     response_type: Literal["token"] = "token"
     text: str
-    log_prob: float
+    log_prob: float | None
     token_id: int