Use chat moderation endpoint

app.py

@@ -1,4 +1,3 @@
-import asyncio
 import json
 import os
 
@@ -41,75 +40,49 @@ alinia_guardrail = httpx.AsyncClient(
 mistral_client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
 
 
-async def get_mistral_moderation(user_content, assistant_content):
-    def sync_moderation(inputs):
-        return mistral_client.classifiers.moderate_chat(
-            model="mistral-moderation-latest", inputs=inputs
-        )
-
-    inputs_assistant = [
-        {"role": "user", "content": user_content},
-        {"role": "assistant", "content": assistant_content},
-    ]
-
-    inputs_user = [{"role": "user", "content": user_content}]
-
+async def mistral_moderate(chat_history):
     try:
-        response_full, response_user_only = await asyncio.gather(
-            asyncio.to_thread(sync_moderation, inputs_assistant),
-            asyncio.to_thread(sync_moderation, inputs_user),
+        response_full = await mistral_client.classifiers.moderate_chat_async(
+            model="mistral-moderation-latest", inputs=chat_history[-2:]
         )
-
-        return {
-            "full_interaction": response_full.results,
-            "user_only": response_user_only.results,
-        }
+        response_user_only = await mistral_client.classifiers.moderate_chat_async(
+            model="mistral-moderation-latest", inputs=chat_history[-1:]
+        )
+        result = {}
+        if response_full.results:
+            result["full_interaction"] = response_full.results[0].category_scores
+        if response_user_only.results:
+            result["user_only"] = response_user_only.results[0].category_scores
     except Exception as e:
-        print(f"Mistral moderation error: {e!s}")
-        return {"error": str(e)}
+        message = f"Mistral moderate failed: {e!s}"
+        print(message)
+        result = message
+    return result
 
 
-async def check_safety(message: str, metadata: dict) -> dict:
+async def alinia_moderate(chat_history, chat_model_id, mistral_moderation) -> dict:
     try:
-        user_content = (
-            metadata["messages"][-2]["content"]
-            if len(metadata.get("messages", [])) >= 2
-            else ""
-        )
-        # Mistral moderation results
-        try:
-            mistral_response = await get_mistral_moderation(user_content, message)
-            mistral_results = mistral_response
-        except Exception as e:
-            print(f"[Mistral moderation error]: {e!s}")
-            mistral_results = None
-
         resp = await alinia_guardrail.post(
-            "/moderations/",
+            "/chat/moderations",
             json={
-                "input": message,
+                "messages": chat_history[-2:],
                 "metadata": {
                     "app": "slmdr",
                     "app_environment": "stable",
-                    "chat_model_id": MODEL_ARGS["model"],
-                    "mistral_results": json.loads(
-                        json.dumps(mistral_results, default=str)
-                    ),
-                }
-                | metadata,
+                    "chat_model_id": chat_model_id,
+                    "mistral_results": json.dumps(mistral_moderation, default=str),
+                },
                 "detection_config": {"safety": True},
             },
         )
         resp.raise_for_status()
-        result = resp.json()
-        selected_results = result["result"]["category_details"]["safety"]
-        selected_results = {
-            key.title(): value for key, value in selected_results.items()
-        }
-        return selected_results
+        moderation = resp.json()["result"]["category_details"]["safety"]
+        result = {key.title(): value for key, value in moderation.items()}
     except Exception as e:
-        print(f"Safety check error: {e!s}")
-        return {"Error": str(e)}
+        message = f"Alinia moderate failed: {e!s}"
+        print(message)
+        result = {"Error": 1}
+    return result
 
 
 def user(message, chat_history):
@@ -118,16 +91,15 @@ def user(message, chat_history):
 
 
 async def assistant(chat_history, system_prompt, model_name):
-    try:
-        client = CHAT_CLIENTS[model_name]
-
-        if chat_history[0]["role"] != "system":
-            chat_history = [{"role": "system", "content": system_prompt}, *chat_history]
+    client = CHAT_CLIENTS[model_name]
+    alinia_moderation = {}
 
-        chat_history.append({"role": "assistant", "content": ""})
+    if chat_history[0]["role"] != "system":
+        chat_history = [{"role": "system", "content": system_prompt}, *chat_history]
 
-        print(chat_history)
+    chat_history.append({"role": "assistant", "content": ""})
 
+    try:
         stream = await client.chat.completions.create(
             **MODEL_ARGS, messages=chat_history
         )
@@ -135,17 +107,18 @@ async def assistant(chat_history, system_prompt, model_name):
         async for chunk in stream:
             if chunk.choices[0].delta.content is not None:
                 chat_history[-1]["content"] += chunk.choices[0].delta.content
-                yield chat_history, ""
-
-        # metadata = {
-        #     "messages": chat_history + [{"role": "assistant", "content": full_response}]
-        # }
-        safety_results = await check_safety(chat_history[-1]["content"], {})
-        yield chat_history, safety_results
+                yield chat_history, alinia_moderation
 
+        mistral_moderation = await mistral_moderate(chat_history)
+        alinia_moderation = await alinia_moderate(
+            chat_history,
+            chat_model_id=model_name,
+            mistral_moderation=mistral_moderation,
+        )
     except Exception as e:
-        chat_history.append({"role": "assistant", "content": f"Error occurred: {e!s}"})
-        yield chat_history, ""
+        chat_history[-1]["content"] = f"Error occurred: {e!s}"
+
+    yield chat_history, alinia_moderation
 
 
 with gr.Blocks(title="🦎 Salamandra & Oranguten") as demo:
@@ -187,12 +160,12 @@ with gr.Blocks(title="🦎 Salamandra & Oranguten") as demo:
         outputs=[chatbot, response_safety],
     )
 
+    model_selector.change(lambda: [], outputs=chatbot)
     system_prompt_selector.change(
         lambda example_name: EXAMPLE_PROMPTS[example_name],
         inputs=system_prompt_selector,
         outputs=system_prompt,
     )
-
     system_prompt.change(lambda: [], outputs=chatbot)
 
     new_chat.click(