FEAT: Gradio HiTL Scorer

doc/api.rst

@@ -365,6 +365,7 @@ API Reference
     FloatScaleThresholdScorer
     GandalfScorer
     HumanInTheLoopScorer
+    HumanInTheLoopScorerGradio
     LikertScalePaths
     MarkdownInjectionScorer
     PromptShieldScorer

pyproject.toml

@@ -37,7 +37,7 @@ classifiers = [
 requires-python = ">=3.10, <3.13"
 dependencies = [
     "aioconsole>=0.7.1",
-    "aiofiles>=24.1.0",
+    "aiofiles==23.2.1", # Pin the version to downgrade aiofiles to make sure it works with Gradio.
     "appdirs>=1.4.0",
     "art==6.1.0",
     "azure-cognitiveservices-speech>=1.36.0",
@@ -113,6 +113,12 @@ playwright = [
     "ollama>=0.4.4"
 ]
 
+gradio = [
+    "gradio>=5.16.0",
+    "rpyc>=6.0.1",
+    "pywebview>=5.4"
+]
+
 all = [
     "accelerate==0.34.2",
     "azureml-mlflow==1.57.0",
@@ -138,6 +144,9 @@ all = [
     "flask>=3.1.0",
     "ollama>=0.4.4",
     "types-PyYAML>=6.0.12.9",
+    "gradio>=5.16.0",
+    "rpyc>=6.0.1",
+    "pywebview>=5.4"
 ]
 
 [project.scripts]

pyrit/score/__init__.py

@@ -5,6 +5,7 @@
 from pyrit.score.float_scale_threshold_scorer import FloatScaleThresholdScorer
 from pyrit.score.gandalf_scorer import GandalfScorer
 from pyrit.score.human_in_the_loop_scorer import HumanInTheLoopScorer
+from pyrit.score.human_in_the_loop_gradio import HumanInTheLoopScorerGradio
 from pyrit.score.insecure_code_scorer import InsecureCodeScorer
 from pyrit.score.markdown_injection import MarkdownInjectionScorer
 from pyrit.score.prompt_shield_scorer import PromptShieldScorer
@@ -23,6 +24,7 @@
     "FloatScaleThresholdScorer",
     "GandalfScorer",
     "HumanInTheLoopScorer",
+    "HumanInTheLoopScorerGradio",
     "InsecureCodeScorer",
     "LikertScalePaths",
     "MarkdownInjectionScorer",

pyrit/score/human_in_the_loop_gradio.py

@@ -0,0 +1,47 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import asyncio
+from pyrit.score.scorer import Scorer
+from pyrit.models import Score, PromptRequestPiece
+from typing import Optional
+
+class HumanInTheLoopScorerGradio(Scorer):
+    """
+    Create scores from manual human input using Gradio and adds them to the database.
+
+    Parameters:
+        scorer (Scorer): The scorer to use for the initial scoring.
+        re_scorers (list[Scorer]): The scorers to use for re-scoring.
+        open_browser(bool): The scorer will open the Gradio interface in a browser instead of opening it in PyWebview
+    """
+
+    def __init__(self, *, open_browser=False, scorer: Scorer = None, re_scorers: list[Scorer] = None) -> None:
+        # Import here to avoid importing rpyc in the main module that might not be installed
+        from pyrit.ui.rpc import AppRPCServer
+
+        self._scorer = scorer
+        self._re_scorers = re_scorers
+        self._rpc_server = AppRPCServer(open_browser=open_browser)
+        self._rpc_server.start()
+
+
+    async def score_async(self, request_response: PromptRequestPiece, *, task: Optional[str] = None) -> list[Score]:
+        try:
+            return await asyncio.to_thread(self.score_prompt_manually, request_response, task=task)
+        except asyncio.CancelledError:
+            self._rpc_server.stop()
+            raise
+
+
+    def score_prompt_manually(self, request_prompt: PromptRequestPiece, *, task: Optional[str] = None) -> list[Score]:
+        self._rpc_server.wait_for_client()
+        self._rpc_server.send_score_prompt(request_prompt)
+        score = self._rpc_server.wait_for_score()
+        return [score]
+
+    def validate(self, request_response: PromptRequestPiece, *, task: Optional[str] = None):
+        pass
+
+    def __del__(self):
+        self._rpc_server.stop()

pyrit/ui/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.

pyrit/ui/app.py

@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import os
+import sys
+import subprocess
+import traceback
+
+GLOBAL_MUTEX_NAME = "PyRIT-Gradio"
+
+def launch_app(open_browser=False):
+    # Launch a new process to run the gradio UI.
+    # Locate the python executable and run this file.
+    current_path = os.path.abspath(__file__)
+    python_path = sys.executable
+
+    # Start a new process to run it
+    subprocess.Popen([python_path, current_path, str(open_browser)], creationflags=subprocess.CREATE_NEW_CONSOLE)
+
+def is_app_running():
+    if sys.platform != "win32":
+        raise NotImplementedError("This function is only supported on Windows.")
+        return True
+
+    import ctypes.wintypes
+
+    SYNCHRONIZE = 0x00100000
+    mutex = ctypes.windll.kernel32.OpenMutexW(SYNCHRONIZE, False, GLOBAL_MUTEX_NAME)
+    if not mutex:
+        return False
+
+    # Close the handle to the mutex
+    ctypes.windll.kernel32.CloseHandle(mutex)
+    return True
+
+if __name__ == "__main__":
+    def create_mutex():
+        if sys.platform != "win32":
+            raise NotImplementedError("This function is only supported on Windows.")
+
+        # TODO make sure to add cross-platform support for this.
+        import ctypes.wintypes
+        mutex = ctypes.windll.kernel32.CreateMutexW(None, False, GLOBAL_MUTEX_NAME)
+        last_error = ctypes.windll.kernel32.GetLastError()
+        if last_error == 183:  # ERROR_ALREADY_EXISTS
+            return False
+        return True
+
+    if not create_mutex():
+        print("Gradio UI is already running.")
+        sys.exit(1)
+    print("Starting Gradio Interface please wait...")
+    try:
+        open_browser = False
+        if len(sys.argv) > 1:
+            open_browser = sys.argv[1] == "True"
+
+        from scorer import GradioApp
+        app = GradioApp()
+        app.start_gradio(open_browser=open_browser)
+    except:
+        # Print the error message and traceback
+        print(traceback.format_exc())
+        input("Press Enter to exit.")

pyrit/ui/connection_status.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import gradio as gr
+
+from rpc_client import RPCClient
+
+class ConnectionStatusHandler:
+    def __init__(self,
+                 is_connected_state: gr.State,
+                 rpc_client: RPCClient):
+        self.state = is_connected_state
+        self.server_disconnected = False
+        self.rpc_client = rpc_client
+        self.next_prompt = ""
+
+    def setup(self, *, main_interface: gr.Column, loading_animation: gr.Column, next_prompt_state: gr.State):
+        self.state.change(fn=self._on_state_change, inputs=[self.state], outputs=[main_interface, loading_animation, next_prompt_state])
+
+        connection_status_timer = gr.Timer(1)
+        connection_status_timer.tick(
+            fn=self._check_connection_status,
+            inputs=[self.state],
+            outputs=[self.state]
+        ).then(
+            fn=self._reconnect_if_needed,
+            outputs=[self.state]
+        )
+
+    def set_ready(self):
+        self.server_disconnected = False
+
+    def set_disconnected(self):
+        self.server_disconnected = True
+
+    def set_next_prompt(self, next_prompt: str):
+        self.next_prompt = next_prompt
+
+    def _on_state_change(self, is_connected: bool):
+        print("Connection status changed to: ", is_connected, " - ", self.next_prompt)
+        if is_connected:
+            return [gr.Column(visible=True), gr.Row(visible=False), self.next_prompt]
+        return [gr.Column(visible=False), gr.Row(visible=True), self.next_prompt]
+
+    def _check_connection_status(self, is_connected: bool):
+        if self.server_disconnected or not is_connected:
+            print("Gradio disconnected")
+            return False
+        return True
+
+    def _reconnect_if_needed(self):
+        if self.server_disconnected:
+            print("Attempting to reconnect")
+            self.rpc_client.reconnect()
+            prompt = self.rpc_client.wait_for_prompt()
+            self.next_prompt = str(prompt.original_value)
+            self.server_disconnected = False
+        return True