Cancelling pending chat messages of user on submission of a new message (LAION-AI#2141)

Author: yk

inference/server/oasst_inference_server/chat_repository.py

@@ -6,6 +6,7 @@
 import sqlmodel
 from loguru import logger
 from oasst_inference_server import database, models
+from oasst_inference_server.schemas import chat as chat_schema
 from oasst_shared.schemas import inference
 
 
@@ -30,6 +31,9 @@ async def start_work(
         logger.debug(f"Starting work on message {message_id}")
         message = await self.get_assistant_message_by_id(message_id)
 
+        if message.state == inference.MessageState.cancelled:
+            raise chat_schema.MessageCancelledException(message_id=message_id)
+
         if message.state != inference.MessageState.pending:
             raise fastapi.HTTPException(status_code=400, detail="Message is not pending")
 

inference/server/oasst_inference_server/routes/workers.py

@@ -8,6 +8,7 @@
 from fastapi import Depends
 from loguru import logger
 from oasst_inference_server import chat_repository, database, deps, models, queueing, worker_utils
+from oasst_inference_server.schemas import chat as chat_schema
 from oasst_inference_server.settings import settings
 from oasst_shared.schemas import inference
 
@@ -148,6 +149,8 @@ def _add_receive(ftrs: set):
                         work_request_map[work_request.id] = WorkRequestContainer(
                             work_request=work_request, message_id=message_id
                         )
+                    except chat_schema.MessageCancelledException as e:
+                        logger.warning(f"Message was cancelled before work could be initiated: {e.message_id=}")
                     finally:
                         _add_dequeue(pending_futures)
                 else:

inference/server/oasst_inference_server/schemas/chat.py

@@ -70,3 +70,9 @@ class ChatRead(ChatListRead):
 
 class ListChatsResponse(pydantic.BaseModel):
     chats: list[ChatListRead]
+
+
+class MessageCancelledException(Exception):
+    def __init__(self, message_id: str):
+        super().__init__(f"Message {message_id} was cancelled")
+        self.message_id = message_id

inference/server/oasst_inference_server/user_chat_repository.py

@@ -59,6 +59,7 @@ async def create_chat(self) -> models.DbChat:
 
     async def add_prompter_message(self, chat_id: str, parent_id: str | None, content: str) -> models.DbMessage:
         logger.info(f"Adding prompter message {len(content)=} to chat {chat_id}")
+
         chat: models.DbChat = (
             await self.session.exec(
                 sqlmodel.select(models.DbChat)
@@ -110,6 +111,29 @@ async def initiate_assistant_message(
         self, parent_id: str, work_parameters: inference.WorkParameters
     ) -> models.DbMessage:
         logger.info(f"Adding stub assistant message to {parent_id=}")
+
+        # find and cancel all pending messages by this user
+        pending_msg_query = (
+            sqlmodel.select(models.DbMessage)
+            .where(
+                models.DbMessage.role == "assistant",
+                models.DbMessage.state == inference.MessageState.pending,
+            )
+            .join(models.DbChat)
+            .where(
+                models.DbChat.user_id == self.user_id,
+            )
+        )
+
+        pending_msgs: list[models.DbMessage] = (await self.session.exec(pending_msg_query)).all()
+        for pending_msg in pending_msgs:
+            logger.warning(
+                f"User {self.user_id} has a pending message {pending_msg.id} in chat {pending_msg.chat_id}. Cancelling..."
+            )
+            pending_msg.state = inference.MessageState.cancelled
+            await self.session.commit()
+            logger.debug(f"Cancelled message {pending_msg.id} in chat {pending_msg.chat_id}.")
+
         query = (
             sqlmodel.select(models.DbMessage)
             .options(sqlalchemy.orm.selectinload(models.DbMessage.chat))

oasst-shared/oasst_shared/schemas/inference.py

@@ -144,6 +144,7 @@ class MessageState(str, enum.Enum):
     in_progress = "in_progress"
     complete = "complete"
     aborted_by_worker = "aborted_by_worker"
+    cancelled = "cancelled"
 
 
 class MessageRead(pydantic.BaseModel):