Merge branch 'main' of github.com:LAION-AI/Open-Chat-GPT

LAION-AI · Jan 15, 2023 · f497bf1 · f497bf1
2 parents 5140928 + de28d67
commit f497bf1
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py
@@ -150,7 +150,7 @@ def store_text_reply(
         self._validate_task(task)
 
         # If there's no parent message assume user started new conversation
-        role = "prompter"
+        role = None
         depth = 0
 
         if task.parent_message_id:
@@ -170,10 +170,23 @@ def store_text_reply(
             self.db.add(parent_message)
 
             depth = parent_message.depth + 1
-            if parent_message.role == "assistant":
-                role = "prompter"
-            else:
-                role = "assistant"
+
+        task_payload: db_payload.TaskPayload = task.payload.payload
+        if isinstance(task_payload, db_payload.InitialPromptPayload):
+            role = "prompter"
+        elif isinstance(task_payload, db_payload.PrompterReplyPayload):
+            role = "prompter"
+        elif isinstance(task_payload, db_payload.AssistantReplyPayload):
+            role = "assistant"
+        elif isinstance(task_payload, db_payload.SummarizationStoryPayload):
+            raise NotImplementedError("SummarizationStory task not implemented.")
+        else:
+            raise OasstError(
+                f"Unexpected task payload type: {type(task_payload).__name__}",
+                OasstErrorCode.TASK_UNEXPECTED_PAYLOAD_TYPE_,
+            )
+
+        assert role in ("assistant", "prompter")
 
         # create reply message
         new_message_id = uuid4()

diff --git a/oasst-shared/oasst_shared/exceptions/oasst_api_error.py b/oasst-shared/oasst_shared/exceptions/oasst_api_error.py
@@ -59,6 +59,7 @@ class OasstErrorCode(IntEnum):
     TASK_ALREADY_DONE = 2105
     TASK_NOT_COLLECTIVE = 2106
     TASK_NOT_ASSIGNED_TO_USER = 2106
+    TASK_UNEXPECTED_PAYLOAD_TYPE_ = 2107
     USER_NOT_FOUND = 2200
 
     # 3000-4000: external resources