From ed568088484a84074cd0f30d9aef04766b2ab413 Mon Sep 17 00:00:00 2001
From: Hristo Stoychev <hristo@metaview.ai>
Date: Thu, 2 Nov 2023 15:28:35 +0200
Subject: [PATCH] Add support for diarization_type, speakers_count,
 remove_atmospherics

---
 src/rev_ai/apiclient.py               | 20 +++++++++++++++-----
 src/rev_ai/models/asynchronous/job.py |  8 ++++++--
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/rev_ai/apiclient.py b/src/rev_ai/apiclient.py
index ea94e809..44cad5cd 100644
--- a/src/rev_ai/apiclient.py
+++ b/src/rev_ai/apiclient.py
@@ -66,7 +66,8 @@ def submit_job_url(
             notification_config=None,
             skip_postprocessing=False,
             remove_atmospherics=False,
-            speakers_count=None):
+            speakers_count=None,
+            diarization_type=None):
         """Submit media given a URL for transcription.
         The audio data is downloaded from the URL
         :param media_url: web location of the media file
@@ -116,6 +117,7 @@ def submit_job_url(
         :param remove_atmospherics: Atmospherics such as <laugh>, <affirmative>, etc. will not
             appear in the transcript.
         :param speakers_count: Use to specify the total number of unique speakers in the audio.
+        :param diarization_type: Use to specify diarization type.
         :returns: raw response data
         :raises: HTTPError
         """
@@ -128,7 +130,8 @@ def submit_job_url(
                                                    verbatim, rush, test_mode,
                                                    segments_to_transcribe, speaker_names,
                                                    source_config, notification_config,
-                                                   skip_postprocessing)
+                                                   skip_postprocessing, remove_atmospherics,
+                                                   speakers_count, diarization_type)
 
         response = self._make_http_request(
             "POST",
@@ -161,7 +164,8 @@ def submit_job_local_file(
             notification_config=None,
             skip_postprocessing=False,
             remove_atmospherics=False,
-            speakers_count=None):
+            speakers_count=None,
+            diarization_type=None):
         """Submit a local file for transcription.
         Note that the content type is inferred if not provided.
 
@@ -208,6 +212,7 @@ def submit_job_local_file(
         :param remove_atmospherics: Atmospherics such as <laugh>, <affirmative>, etc. will not
             appear in the transcript.
         :param speakers_count: Use to specify the total number of unique speakers in the audio.
+        :param diarization_type: Use to specify diarization type.
         :returns: raw response data
         :raises: HTTPError, ValueError
         """
@@ -222,7 +227,9 @@ def submit_job_local_file(
                                                    language, custom_vocabulary_id, transcriber,
                                                    verbatim, rush, test_mode,
                                                    segments_to_transcribe, speaker_names, None,
-                                                   notification_config, skip_postprocessing)
+                                                   notification_config, skip_postprocessing,
+                                                   remove_atmospherics, speakers_count,
+                                                   diarization_type)
 
         with open(filename, 'rb') as f:
             files = {
@@ -475,7 +482,8 @@ def _create_job_options_payload(
             notification_config=None,
             skip_postprocessing=False,
             remove_atmospherics=None,
-            speakers_count=None):
+            speakers_count=None,
+            diarization_type=None):
         payload = {}
         if media_url:
             payload['media_url'] = media_url
@@ -524,6 +532,8 @@ def _create_job_options_payload(
             payload['remove_atmospherics'] = remove_atmospherics
         if speakers_count:
             payload['speakers_count'] = speakers_count
+        if diarization_type:
+            payload['diarization_type'] = diarization_type
         return payload
 
     def _create_captions_query(self, speaker_channel):
diff --git a/src/rev_ai/models/asynchronous/job.py b/src/rev_ai/models/asynchronous/job.py
index 609833b1..006da711 100644
--- a/src/rev_ai/models/asynchronous/job.py
+++ b/src/rev_ai/models/asynchronous/job.py
@@ -28,7 +28,8 @@ def __init__(
             rush=None,
             segments_to_transcribe=None,
             remove_atmospherics=None,
-            speakers_count=None):
+            speakers_count=None,
+            diarization_type=None):
         """
         :param id_: unique id of job
         :param created_on: date and time at which this job was started
@@ -59,6 +60,7 @@ def __init__(
         :param remove_atmospherics: Atmospherics such as <laugh>, <affirmative>, etc. will noT
             appear in the transcript.
         :param speakers_count: Use to specify the total number of unique speakers in the audio.
+        :param diarization_type: Use to specify diarization type.
         """
         self.id = id_
         self.created_on = created_on
@@ -85,6 +87,7 @@ def __init__(
         self.segments_to_transcribe = segments_to_transcribe
         self.remove_atmospherics = remove_atmospherics
         self.speakers_count = speakers_count
+        self.diarization_type = diarization_type
 
     def __eq__(self, other):
         """Override default equality operator"""
@@ -120,5 +123,6 @@ def from_json(cls, json):
             rush=json.get('rush'),
             segments_to_transcribe=json.get('segments_to_transcribe'),
             remove_atmospherics=json.get('remove_atmospherics'),
-            speakers_count=json.get('speakers_count')
+            speakers_count=json.get('speakers_count'),
+            diarization_type=json.get('diarization_type')
         )