From 472011669bbad87af08b4d251525e1bb9fcb7c08 Mon Sep 17 00:00:00 2001
From: AssemblyAI <engineering.sdk@assemblyai.com>
Date: Thu, 1 Jun 2023 14:53:48 +0200
Subject: [PATCH 1/2] Project import generated by Copybara.

GitOrigin-RevId: b1a1ba89e2e7565dee26cd8df86c7a04c643f9aa
---
 assemblyai/types.py        | 37 ++++++++++++++++++++++++++++++++-----
 setup.py                   |  2 +-
 tests/unit/test_domains.py |  1 +
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/assemblyai/types.py b/assemblyai/types.py
index 12789bc..7f64907 100644
--- a/assemblyai/types.py
+++ b/assemblyai/types.py
@@ -336,6 +336,9 @@ class RawTranscriptionConfig(BaseModel):
     speaker_labels: Optional[bool]
     "Enable Speaker Diarization."
 
+    speakers_expected: Optional[int]
+    "The number of speakers you expect to be in your audio file."
+
     # content_safety: bool = False
     # "Enable Content Safety Detection."
 
@@ -406,6 +409,7 @@ def __init__(
         redact_pii_policies: Optional[PIIRedactionPolicy] = None,
         redact_pii_sub: Optional[PIISubstitutionPolicy] = None,
         speaker_labels: Optional[bool] = None,
+        speakers_expected: Optional[int] = None,
         # content_safety: bool = False,
         # iab_categories: bool = False,
         custom_spelling: Optional[Dict[str, Union[str, Sequence[str]]]] = None,
@@ -439,6 +443,7 @@ def __init__(
             redact_pii_policies: The list of PII Redaction policies to enable.
             redact_pii_sub: The replacement logic for detected PII.
             speaker_labels: Enable Speaker Diarization.
+            speakers_expected: The number of speakers you expect to hear in your audio file. Up to 10 speakers are supported.
             content_safety: Enable Content Safety Detection.
             iab_categories: Enable Topic Detection.
             custom_spelling: Customize how words are spelled and formatted using to and from values.
@@ -480,7 +485,7 @@ def __init__(
             redact_pii_policies,
             redact_pii_sub,
         )
-        self.speaker_labels = speaker_labels
+        self.set_speaker_diarization(speaker_labels, speakers_expected)
         # self.content_safety = content_safety
         # self.iab_categories = iab_categories
         self.set_custom_spelling(custom_spelling, override=True)
@@ -633,11 +638,11 @@ def speaker_labels(self) -> Optional[bool]:
 
         return self._raw_transcription_config.speaker_labels
 
-    @speaker_labels.setter
-    def speaker_labels(self, enable: Optional[bool]) -> None:
-        "Enable Speaker Diarization feature."
+    @property
+    def speakers_expected(self) -> Optional[int]:
+        "Returns the number of speakers expected to be in the audio file. Used in combination with the `speaker_labels` parameter."
 
-        self._raw_transcription_config.speaker_labels = enable
+        return self._raw_transcription_config.speakers_expected
 
     # @property
     # def content_safety(self) -> bool:
@@ -799,6 +804,28 @@ def set_casing_and_formatting(
 
         return self
 
+    def set_speaker_diarization(
+        self,
+        enable: bool = True,
+        speakers_expected: Optional[int] = None,
+    ) -> Self:
+        """
+        Whether to enable Speaker Diarization on the transcript.
+
+        Args:
+            `enable`: Enable Speaker Diarization
+            `speakers_expected`: The number of speakers in the audio file.
+        """
+
+        if not enable:
+            self._raw_transcription_config.speaker_labels = None
+            self._raw_transcription_config.speakers_expected = None
+        else:
+            self._raw_transcription_config.speaker_labels = True
+            self._raw_transcription_config.speakers_expected = speakers_expected
+
+        return self
+
     def set_webhook(
         self,
         url: Optional[str],
diff --git a/setup.py b/setup.py
index b0e1cf1..d24ab01 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name="assemblyai",
-    version="0.4.2",
+    version="0.4.1",
     description="AssemblyAI Python SDK",
     author="AssemblyAI",
     author_email="engineering.sdk@assemblyai.com",
diff --git a/tests/unit/test_domains.py b/tests/unit/test_domains.py
index 5730ad0..7ef38bd 100644
--- a/tests/unit/test_domains.py
+++ b/tests/unit/test_domains.py
@@ -18,6 +18,7 @@ def test_configuration_drift():
         "set_redact_pii",  # PII redaction
         "set_summarize",  # summarization
         "set_webhook",  # webhook
+        "set_speaker_diarization",  # speaker diarization
     }
 
     # get all members

From 43f13ece4430ac01a20b8bafc15cb5e6fa02de0c Mon Sep 17 00:00:00 2001
From: Soheyl <soheyl@assemblyai.com>
Date: Fri, 2 Jun 2023 12:13:10 +0200
Subject: [PATCH 2/2] build: bump version to `0.5.0` due breaking change

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index d24ab01..7d615d0 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name="assemblyai",
-    version="0.4.1",
+    version="0.5.0",
     description="AssemblyAI Python SDK",
     author="AssemblyAI",
     author_email="engineering.sdk@assemblyai.com",