feat: add content_safety functionality

Co-authored-by: Soheyl <[email protected]>
GitOrigin-RevId: ade523b5a3e103d0f57fd1bdc424d80b429214d1

Author: dmccrystal

README.md

@@ -308,6 +308,54 @@ config=aai.TranscriptionConfig(
 )
 ```
 
+</details>
+<details>
+  <summary>Detect Sensitive Content in a Transcript</summary>
+
+```python
+import assemblyai as aai
+
+transcriber = aai.Transcriber()
+transcript = transcriber.transcribe(
+  "https://example.org/audio.mp3",
+  config=aai.TranscriptionConfig(content_safety=True)
+)
+
+
+# Get the parts of the transcript which were flagged as sensitive
+for result in transcript.content_safety_labels.results:
+  print(result.text)  # sensitive text snippet
+  print(result.timestamp.start)
+  print(result.timestamp.end)
+
+  for label in result.labels:
+    print(label.label)  # content safety category
+    print(label.confidence) # model's confidence that the text is in this category
+    print(label.severity) # severity of the text in relation to the category
+
+# Get the confidence of the most common labels in relation to the entire audio file
+for label, confidence in transcript.content_safety_labels.summary.items():
+  print(f"{confidence * 100}% confident that the audio contains {label}")
+
+# Get the overall severity of the most common labels in relation to the entire audio file
+for label, severity_confidence in transcript.content_safety_labels.severity_score_summary.items():
+  print(f"{severity_confidence.low * 100}% confident that the audio contains low-severity {label}")
+  print(f"{severity_confidence.medium * 100}% confident that the audio contains mid-severity {label}")
+  print(f"{severity_confidence.high * 100}% confident that the audio contains high-severity {label}")
+
+```
+
+[Read more about the content safety categories.](https://www.assemblyai.com/docs/Models/content_moderation#all-labels-supported-by-the-model)
+
+By default, the content safety model will only include labels with a confidence greater than 0.5 (50%). To change this, pass `content_safety_confidence` (as an integer percentage between 25 and 100, inclusive) to the `TranscriptionConfig`:
+
+```python
+config=aai.TranscriptionConfig(
+  content_safety=True,
+  content_safety_confidence=80,  # only include labels with a confidence greater than 80%
+)
+```
+
 </details>
 
 ---

assemblyai/api.py

@@ -34,7 +34,6 @@ def create_transcript(
             by_alias=True,
         ),
     )
-
     if response.status_code != httpx.codes.ok:
         raise types.TranscriptError(
             f"failed to transcript url {request.audio_url}: {_get_error_message(response)}"

assemblyai/transcriber.py

@@ -214,6 +214,10 @@ def summary(self) -> Optional[str]:
     def chapters(self) -> Optional[List[types.Chapter]]:
         return self._impl.transcript.chapters
 
+    @property
+    def content_safety_labels(self) -> Optional[types.ContentSafetyResponse]:
+        return self._impl.transcript.content_safety_labels
+
     @property
     def status(self) -> types.TranscriptStatus:
         "The current status of the transcript"

assemblyai/types.py

@@ -339,8 +339,11 @@ class RawTranscriptionConfig(BaseModel):
     speakers_expected: Optional[int]
     "The number of speakers you expect to be in your audio file."
 
-    # content_safety: bool = False
-    # "Enable Content Safety Detection."
+    content_safety: Optional[bool]
+    "Enable Content Safety Detection."
+
+    content_safety_confidence: Optional[int]
+    "The minimum confidence level for a content safety label to be produced."
 
     # iab_categories: bool = False
     # "Enable Topic Detection."
@@ -410,7 +413,8 @@ def __init__(
         redact_pii_sub: Optional[PIISubstitutionPolicy] = None,
         speaker_labels: Optional[bool] = None,
         speakers_expected: Optional[int] = None,
-        # content_safety: bool = False,
+        content_safety: Optional[bool] = None,
+        content_safety_confidence: Optional[int] = None,
         # iab_categories: bool = False,
         custom_spelling: Optional[Dict[str, Union[str, Sequence[str]]]] = None,
         disfluencies: Optional[bool] = None,
@@ -486,7 +490,7 @@ def __init__(
             redact_pii_sub,
         )
         self.set_speaker_diarization(speaker_labels, speakers_expected)
-        # self.content_safety = content_safety
+        self.set_content_safety(content_safety, content_safety_confidence)
         # self.iab_categories = iab_categories
         self.set_custom_spelling(custom_spelling, override=True)
         self.disfluencies = disfluencies
@@ -644,17 +648,51 @@ def speakers_expected(self) -> Optional[int]:
 
         return self._raw_transcription_config.speakers_expected
 
-    # @property
-    # def content_safety(self) -> bool:
-    #     "Returns the status of the Content Safety feature."
+    @property
+    def content_safety(self) -> Optional[bool]:
+        "Returns the status of the Content Safety feature."
+
+        return self._raw_transcription_config.content_safety
+
+    @property
+    def content_safety_confidence(self) -> Optional[int]:
+        "The minimum confidence level for a content safety label to be produced. Used in combination with the `content_safety` parameter."
 
-    #     return self._raw_transcription_config.content_safety
+        return self._raw_transcription_config.content_safety_confidence
 
-    # @content_safety.setter
-    # def content_safety(self, enable: bool) -> None:
-    #     "Enable Content Safety feature."
+    def set_content_safety(
+        self,
+        enable: Optional[bool] = True,
+        content_safety_confidence: Optional[int] = None,
+    ) -> Self:
+        """Enable Content Safety feature.
+
+        Args:
+            `enable`: Whether or not to enable the Content Safety feature.
+            `content_safety_confidence`: The minimum confidence level for a content safety label to be produced.
 
-    #     self._raw_transcription_config.content_safety = enable
+        Raises:
+            `ValueError`: Raised if `content_safety_confidence` is not between 25 and 100 (inclusive).
+        """
+
+        if not enable:
+            self._raw_transcription_config.content_safety = None
+            self._raw_transcription_config.content_safety_confidence = None
+            return self
+
+        if content_safety_confidence is not None and (
+            content_safety_confidence < 25 or content_safety_confidence > 100
+        ):
+            raise ValueError(
+                "content_safety_confidence must be between 25 and 100 (inclusive)."
+            )
+
+        self._raw_transcription_config.content_safety = enable
+        self._raw_transcription_config.content_safety_confidence = (
+            content_safety_confidence
+        )
+
+        return self
 
     # @property
     # def iab_categories(self) -> bool:
@@ -1162,7 +1200,7 @@ class AutohighlightResponse(BaseModel):
 class ContentSafetyLabelResult(BaseModel):
     label: ContentSafetyLabel
     confidence: float
-    severity: float
+    severity: Optional[float]
 
 
 class ContentSafetySeverityScore(BaseModel):
@@ -1180,8 +1218,10 @@ class ContentSafetyResult(BaseModel):
 class ContentSafetyResponse(BaseModel):
     status: StatusResult
     results: Optional[List[ContentSafetyResult]]
-    summary: Optional[Dict[str, float]]
-    severity_score_summary: Optional[Dict[str, ContentSafetySeverityScore]]
+    summary: Optional[Dict[ContentSafetyLabel, float]]
+    severity_score_summary: Optional[
+        Dict[ContentSafetyLabel, ContentSafetySeverityScore]
+    ]
 
 
 class IABLabelResult(BaseModel):
@@ -1308,8 +1348,11 @@ class BaseTranscript(BaseModel):
     speaker_labels: Optional[bool]
     "Enable Speaker Diarization."
 
-    # content_safety: bool = False
-    # "Enable Content Safety Detection."
+    content_safety: Optional[bool]
+    "Enable Content Safety Detection."
+
+    content_safety_confidence: Optional[int]
+    "The minimum confidence level for a content safety label to be produced."
 
     # iab_categories: bool = False
     # "Enable Topic Detection."
@@ -1401,8 +1444,8 @@ class TranscriptResponse(BaseTranscript):
     # auto_highlights_result: Optional[AutohighlightResponse] = None
     # "The list of results when enabling Automatic Transcript Highlights"
 
-    # content_safety_labels: Optional[ContentSafetyResponse] = None
-    # "The list of results when Content Safety is enabled"
+    content_safety_labels: Optional[ContentSafetyResponse]
+    "The list of results when Content Safety is enabled"
 
     # iab_categories_result: Optional[IABResponse] = None
     # "The list of results when Topic Detection is enabled"
@@ -1416,15 +1459,18 @@ class TranscriptResponse(BaseTranscript):
     # entities: Optional[List[Entity]] = None
     # "When Entity Detection is enabled, the list of detected Entities"
 
-    # def __init__(self, **data: Any):
-    #     # cleanup the response before creating the object
-    #     if data.get("iab_categories_result") == {}:
-    #         data["iab_categories_result"] = None
+    def __init__(self, **data: Any):
+        # cleanup the response before creating the object
+        # if data.get("iab_categories_result") == {}:
+        #     data["iab_categories_result"] = None
 
-    #     if data.get("content_safety_labels") == {}:
-    #         data["content_safety_labels"] = None
+        if data.get("content_safety_labels") == {} or (
+            not data.get("content_safety")
+            and data.get("content_safety_labels", {}).get("status") == "unavailable"
+        ):
+            data["content_safety_labels"] = None
 
-    #     super().__init__(**data)
+        super().__init__(**data)
 
 
 class LemurModel(str, Enum):

tests/unit/factories.py

@@ -3,6 +3,7 @@
 from AssemblyAI's API.
 """
 
+from enum import Enum
 from functools import partial
 from typing import Any, Callable, Dict
 
@@ -13,6 +14,14 @@
 from assemblyai import types
 
 
+class TimestampFactory(factory.Factory):
+    class Meta:
+        model = aai.Timestamp
+
+    start = factory.Faker("pyint")
+    end = factory.Faker("pyint")
+
+
 class WordFactory(factory.Factory):
     class Meta:
         model = aai.Word
@@ -234,6 +243,8 @@ def convert_dict_from_stub(stub: factory.base.StubObject) -> Dict[str, Any]:
                     if stub_is_list(value)
                     else convert_dict_from_stub(value)
                 )
+            elif isinstance(value, Enum):
+                stub_dict[key] = value.value
         return stub_dict
 
     def dict_factory(f, **kwargs):