Inference worker inform backend on safety intervention (LAION-AI#2505)

Author: olliestanley

docker/inference/Dockerfile.safety

@@ -55,7 +55,7 @@ COPY --chown="${APP_USER}:${APP_USER}" --from=build /build/lib
 COPY --chown="${APP_USER}:${APP_USER}"              ./${APP_RELATIVE_PATH}/*.py .
 
 
-CMD python3 __main__.py
+CMD python3 main.py
 
 FROM base-env as prod
 ARG APP_USER

inference/safety/requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/LAION-AI/blade2blade@8fd43bcbc5ff35fd59663c77ef08b3ec6c239dd4#egg=blade2blade
+blade2blade
 fastapi
 loguru
 pydantic

inference/server/alembic/versions/2023_04_29_1739-ea19bbc743f9_add_safe_content_to_message.py

@@ -0,0 +1,34 @@
+"""Add safe_content to message
+
+Revision ID: ea19bbc743f9
+Revises: 401eef162771
+Create Date: 2023-04-14 22:37:41.373382
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "ea19bbc743f9"
+down_revision = "401eef162771"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("message", sa.Column("safe_content", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
+    op.add_column("message", sa.Column("safety_level", sa.Integer(), nullable=True))
+    op.add_column("message", sa.Column("safety_label", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
+    op.add_column("message", sa.Column("safety_rots", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("message", "safe_content")
+    op.drop_column("message", "safety_level")
+    op.drop_column("message", "safety_label")
+    op.drop_column("message", "safety_rots")
+    # ### end Alembic commands ###

inference/server/oasst_inference_server/chat_repository.py

@@ -26,6 +26,15 @@ async def get_assistant_message_by_id(self, message_id: str) -> models.DbMessage
         message = (await self.session.exec(query)).one()
         return message
 
+    async def get_prompter_message_by_id(self, message_id: str) -> models.DbMessage:
+        query = (
+            sqlmodel.select(models.DbMessage)
+            .options(sqlalchemy.orm.selectinload(models.DbMessage.reports))
+            .where(models.DbMessage.id == message_id, models.DbMessage.role == "prompter")
+        )
+        message = (await self.session.exec(query)).one()
+        return message
+
     async def start_work(
         self, *, message_id: str, worker_id: str, worker_config: inference.WorkerConfig
     ) -> models.DbMessage:

inference/server/oasst_inference_server/models/chat.py

@@ -23,6 +23,11 @@ class DbMessage(SQLModel, table=True):
     content: str | None = Field(None)
     error: str | None = Field(None)
 
+    safe_content: str | None = Field(None)
+    safety_level: int | None = Field(None)
+    safety_label: str | None = Field(None)
+    safety_rots: str | None = Field(None)
+
     state: inference.MessageState = Field(inference.MessageState.manual)
     work_parameters: inference.WorkParameters = Field(None, sa_column=sa.Column(pg.JSONB))
     work_begin_at: datetime.datetime | None = Field(None)
@@ -59,6 +64,10 @@ def to_read(self) -> inference.MessageRead:
             score=self.score,
             work_parameters=self.work_parameters,
             reports=[r.to_read() for r in self.reports],
+            safe_content=self.safe_content,
+            safety_level=self.safety_level,
+            safety_label=self.safety_label,
+            safety_rots=self.safety_rots,
         )
 
 

inference/server/oasst_inference_server/routes/chats.py

@@ -234,6 +234,14 @@ async def event_generator(chat_id: str, message_id: str, worker_compat_hash: str
                     )
                     break
 
+                if response_packet.response_type == "safe_prompt":
+                    logger.info(f"Received safety intervention for {chat_id}")
+                    yield {
+                        "data": chat_schema.SafePromptResponseEvent(
+                            safe_prompt=response_packet.safe_prompt,
+                        ).json(),
+                    }
+
                 if response_packet.response_type == "internal_error":
                     yield {
                         "data": chat_schema.ErrorResponseEvent(

inference/server/oasst_inference_server/routes/workers.py

@@ -211,6 +211,13 @@ def _add_receive(ftrs: set):
                                     response=worker_response,
                                 )
                                 await _update_session(worker_response.metrics)
+                            case "safe_prompt":
+                                logger.info("Received safe prompt response")
+                                worker_response = cast(inference.SafePromptResponse, worker_response)
+                                await handle_safe_prompt_response(
+                                    response=worker_response,
+                                    work_request_map=work_request_map,
+                                )
                             case _:
                                 raise RuntimeError(f"Unknown response type: {worker_response.response_type}")
                     finally:
@@ -387,6 +394,27 @@ async def handle_general_error_response(
     logger.warning(f"Got general error {response=}")
 
 
+async def handle_safe_prompt_response(
+    response: inference.SafePromptResponse,
+    work_request_map: WorkRequestContainerMap,
+):
+    """
+    Handle the case where the worker informs the server that the safety model has intervened and modified the user prompt to be safe.
+    """
+    work_response_container = get_work_request_container(work_request_map, response.request_id)
+    message_id = work_response_container.message_id
+
+    async with deps.manual_create_session() as session:
+        cr = chat_repository.ChatRepository(session=session)
+        message = await cr.get_assistant_message_by_id(message_id)
+        prompt = await cr.get_prompter_message_by_id(message.parent_id)
+        prompt.safe_content = response.safe_prompt
+        prompt.safety_level = response.safety_parameters.level
+        prompt.safety_label = response.safety_label
+        prompt.safety_rots = response.safety_rots
+        await session.commit()
+
+
 async def handle_timeout(message: inference.MessageRead):
     response = inference.InternalErrorResponse(
         error="Timeout",

inference/server/oasst_inference_server/schemas/chat.py

@@ -38,8 +38,15 @@ class MessageResponseEvent(pydantic.BaseModel):
     message: inference.MessageRead
 
 
+class SafePromptResponseEvent(pydantic.BaseModel):
+    event_type: Literal["safe_prompt"] = "safe_prompt"
+    safe_prompt: str
+    message: inference.MessageRead
+
+
 ResponseEvent = Annotated[
-    Union[TokenResponseEvent, ErrorResponseEvent, MessageResponseEvent], pydantic.Field(discriminator="event_type")
+    Union[TokenResponseEvent, ErrorResponseEvent, MessageResponseEvent, SafePromptResponseEvent],
+    pydantic.Field(discriminator="event_type"),
 ]
 
 

inference/worker/work.py

@@ -80,24 +80,22 @@ def _prepare_message(message: inference.MessageRead) -> str:
     return prompt, parameters
 
 
-def prepare_safe_prompt(prompt: str, label: str, rots: str):
+def prepare_safe_prompt(prompt: str, label: str, rots: str) -> str:
     pre_prompt = f"Answer the following request with {label} as responsible chatbot that believes that {rots}: "
     input_list = prompt.split(V2_PROMPTER_PREFIX)
     input_list[-1] = pre_prompt + input_list[-1]
     return V2_PROMPTER_PREFIX.join(input_list)
 
 
-def get_safety_opinion(prompt: str, safety_opinion: str, safety_level: int):
+def is_safety_triggered(safety_label: str, safety_level: int) -> bool:
+    return ("caution" in safety_label and safety_level > 1) or ("intervention" in safety_label and safety_level > 0)
+
+
+def parse_safety_response(safety_opinion: str) -> tuple[str, str]:
     safety_opinion = re.sub(r"<pad>|</s>", "", safety_opinion).split("<sep>")
     label, rots = safety_opinion[0], "and".join([x.strip(".") for x in safety_opinion[1:]])
     label = label.replace("<pad>", "").strip()
-
-    if "caution" in label and safety_level > 1:
-        return prepare_safe_prompt(prompt, label, rots)
-    elif "intervention" in label and safety_level > 0:
-        return prepare_safe_prompt(prompt, label, rots)
-    else:
-        return prompt
+    return label, rots
 
 
 def handle_work_request(
@@ -115,8 +113,23 @@ def handle_work_request(
     if settings.enable_safety and work_request.safety_parameters.level:
         safety_request = inference.SafetyRequest(inputs=prompt, parameters=work_request.safety_parameters)
         safety_response = get_safety_server_response(safety_request)
-        prompt = get_safety_opinion(prompt, safety_response.outputs, work_request.safety_parameters.level)
-        logger.debug(f"Safe prompt: {prompt}")
+        safety_label, safety_rots = parse_safety_response(safety_response.outputs)
+
+        if is_safety_triggered(safety_label, work_request.safety_parameters.level):
+            prompt = prepare_safe_prompt(prompt, safety_label, safety_rots)
+
+            utils.send_response(
+                ws,
+                inference.SafePromptResponse(
+                    request_id=work_request.id,
+                    safe_prompt=prompt,
+                    safety_parameters=work_request.safety_parameters,
+                    safety_label=safety_label,
+                    safety_rots=safety_rots,
+                ),
+            )
+
+            logger.debug(f"Safe prompt: {prompt}")
 
     stream_response = None
     token_buffer = utils.TokenBuffer(stop_sequences=parameters.stop)

oasst-shared/oasst_shared/schemas/inference.py

@@ -169,6 +169,10 @@ class MessageRead(pydantic.BaseModel):
     reports: list[Report] = []
     # work parameters will be None on user prompts
     work_parameters: WorkParameters | None
+    safe_content: str | None
+    safety_level: int | None
+    safety_label: str | None
+    safety_rots: str | None
 
     @property
     def is_assistant(self) -> bool:
@@ -240,6 +244,14 @@ class PongResponse(WorkerResponseBase):
     metrics: WorkerMetricsInfo | None = None
 
 
+class SafePromptResponse(WorkerResponseBase):
+    response_type: Literal["safe_prompt"] = "safe_prompt"
+    safe_prompt: str
+    safety_parameters: SafetyParameters
+    safety_label: str
+    safety_rots: str
+
+
 class TokenResponse(WorkerResponseBase):
     response_type: Literal["token"] = "token"
     text: str
@@ -298,6 +310,7 @@ class GeneralErrorResponse(WorkerResponseBase):
         PongResponse,
         InternalFinishedMessageResponse,
         InternalErrorResponse,
+        SafePromptResponse,
     ],
     pydantic.Field(discriminator="response_type"),
 ]