From c028c4b58455b648a4c69d0129d7f11ce2921d20 Mon Sep 17 00:00:00 2001
From: Aaron Boxer <aaron.boxer@collabora.com>
Date: Mon, 11 May 2026 19:37:10 -0400
Subject: [PATCH] feat: add segment_post_processor hook for external plugins

Add a minimal, non-breaking hook to WhisperLive that allows external
projects to post-process transcription segments before they are sent
to the client.

Changes:
- ServeClientBase: add segment_post_processor attribute (default None)
- ServeClientBase.send_transcription_to_client: apply post_processor
  per-segment with error handling (falls back to original segment)
- TranscriptionServer: add segment_post_processor parameter to run()
  and wire it to each client on creation

This enables downstream projects to plug in custom processing
(e.g. formatting, PII redaction, diarization tagging) without
modifying WhisperLive core code.
---
 whisper_live/backend/base.py | 21 +++++++++++++++++++++
 whisper_live/server.py       | 14 +++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/whisper_live/backend/base.py b/whisper_live/backend/base.py
index c6e6a7e..4d926dc 100644
--- a/whisper_live/backend/base.py
+++ b/whisper_live/backend/base.py
@@ -63,6 +63,13 @@ class ServeClientBase(object):
         self.end_time_for_same_output = None
         self.translation_queue = translation_queue
 
+        # Optional post-processing callable for segments.
+        # If set, called with a segment dict and must return a segment dict.
+        # Allows external projects to plug in custom post-processing
+        # (e.g. PII redaction, formatting, diarization) without modifying
+        # WhisperLive's core code.
+        self.segment_post_processor = None
+
         # threading
         self.lock = threading.Lock()
 
@@ -254,9 +261,23 @@ class ServeClientBase(object):
         This method formats the transcription segments into a JSON object and attempts to send
         this object to the client. If an error occurs during the send operation, it logs the error.
 
+        If a ``segment_post_processor`` callable is set, each segment is passed through it
+        before sending. The callable receives a segment dict and must return a segment dict.
+
         Returns:
             segments (list): A list of transcription segments to be sent to the client.
         """
+        if self.segment_post_processor is not None:
+            processed = []
+            for seg in segments:
+                try:
+                    result = self.segment_post_processor(seg)
+                    processed.append(result if result is not None else seg)
+                except Exception as e:
+                    logging.error(f"[ERROR]: segment_post_processor failed: {e}")
+                    processed.append(seg)
+            segments = processed
+
         try:
             self.websocket.send(
                 json.dumps({
diff --git a/whisper_live/server.py b/whisper_live/server.py
index e420a47..9798618 100644
--- a/whisper_live/server.py
+++ b/whisper_live/server.py
@@ -176,6 +176,7 @@ class TranscriptionServer:
         self.single_model = False
         self.batch_config = None
         self.raw_pcm_input = False
+        self.segment_post_processor = None
 
     def initialize_client(
         self, websocket, options, faster_whisper_custom_model_path,
@@ -316,6 +317,10 @@ class TranscriptionServer:
         if client is None:
             raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.")
 
+        # Attach segment post-processor if configured
+        if self.segment_post_processor is not None:
+            client.segment_post_processor = self.segment_post_processor
+
         if translation_client:
             client.translation_client = translation_client
             client.translation_thread = translation_thread
@@ -477,7 +482,8 @@ class TranscriptionServer:
             batch_max_size=8,
             batch_window_ms=50,
             raw_pcm_input=False,
-            metrics_port: int = 0):
+            metrics_port: int = 0,
+            segment_post_processor=None):
         """
         Run the transcription server.
 
@@ -493,6 +499,11 @@ class TranscriptionServer:
             batch_window_ms (int): Maximum time in milliseconds to wait for
                 the batch to fill after the first request arrives. Defaults
                 to 50.
+            segment_post_processor (callable, optional): A callable that receives
+                a transcription segment dict and returns a modified segment dict.
+                Applied to every segment before sending to the client. Useful for
+                plugging in custom post-processing (e.g. formatting, redaction).
+                Defaults to None.
         """
         self.cache_path = cache_path
         self.raw_pcm_input = raw_pcm_input
@@ -506,6 +517,7 @@ class TranscriptionServer:
         if batch_enabled and batch_window_ms < 0:
             raise ValueError(f"batch_window_ms must be >= 0, got {batch_window_ms}")
 
+        self.segment_post_processor = segment_post_processor
         self.client_manager = ClientManager(max_clients, max_connection_time)
         if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path):
             if "/" not in faster_whisper_custom_model_path: