Project import generated by Copybara.

GitOrigin-RevId: ea71d0332ba79dc276e93de36daf33ebb6066493

Author: AssemblyAI

README.md

@@ -23,7 +23,7 @@ With a single API call, get access to AI models built on the latest AI breakthro
 - [Example](#examples)
   - [Core Examples](#core-examples)
   - [LeMUR Examples](#lemur-examples)
-  - [Audio Intelligence+ Examples](#audio-intelligence-examples)
+  - [Audio Intelligence Examples](#audio-intelligence-examples)
 - [Playgrounds](#playgrounds)
 - [Advanced](#advanced-todo)
 
@@ -159,35 +159,6 @@ print(transcript.text)
 
 </details>
 
-<details>
-  <summary>Summarize the content of a transcript</summary>
-
-```python
-import assemblyai as aai
-
-transcriber = aai.Transcriber()
-transcript = transcriber.transcribe(
-  "https://example.org/audio.mp3",
-  config=aai.TranscriptionConfig(summarize=True)
-)
-
-print(transcript.summary)
-```
-
-By default, the summarization model will be `informative` and the summarization type will be `bullets`. [Read more about summarization models and types here](https://www.assemblyai.com/docs/Models/summarization#types-and-models).
-
-To change the model and/or type, pass additional parameters to the `TranscriptionConfig`:
-
-```python
-config=aai.TranscriptionConfig(
-  summarize=True,
-  summary_model=aai.SummarizationModel.catchy,
-  summary_type=aai.Summarizationtype.headline
-)
-```
-
-</details>
-
 ---
 
 ### **LeMUR Examples**
@@ -260,7 +231,7 @@ for result in result:
 
 ---
 
-### **Audio Intelligence+ Examples**
+### **Audio Intelligence Examples**
 
 <details>
   <summary>PII Redact a Transcript</summary>
@@ -286,6 +257,57 @@ transcriber = aai.Transcriber()
 transcript = transcriber.transcribe("https://example.org/audio.mp3", config)
 ```
 
+</details>
+<details>
+  <summary>Summarize the content of a transcript over time</summary>
+
+```python
+import assemblyai as aai
+
+transcriber = aai.Transcriber()
+transcript = transcriber.transcribe(
+  "https://example.org/audio.mp3",
+  config=aai.TranscriptionConfig(auto_chapters=True)
+)
+
+for chapter in transcript.chapters:
+  print(f"Summary: {chapter.summary}")  # A one paragraph summary of the content spoken during this timeframe
+  print(f"Start: {chapter.start}, End: {chapter.end}")  # Timestamps (in milliseconds) of the chapter
+  print(f"Healine: {chapter.headline}")  # A single sentence summary of the content spoken during this timeframe
+  print(f"Gist: {chapter.gist}")  # An ultra-short summary, just a few words, of the content spoken during this timeframe
+```
+
+[Read more about auto chapters here.](https://www.assemblyai.com/docs/Models/auto_chapters)
+
+</details>
+
+<details>
+  <summary>Summarize the content of a transcript</summary>
+
+```python
+import assemblyai as aai
+
+transcriber = aai.Transcriber()
+transcript = transcriber.transcribe(
+  "https://example.org/audio.mp3",
+  config=aai.TranscriptionConfig(summarization=True)
+)
+
+print(transcript.summary)
+```
+
+By default, the summarization model will be `informative` and the summarization type will be `bullets`. [Read more about summarization models and types here](https://www.assemblyai.com/docs/Models/summarization#types-and-models).
+
+To change the model and/or type, pass additional parameters to the `TranscriptionConfig`:
+
+```python
+config=aai.TranscriptionConfig(
+  summarization=True,
+  summary_model=aai.SummarizationModel.catchy,
+  summary_type=aai.SummarizationType.headline
+)
+```
+
 </details>
 
 ---
@@ -297,7 +319,6 @@ Visit one of our Playgrounds:
 - [LeMUR Playground](https://www.assemblyai.com/playground/v2/source)
 - [Transcription Playground](https://www.assemblyai.com/playground)
 
-
 # Advanced
 
 ## How the SDK handles Default Configurations
@@ -329,7 +350,6 @@ transcriber = aai.Transcriber()
 transcriber.config = aai.TranscriptionConfig(punctuate=False, format_text=False)
 ```
 
-
 In case you want to override the `Transcriber`'s configuration for a specific operation with a different one, you can do so via the `config` parameter of a `.transcribe*(...)` method:
 
 ```python

assemblyai/transcriber.py

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

assemblyai/types.py

@@ -354,8 +354,8 @@ class RawTranscriptionConfig(BaseModel):
     # sentiment_analysis: bool = False
     # "Enable Sentiment Analysis."
 
-    # auto_chapters: bool = False
-    # "Enable Auto Chapters."
+    auto_chapters: Optional[bool]
+    "Enable Auto Chapters."
 
     # entity_detection: bool = False
     # "Enable Entity Detection."
@@ -415,7 +415,7 @@ def __init__(
         custom_spelling: Optional[Dict[str, Union[str, Sequence[str]]]] = None,
         disfluencies: Optional[bool] = None,
         # sentiment_analysis: bool = False,
-        # auto_chapters: bool = False,
+        auto_chapters: Optional[bool] = None,
         # entity_detection: bool = False,
         summarization: Optional[bool] = None,
         summary_model: Optional[SummarizationModel] = None,
@@ -491,7 +491,7 @@ def __init__(
         self.set_custom_spelling(custom_spelling, override=True)
         self.disfluencies = disfluencies
         # self.sentiment_analysis = sentiment_analysis
-        # self.auto_chapters = auto_chapters
+        self.auto_chapters = auto_chapters
         # self.entity_detection = entity_detection
         self.set_summarize(
             summarization,
@@ -707,17 +707,23 @@ def disfluencies(self, enable: Optional[bool]) -> None:
 
     #     self._raw_transcription_config.sentiment_analysis = enable
 
-    # @property
-    # def auto_chapters(self) -> bool:
-    #     "Returns the status of the Auto Chapters feature."
+    @property
+    def auto_chapters(self) -> bool:
+        "Returns the status of the Auto Chapters feature."
+
+        return self._raw_transcription_config.auto_chapters
 
-    #     return self._raw_transcription_config.auto_chapters
+    @auto_chapters.setter
+    def auto_chapters(self, enable: bool) -> None:
+        "Enable Auto Chapters."
 
-    # @auto_chapters.setter
-    # def auto_chapters(self, enable: bool) -> None:
-    #     "Enable Auto Chapters."
+        # Validate required params are also set
+        if self.punctuate == False:
+            raise ValueError(
+                "If `auto_chapters` is enabled, then `punctuate` must not be disabled"
+            )
 
-    #     self._raw_transcription_config.auto_chapters = enable
+        self._raw_transcription_config.auto_chapters = enable
 
     # @property
     # def entity_detection(self) -> bool:
@@ -1317,8 +1323,8 @@ class BaseTranscript(BaseModel):
     # sentiment_analysis: bool = False
     # "Enable Sentiment Analysis."
 
-    # auto_chapters: bool = False
-    # "Enable Auto Chapters."
+    auto_chapters: Optional[bool]
+    "Enable Auto Chapters."
 
     # entity_detection: bool = False
     # "Enable Entity Detection."
@@ -1401,7 +1407,7 @@ class TranscriptResponse(BaseTranscript):
     # iab_categories_result: Optional[IABResponse] = None
     # "The list of results when Topic Detection is enabled"
 
-    # chapters: Optional[List[Chapter]] = None
+    chapters: Optional[List[Chapter]]
     # "When Auto Chapters is enabled, the list of Auto Chapters results"
 
     # sentiment_analysis_results: Optional[List[Sentiment]] = None

tests/e2e/__init__.py

@@ -0,0 +1,69 @@
+import pytest
+
+import assemblyai as aai
+
+
+def test_auto_chapters_disabled_by_default():
+    """
+    Tests that excluding `auto_chapters` from the `TranscriptionConfig` will
+    result in the default behavior of it being disabled
+    """
+    transcript = aai.Transcriber().transcribe(
+        data="https://assemblyai-test.s3.us-west-2.amazonaws.com/sdk/rogan-1min.mp3",
+        config=aai.TranscriptionConfig(),
+    )
+
+    assert transcript.status == aai.TranscriptStatus.completed
+    assert transcript.error is None
+    assert transcript.config.auto_chapters in (None, False)
+    assert transcript.chapters is None
+
+
+def test_auto_chapters_enabled():
+    """
+    Tests that including `auto_chapters=True` in the `TranscriptionConfig`
+    will enable the auto_chapters feature with a meaningful response
+    """
+    transcript = aai.Transcriber().transcribe(
+        data="https://assemblyai-test.s3.us-west-2.amazonaws.com/sdk/rogan-1min.mp3",
+        config=aai.TranscriptionConfig(auto_chapters=True),
+    )
+
+    assert transcript.status == aai.TranscriptStatus.completed
+    assert transcript.error is None
+
+    assert transcript.config.auto_chapters == True
+
+    assert transcript.chapters is not None
+    assert isinstance(transcript.chapters, list)
+    assert len(transcript.chapters) > 0
+
+    last_end_timestamp = 0
+    for chapter in transcript.chapters:
+        assert isinstance(chapter, aai.types.Chapter)
+        assert len(chapter.summary.strip()) > 0
+        assert len(chapter.headline.strip()) > 0
+        assert len(chapter.gist.strip()) > 0
+
+        assert chapter.start >= last_end_timestamp
+        assert chapter.end > chapter.start
+
+        last_end_timestamp = chapter.end
+
+
+def test_auto_chapters_failed():
+    """
+    Test that failure to produce auto_chapters will result in an error that
+    is properly wrapped by the `Transcript` object. In this case, the error
+    is that the French language model does not support auto_chapters.
+    """
+    transcript = aai.Transcriber().transcribe(
+        data="https://assemblyai-test.s3.us-west-2.amazonaws.com/sdk/rogan-1min.mp3",
+        config=aai.TranscriptionConfig(
+            auto_chapters=True, language_code=aai.LanguageCode.fr
+        ),
+    )
+
+    assert transcript.status == aai.TranscriptStatus.error
+    assert transcript.error is not None
+    assert "auto_chapters" in transcript.error

tests/e2e/test_auto_chapters_e2e.py

@@ -0,0 +1,83 @@
+import pytest
+
+import assemblyai.developer_tools.python.sdk as aai
+
+
+@pytest.mark.summarization
+@pytest.mark.parametrize(
+    "summary_model, summary_type",
+    [
+        (None, None),  # default is ("informative", "bullets")
+        (
+            aai.SummarizationModel.conversational,
+            aai.SummarizationType.paragraph,
+        ),
+        (
+            aai.SummarizationModel.conversational,
+            aai.SummarizationType.headline,
+        ),
+        (
+            aai.SummarizationModel.conversational,
+            aai.SummarizationType.bullets,
+        ),
+        (
+            aai.SummarizationModel.conversational,
+            aai.SummarizationType.bullets_verbose,
+        ),
+        (aai.SummarizationModel.catchy, aai.SummarizationType.headline),
+        (aai.SummarizationModel.catchy, aai.SummarizationType.gist),
+        (aai.SummarizationModel.informative, aai.SummarizationType.paragraph),
+        (aai.SummarizationModel.informative, aai.SummarizationType.headline),
+        (aai.SummarizationModel.informative, aai.SummarizationType.bullets),
+        (
+            aai.SummarizationModel.informative,
+            aai.SummarizationType.bullets_verbose,
+        ),
+    ],
+)
+def test_summarization_e2e(
+    summary_model: aai.SummarizationModel,
+    summary_type: aai.SummarizationType,
+):
+    """Test all combinations of transcription with summarization."""
+    config = aai.TranscriptionConfig(
+        summarization=True, summary_model=summary_model, summary_type=summary_type
+    )
+    if summary_model == aai.SummarizationModel.conversational:
+        config.set_speaker_diarization(True)
+
+    transcript = aai.Transcriber().transcribe(
+        data="https://assemblyai-test.s3.us-west-2.amazonaws.com/e2e_tests/summarization/kelley.wav",
+        config=config,
+    )
+
+    # Assign defaults
+    if not summary_type:
+        summary_type = aai.SummarizationType.bullets
+    if not summary_model:
+        summary_model = aai.SummarizationModel.informative
+
+    # Check that summarization was enabled
+    assert transcript.config.summarization
+
+    # Check that the response has a successful status
+    assert transcript.status == aai.TranscriptStatus.completed
+
+    # Check that the summary model and type match the request
+    assert transcript.config.summary_model == summary_model
+    assert transcript.config.summary_type == summary_type
+
+    # Check that there is no error message
+    assert not transcript.error
+
+    # Check that a summary exists on the transcript
+    assert transcript.summary
+
+    # Check that bulleted summaries start with dashes
+    if summary_type in [
+        aai.SummarizationType.bullets,
+        aai.SummarizationType.bullets_verbose,
+    ]:
+        assert transcript.summary.startswith("- ")
+    else:
+        assert not transcript.summary.startswith("- ")

tests/e2e/test_summarization_e2e.py

@@ -37,6 +37,17 @@ class Meta:
     words = factory.List([factory.SubFactory(UtteranceWordFactory)])
 
 
+class ChapterFactory(factory.Factory):
+    class Meta:
+        model = types.Chapter
+
+    summary = factory.Faker("sentence")
+    headline = factory.Faker("sentence")
+    gist = factory.Faker("sentence")
+    start = factory.Faker("pyint")
+    end = factory.Faker("pyint")
+
+
 class BaseTranscriptFactory(factory.Factory):
     class Meta:
         model = types.BaseTranscript