[NEW] Solving merge conflicts

backend/oasst_backend/api/v1/frontend_messages.py

@@ -16,7 +16,7 @@ def get_message_by_frontend_id(
     """
     Get a message by its frontend ID.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message_by_frontend_message_id(message_id)
     return utils.prepare_message(message)
 
@@ -29,7 +29,7 @@ def get_conv_by_frontend_id(
     Get a conversation from the tree root and up to the message with given frontend ID.
     """
 
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message_by_frontend_message_id(message_id)
     messages = pr.fetch_message_conversation(message)
     return utils.prepare_conversation(messages)
@@ -43,7 +43,7 @@ def get_tree_by_frontend_id(
     Get all messages belonging to the same message tree.
     Message is identified by its frontend ID.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message_by_frontend_message_id(message_id)
     tree = pr.fetch_message_tree(message.message_tree_id)
     return utils.prepare_tree(tree, message.message_tree_id)
@@ -56,7 +56,7 @@ def get_children_by_frontend_id(
     """
     Get all messages belonging to the same message tree.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message_by_frontend_message_id(message_id)
     messages = pr.fetch_message_children(message.id)
     return utils.prepare_message_list(messages)
@@ -70,7 +70,7 @@ def get_descendants_by_frontend_id(
     Get a subtree which starts with this message.
     The message is identified by its frontend ID.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message_by_frontend_message_id(message_id)
     descendants = pr.fetch_message_descendants(message)
     return utils.prepare_tree(descendants, message.id)
@@ -84,7 +84,7 @@ def get_longest_conv_by_frontend_id(
     Get the longest conversation from the tree of the message.
     The message is identified by its frontend ID.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message_by_frontend_message_id(message_id)
     conv = pr.fetch_longest_conversation(message.message_tree_id)
     return utils.prepare_conversation(conv)
@@ -98,7 +98,7 @@ def get_max_children_by_frontend_id(
     Get message with the most children from the tree of the provided message.
     The message is identified by its frontend ID.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message_by_frontend_message_id(message_id)
     message, children = pr.fetch_message_with_max_children(message.message_tree_id)
     return utils.prepare_tree([message, *children], message.id)

backend/oasst_backend/api/v1/frontend_users.py

@@ -29,7 +29,7 @@ def query_frontend_user_messages(
     """
     Query frontend user messages.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     messages = pr.query_messages(
         username=username,
         api_client_id=api_client_id,
@@ -47,6 +47,6 @@ def query_frontend_user_messages(
 def mark_frontend_user_messages_deleted(
     username: str, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db)
 ):
-    pr = PromptRepository(db, api_client, None)
+    pr = PromptRepository(db, api_client)
     messages = pr.query_messages(username=username, api_client_id=api_client.id)
     pr.mark_messages_deleted(messages)

backend/oasst_backend/api/v1/leaderboards.py

@@ -1,7 +1,8 @@
 from fastapi import APIRouter, Depends
 from oasst_backend.api import deps
 from oasst_backend.models import ApiClient
-from oasst_backend.prompt_repository import PromptRepository
+from oasst_backend.user_repository import UserRepository
+from oasst_shared.schemas.protocol import LeaderboardStats
 from sqlmodel import Session
 
 router = APIRouter()
@@ -11,15 +12,15 @@ router = APIRouter()
 def get_assistant_leaderboard(
     db: Session = Depends(deps.get_db),
     api_client: ApiClient = Depends(deps.get_trusted_api_client),
-):
-    pr = PromptRepository(db, api_client, None)
-    return pr.get_user_leaderboard(role="assistant")
+) -> LeaderboardStats:
+    ur = UserRepository(db, api_client)
+    return ur.get_user_leaderboard(role="assistant")
 
 
 @router.get("/create/prompter")
 def get_prompter_leaderboard(
     db: Session = Depends(deps.get_db),
     api_client: ApiClient = Depends(deps.get_trusted_api_client),
-):
-    pr = PromptRepository(db, api_client, None)
-    return pr.get_user_leaderboard(role="prompter")
+) -> LeaderboardStats:
+    ur = UserRepository(db, api_client)
+    return ur.get_user_leaderboard(role="prompter")

backend/oasst_backend/api/v1/messages.py

@@ -29,7 +29,7 @@ def query_messages(
     """
     Query messages.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     messages = pr.query_messages(
         username=username,
         api_client_id=api_client_id,
@@ -51,7 +51,7 @@ def get_message(
     """
     Get a message by its internal ID.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message(message_id)
     return utils.prepare_message(message)
 
@@ -64,7 +64,7 @@ def get_conv(
     Get a conversation from the tree root and up to the message with given internal ID.
     """
 
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     messages = pr.fetch_message_conversation(message_id)
     return utils.prepare_conversation(messages)
 
@@ -76,7 +76,7 @@ def get_tree(
     """
     Get all messages belonging to the same message tree.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message(message_id)
     tree = pr.fetch_message_tree(message.message_tree_id)
     return utils.prepare_tree(tree, message.message_tree_id)
@@ -89,7 +89,7 @@ def get_children(
     """
     Get all messages belonging to the same message tree.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     messages = pr.fetch_message_children(message_id)
     return utils.prepare_message_list(messages)
 
@@ -101,7 +101,7 @@ def get_descendants(
     """
     Get a subtree which starts with this message.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message(message_id)
     descendants = pr.fetch_message_descendants(message)
     return utils.prepare_tree(descendants, message.id)
@@ -114,7 +114,7 @@ def get_longest_conv(
     """
     Get the longest conversation from the tree of the message.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     message = pr.fetch_message(message_id)
     conv = pr.fetch_longest_conversation(message.message_tree_id)
     return utils.prepare_conversation(conv)
@@ -127,7 +127,7 @@ def get_max_children(
     """
     Get message with the most children from the tree of the provided message.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     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)
@@ -137,5 +137,5 @@ def get_max_children(
 def mark_message_deleted(
     message_id: UUID, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db)
 ):
-    pr = PromptRepository(db, api_client, None)
+    pr = PromptRepository(db, api_client)
     pr.mark_messages_deleted(message_id)

backend/oasst_backend/api/v1/stats.py

@@ -13,5 +13,5 @@ def get_message_stats(
     db: Session = Depends(deps.get_db),
     api_client: ApiClient = Depends(deps.get_trusted_api_client),
 ):
-    pr = PromptRepository(db, api_client, None)
+    pr = PromptRepository(db, api_client)
     return pr.get_stats()

backend/oasst_backend/api/v1/tasks.py

@@ -8,7 +8,7 @@ from loguru import logger
 from oasst_backend.api import deps
 from oasst_backend.api.v1.utils import prepare_conversation
 from oasst_backend.config import settings
-from oasst_backend.prompt_repository import PromptRepository
+from oasst_backend.prompt_repository import PromptRepository, TaskRepository
 from oasst_backend.utils.hugging_face import HF_embeddingModel, HF_url, HuggingFaceAPI
 from oasst_shared.exceptions import OasstError, OasstErrorCode
 from oasst_shared.schemas import protocol as protocol_schema
@@ -192,9 +192,9 @@ def request_task(
     api_client = deps.api_auth(api_key, db)
 
     try:
-        pr = PromptRepository(db, api_client, request.user)
+        pr = PromptRepository(db, api_client, client_user=request.user)
         task, message_tree_id, parent_message_id = generate_task(request, pr)
-        pr.store_task(task, message_tree_id, parent_message_id, request.collective)
+        pr.task_repository.store_task(task, message_tree_id, parent_message_id, request.collective)
 
     except OasstError:
         raise
@@ -219,11 +219,11 @@ def tasks_acknowledge(
     api_client = deps.api_auth(api_key, db)
 
     try:
-        pr = PromptRepository(db, api_client, user=None)
+        pr = PromptRepository(db, api_client)
 
         # here we store the message id in the database for the task
         logger.info(f"Frontend acknowledges task {task_id=}, {ack_request=}.")
-        pr.bind_frontend_message_id(task_id=task_id, frontend_message_id=ack_request.message_id)
+        pr.task_repository.bind_frontend_message_id(task_id=task_id, frontend_message_id=ack_request.message_id)
 
     except OasstError:
         raise
@@ -247,8 +247,8 @@ 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, user=None)
-        pr.acknowledge_task_failure(task_id)
+        pr = PromptRepository(db, api_client)
+        pr.task_repository.acknowledge_task_failure(task_id)
     except (KeyError, RuntimeError):
         logger.exception("Failed to not acknowledge task.")
         raise OasstError("Failed to not acknowledge task.", OasstErrorCode.TASK_NACK_FAILED)
@@ -267,7 +267,7 @@ async def tasks_interaction(
     api_client = deps.api_auth(api_key, db)
 
     try:
-        pr = PromptRepository(db, api_client, user=interaction.user)
+        pr = PromptRepository(db, api_client, client_user=interaction.user)
 
         match type(interaction):
             case protocol_schema.TextReplyToMessage:
@@ -339,6 +339,6 @@ def close_collective_task(
     api_key: APIKey = Depends(deps.get_api_key),
 ):
     api_client = deps.api_auth(api_key, db)
-    pr = PromptRepository(db, api_client, user=None)
-    pr.close_task(close_task_request.message_id)
+    tr = TaskRepository(db, api_client)
+    tr.close_task(close_task_request.message_id)
     return protocol_schema.TaskDone()

backend/oasst_backend/api/v1/text_labels.py

@@ -25,7 +25,7 @@ def label_text(
 
     try:
         logger.info(f"Labeling text {text_labels=}.")
-        pr = PromptRepository(db, api_client, user=text_labels.user)
+        pr = PromptRepository(db, api_client, client_user=text_labels.user)
         pr.store_text_labels(text_labels)
 
     except Exception:

backend/oasst_backend/api/v1/users.py

@@ -29,7 +29,7 @@ def query_user_messages(
     """
     Query user messages.
     """
-    pr = PromptRepository(db, api_client, user=None)
+    pr = PromptRepository(db, api_client)
     messages = pr.query_messages(
         user_id=user_id,
         api_client_id=api_client_id,
@@ -48,6 +48,6 @@ def query_user_messages(
 def mark_user_messages_deleted(
     user_id: UUID, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db)
 ):
-    pr = PromptRepository(db, api_client, None)
+    pr = PromptRepository(db, api_client)
     messages = pr.query_messages(user_id=user_id)
     pr.mark_messages_deleted(messages)