From 4b3ce24c720b89b07ba5e0faf53565d0f8503675 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20K=C3=B6pf?= Date: Fri, 27 Jan 2023 21:28:57 +0100 Subject: [PATCH] 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 | 34 +++++++++-- .../oasst_backend/api/v1/frontend_users.py | 2 + backend/oasst_backend/api/v1/messages.py | 58 +++++++++++-------- backend/oasst_backend/api/v1/tasks.py | 6 +- backend/oasst_backend/api/v1/users.py | 12 +++- backend/oasst_backend/prompt_repository.py | 6 ++ backend/oasst_backend/user_repository.py | 4 +- 7 files changed, 85 insertions(+), 37 deletions(-) diff --git a/backend/oasst_backend/api/deps.py b/backend/oasst_backend/api/deps.py index 985b72a0..dd10bd20 100644 --- a/backend/oasst_backend/api/deps.py +++ b/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, diff --git a/backend/oasst_backend/api/v1/frontend_users.py b/backend/oasst_backend/api/v1/frontend_users.py index 86f78026..a4ca6380 100644 --- a/backend/oasst_backend/api/v1/frontend_users.py +++ b/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, ) diff --git a/backend/oasst_backend/api/v1/messages.py b/backend/oasst_backend/api/v1/messages.py index b3aace40..1fbaf53c 100644 --- a/backend/oasst_backend/api/v1/messages.py +++ b/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) diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py index 7875ac83..ae498bbb 100644 --- a/backend/oasst_backend/api/v1/tasks.py +++ b/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.") diff --git a/backend/oasst_backend/api/v1/users.py b/backend/oasst_backend/api/v1/users.py index 4c4e0118..dc0a3242 100644 --- a/backend/oasst_backend/api/v1/users.py +++ b/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) diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index d3f655fc..8e2c746d 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/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 diff --git a/backend/oasst_backend/user_repository.py b/backend/oasst_backend/user_repository.py index 28c042e0..6f96748f 100644 --- a/backend/oasst_backend/user_repository.py +++ b/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: