add parsing of additional x-oasst-user header field (#957)

* add parsing of additional x-oasst-user header field

* simplify passing FrontendUserId to PromptRepository, use frontend_user_id on more message enpoints

* remove unrelated change

backend/oasst_backend/api/deps.py

@@ -1,6 +1,6 @@
 from http import HTTPStatus
 from secrets import token_hex
-from typing import Generator
+from typing import Generator, NamedTuple
 
 from fastapi import Depends, Request, Response, Security
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
@@ -19,22 +19,46 @@ def get_db() -> Generator:
         yield db
 
 
-api_key_query = APIKeyQuery(name="api_key", auto_error=False)
-api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
+api_key_query = APIKeyQuery(name="api_key", scheme_name="api-key", auto_error=False)
+api_key_header = APIKeyHeader(name="X-API-Key", scheme_name="api-key", auto_error=False)
+oasst_user_query = APIKeyQuery(name="oasst_user", scheme_name="oasst-user", auto_error=False)
+oasst_user_header = APIKeyHeader(name="x-oasst-user", scheme_name="oasst-user", auto_error=False)
 
 bearer_token = HTTPBearer(auto_error=False)
 
 
-async def get_api_key(
+def get_api_key(
     api_key_query: str = Security(api_key_query),
     api_key_header: str = Security(api_key_header),
-):
+) -> str:
     if api_key_query:
         return api_key_query
     else:
         return api_key_header
 
 
+class FrontendUserId(NamedTuple):
+    auth_method: str
+    username: str
+
+
+def get_frontend_user_id(
+    user_query: str = Security(oasst_user_query),
+    user_header: str = Security(oasst_user_header),
+) -> FrontendUserId:
+    def split_user(v: str) -> tuple[str, str]:
+        if type(v) is str:
+            v = v.split(":", maxsplit=1)
+            if len(v) == 2:
+                return FrontendUserId(auth_method=v[0], username=v[1])
+        return FrontendUserId(auth_method=None, username=None)
+
+    if user_query:
+        return split_user(user_query)
+    else:
+        return split_user(user_header)
+
+
 def create_api_client(
     *,
     session: Session,

backend/oasst_backend/api/v1/frontend_users.py

@@ -104,6 +104,7 @@ def query_frontend_user_messages_cursor(
     max_count: Optional[int] = Query(10, gt=0, le=1000),
     desc: Optional[bool] = False,
     lang: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
@@ -117,6 +118,7 @@ def query_frontend_user_messages_cursor(
         max_count=max_count,
         desc=desc,
         lang=lang,
+        frontend_user=frontend_user,
         api_client=api_client,
         db=db,
     )

backend/oasst_backend/api/v1/messages.py

@@ -18,6 +18,7 @@ router = APIRouter()
 
 @router.get("/", response_model=list[protocol.Message])
 def query_messages(
+    *,
     auth_method: Optional[str] = None,
     username: Optional[str] = None,
     api_client_id: Optional[str] = None,
@@ -28,13 +29,14 @@ def query_messages(
     desc: Optional[bool] = True,
     allow_deleted: Optional[bool] = False,
     lang: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
     """
     Query messages.
     """
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username)
+    pr = PromptRepository(db, api_client, auth_method=frontend_user.auth_method, username=frontend_user.username)
     messages = pr.query_messages_ordered_by_created_date(
         auth_method=auth_method,
         username=username,
@@ -53,6 +55,7 @@ def query_messages(
 
 @router.get("/cursor", response_model=protocol.MessagePage)
 def get_messages_cursor(
+    *,
     before: Optional[str] = None,
     after: Optional[str] = None,
     user_id: Optional[UUID] = None,
@@ -64,6 +67,7 @@ def get_messages_cursor(
     max_count: Optional[int] = Query(10, gt=0, le=1000),
     desc: Optional[bool] = False,
     lang: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
@@ -93,7 +97,7 @@ def get_messages_cursor(
 
     qry_max_count = max_count + 1 if before is None or after is None else max_count
 
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username, user_id=user_id)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     items = pr.query_messages_ordered_by_created_date(
         user_id=user_id,
         auth_method=auth_method,
@@ -137,25 +141,25 @@ def get_messages_cursor(
 
 @router.get("/{message_id}", response_model=protocol.Message)
 def get_message(
+    *,
     message_id: UUID,
-    auth_method: Optional[str] = None,
-    username: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
     """
     Get a message by its internal ID.
     """
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     message = pr.fetch_message(message_id)
     return utils.prepare_message(message)
 
 
 @router.get("/{message_id}/conversation", response_model=protocol.Conversation)
 def get_conv(
+    *,
     message_id: UUID,
-    auth_method: Optional[str] = None,
-    username: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
@@ -163,23 +167,23 @@ def get_conv(
     Get a conversation from the tree root and up to the message with given internal ID.
     """
 
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     messages = pr.fetch_message_conversation(message_id)
     return utils.prepare_conversation(messages)
 
 
 @router.get("/{message_id}/tree", response_model=protocol.MessageTree)
 def get_tree(
+    *,
     message_id: UUID,
-    auth_method: Optional[str] = None,
-    username: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
     """
     Get all messages belonging to the same message tree.
     """
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     message = pr.fetch_message(message_id)
     tree = pr.fetch_message_tree(message.message_tree_id, reviewed=False)
     return utils.prepare_tree(tree, message.message_tree_id)
@@ -187,32 +191,32 @@ def get_tree(
 
 @router.get("/{message_id}/children", response_model=list[protocol.Message])
 def get_children(
+    *,
     message_id: UUID,
-    auth_method: Optional[str] = None,
-    username: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
     """
     Get all messages belonging to the same message tree.
     """
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     messages = pr.fetch_message_children(message_id)
     return utils.prepare_message_list(messages)
 
 
 @router.get("/{message_id}/descendants", response_model=protocol.MessageTree)
 def get_descendants(
+    *,
     message_id: UUID,
-    auth_method: Optional[str] = None,
-    username: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
     """
     Get a subtree which starts with this message.
     """
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     message = pr.fetch_message(message_id)
     descendants = pr.fetch_message_descendants(message)
     return utils.prepare_tree(descendants, message.id)
@@ -220,16 +224,16 @@ def get_descendants(
 
 @router.get("/{message_id}/longest_conversation_in_tree", response_model=protocol.Conversation)
 def get_longest_conv(
+    *,
     message_id: UUID,
-    auth_method: Optional[str] = None,
-    username: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
     """
     Get the longest conversation from the tree of the message.
     """
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     message = pr.fetch_message(message_id)
     conv = pr.fetch_longest_conversation(message.message_tree_id)
     return utils.prepare_conversation(conv)
@@ -237,16 +241,16 @@ def get_longest_conv(
 
 @router.get("/{message_id}/max_children_in_tree", response_model=protocol.MessageTree)
 def get_max_children(
+    *,
     message_id: UUID,
-    auth_method: Optional[str] = None,
-    username: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
     """
     Get message with the most children from the tree of the provided message.
     """
-    pr = PromptRepository(db, api_client, auth_method=auth_method, username=username)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     message = pr.fetch_message(message_id)
     message, children = pr.fetch_message_with_max_children(message.message_tree_id)
     return utils.prepare_tree([message, *children], message.id)
@@ -254,9 +258,13 @@ def get_max_children(
 
 @router.delete("/{message_id}", status_code=HTTP_204_NO_CONTENT)
 def mark_message_deleted(
-    message_id: UUID, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db)
+    *,
+    message_id: UUID,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
+    api_client: ApiClient = Depends(deps.get_trusted_api_client),
+    db: Session = Depends(deps.get_db),
 ):
-    pr = PromptRepository(db, api_client)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     pr.mark_messages_deleted(message_id)
 
 

backend/oasst_backend/api/v1/tasks.py

@@ -77,6 +77,7 @@ def tasks_acknowledge(
     *,
     db: Session = Depends(deps.get_db),
     api_key: APIKey = Depends(deps.get_api_key),
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     task_id: UUID,
     ack_request: protocol_schema.TaskAck,
 ) -> None:
@@ -87,7 +88,7 @@ def tasks_acknowledge(
     api_client = deps.api_auth(api_key, db)
 
     try:
-        pr = PromptRepository(db, api_client)
+        pr = PromptRepository(db, api_client, frontend_user=frontend_user)
 
         # here we store the message id in the database for the task
         logger.info(f"Frontend acknowledges task {task_id=}, {ack_request=}.")
@@ -105,6 +106,7 @@ def tasks_acknowledge_failure(
     *,
     db: Session = Depends(deps.get_db),
     api_key: APIKey = Depends(deps.get_api_key),
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     task_id: UUID,
     nack_request: protocol_schema.TaskNAck,
 ) -> None:
@@ -115,7 +117,7 @@ def tasks_acknowledge_failure(
     try:
         logger.info(f"Frontend reports failure to implement task {task_id=}, {nack_request=}.")
         api_client = deps.api_auth(api_key, db)
-        pr = PromptRepository(db, api_client)
+        pr = PromptRepository(db, api_client, frontend_user=frontend_user)
         pr.task_repository.acknowledge_task_failure(task_id)
     except (KeyError, RuntimeError):
         logger.exception("Failed to not acknowledge task.")

backend/oasst_backend/api/v1/users.py

@@ -225,13 +225,14 @@ def query_user_messages(
     desc: bool = True,
     include_deleted: bool = False,
     lang: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
     """
     Query user messages.
     """
-    pr = PromptRepository(db, api_client, user_id=user_id)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     messages = pr.query_messages_ordered_by_created_date(
         user_id=user_id,
         api_client_id=api_client_id,
@@ -257,6 +258,7 @@ def query_user_messages_cursor(
     max_count: Optional[int] = Query(10, gt=0, le=1000),
     desc: Optional[bool] = False,
     lang: Optional[str] = None,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
 ):
@@ -269,6 +271,7 @@ def query_user_messages_cursor(
         max_count=max_count,
         desc=desc,
         lang=lang,
+        frontend_user=frontend_user,
         api_client=api_client,
         db=db,
     )
@@ -276,9 +279,12 @@ def query_user_messages_cursor(
 
 @router.delete("/{user_id}/messages", status_code=HTTP_204_NO_CONTENT)
 def mark_user_messages_deleted(
-    user_id: UUID, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db)
+    user_id: UUID,
+    frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
+    api_client: ApiClient = Depends(deps.get_trusted_api_client),
+    db: Session = Depends(deps.get_db),
 ):
-    pr = PromptRepository(db, api_client)
+    pr = PromptRepository(db, api_client, frontend_user=frontend_user)
     messages = pr.query_messages_ordered_by_created_date(user_id=user_id, limit=None)
     pr.mark_messages_deleted(messages)
 

backend/oasst_backend/prompt_repository.py

@@ -8,6 +8,7 @@ from uuid import UUID, uuid4
 import oasst_backend.models.db_payload as db_payload
 import sqlalchemy as sa
 from loguru import logger
+from oasst_backend.api.deps import FrontendUserId
 from oasst_backend.journal_writer import JournalWriter
 from oasst_backend.models import (
     ApiClient,
@@ -48,10 +49,15 @@ class PromptRepository:
         user_id: Optional[UUID] = None,
         auth_method: Optional[str] = None,
         username: Optional[str] = None,
+        frontend_user: Optional[FrontendUserId] = None,
     ):
         self.db = db
         self.api_client = api_client
         self.user_repository = user_repository or UserRepository(db, api_client)
+
+        if frontend_user and not auth_method and not username:
+            auth_method, username = frontend_user
+
         if user_id:
             self.user = self.user_repository.get_user(id=user_id)
             self.user_id = self.user.id

backend/oasst_backend/user_repository.py

@@ -118,8 +118,6 @@ class UserRepository:
 
     @managed_tx_method(CommitMode.COMMIT)
     def _lookup_client_user_tx(self, client_user: protocol_schema.User, create_missing: bool = True) -> Optional[User]:
-        if not client_user:
-            return None
         user: User = (
             self.db.query(User)
             .filter(
@@ -146,6 +144,8 @@ class UserRepository:
         return user
 
     def lookup_client_user(self, client_user: protocol_schema.User, create_missing: bool = True) -> Optional[User]:
+        if not client_user:
+            return None
         num_retries = settings.DATABASE_MAX_TX_RETRY_COUNT
         for i in range(num_retries):
             try: