feat: add entity_detection functionality

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

Author: dmccrystal

README.md

@@ -391,6 +391,27 @@ for sentiment_result in transcript.sentiment_analysis_results:
 
 [Read more about sentiment analysis here.](https://www.assemblyai.com/docs/Models/sentiment_analysis)
 
+</details>
+<details>
+  <summary>Identify Entities in a Transcript</summary>
+
+```python
+import assemblyai as aai
+
+transcriber = aai.Transcriber()
+transcript = transcriber.transcribe(
+  "https://example.org/audio.mp3",
+  config=aai.TranscriptionConfig(entity_detection=True)
+)
+
+for entity in transcript.entities:
+  print(entity.text) # i.e. "Dan Gilbert"
+  print(entity.type) # i.e. EntityType.person
+  print(f"Timestamp: {entity.start} - {entity.end}")
+```
+
+[Read more about entity detection here.](https://www.assemblyai.com/docs/Models/entity_detection)
+
 </details>
 
 ---

assemblyai/transcriber.py

@@ -222,6 +222,10 @@ def content_safety_labels(self) -> Optional[types.ContentSafetyResponse]:
     def sentiment_analysis_results(self) -> Optional[List[types.Sentiment]]:
         return self._impl.transcript.sentiment_analysis_results
 
+    @property
+    def entities(self) -> Optional[List[types.Entity]]:
+        return self._impl.transcript.entities
+
     @property
     def status(self) -> types.TranscriptStatus:
         "The current status of the transcript"

assemblyai/types.py

@@ -360,8 +360,8 @@ class RawTranscriptionConfig(BaseModel):
     auto_chapters: Optional[bool]
     "Enable Auto Chapters."
 
-    # entity_detection: bool = False
-    # "Enable Entity Detection."
+    entity_detection: Optional[bool]
+    "Enable Entity Detection."
 
     summarization: Optional[bool]
     "Enable Summarization"
@@ -420,7 +420,7 @@ def __init__(
         disfluencies: Optional[bool] = None,
         sentiment_analysis: Optional[bool] = None,
         auto_chapters: Optional[bool] = None,
-        # entity_detection: bool = False,
+        entity_detection: Optional[bool] = None,
         summarization: Optional[bool] = None,
         summary_model: Optional[SummarizationModel] = None,
         summary_type: Optional[SummarizationType] = None,
@@ -496,7 +496,7 @@ def __init__(
         self.disfluencies = disfluencies
         self.sentiment_analysis = sentiment_analysis
         self.auto_chapters = auto_chapters
-        # self.entity_detection = entity_detection
+        self.entity_detection = entity_detection
         self.set_summarize(
             summarization,
             summary_model,
@@ -763,17 +763,17 @@ def auto_chapters(self, enable: Optional[bool]) -> None:
 
         self._raw_transcription_config.auto_chapters = enable
 
-    # @property
-    # def entity_detection(self) -> bool:
-    #     "Returns whether Entity Detection feature is enabled or not."
+    @property
+    def entity_detection(self) -> bool:
+        "Returns whether Entity Detection feature is enabled or not."
 
-    #     return self._raw_transcription_config.entity_detection
+        return self._raw_transcription_config.entity_detection
 
-    # @entity_detection.setter
-    # def entity_detection(self, enable: bool) -> None:
-    #     "Enable Entity Detection."
+    @entity_detection.setter
+    def entity_detection(self, enable: Optional[bool]) -> None:
+        "Enable Entity Detection."
 
-    #     self._raw_transcription_config.entity_detection = enable
+        self._raw_transcription_config.entity_detection = enable
 
     @property
     def summarization(self) -> Optional[bool]:
@@ -1370,8 +1370,8 @@ class BaseTranscript(BaseModel):
     auto_chapters: Optional[bool]
     "Enable Auto Chapters."
 
-    # entity_detection: bool = False
-    # "Enable Entity Detection."
+    entity_detection: Optional[bool]
+    "Enable Entity Detection."
 
     summarization: Optional[bool]
     "Enable Summarization"
@@ -1457,8 +1457,8 @@ class TranscriptResponse(BaseTranscript):
     sentiment_analysis_results: Optional[List[Sentiment]]
     "When Sentiment Analysis is enabled, the list of Sentiment Analysis results"
 
-    # entities: Optional[List[Entity]] = None
-    # "When Entity Detection is enabled, the list of detected Entities"
+    entities: Optional[List[Entity]]
+    "When Entity Detection is enabled, the list of detected Entities"
 
     def __init__(self, **data: Any):
         # cleanup the response before creating the object

tests/unit/test_entity_detection.py

@@ -0,0 +1,118 @@
+import json
+from typing import Any, Dict, Tuple
+
+import factory
+import httpx
+from pytest_httpx import HTTPXMock
+
+import assemblyai as aai
+from tests.unit import factories
+
+aai.settings.api_key = "test"
+
+
+class EntityFactory(factory.Factory):
+    class Meta:
+        model = aai.types.Entity
+
+    entity_type = factory.Faker("enum", enum_cls=aai.types.EntityType)
+    text = factory.Faker("sentence")
+    start = factory.Faker("pyint")
+    end = factory.Faker("pyint")
+
+
+class EntityDetectionResponseFactory(factories.TranscriptCompletedResponseFactory):
+    entities = factory.List([factory.SubFactory(EntityFactory)])
+
+
+def __submit_mock_request(
+    httpx_mock: HTTPXMock,
+    mock_response: Dict[str, Any],
+    config: aai.TranscriptionConfig,
+) -> Tuple[Dict[str, Any], aai.Transcript]:
+    """
+    Helper function to abstract mock transcriber calls with given `TranscriptionConfig`,
+    and perform some common assertions.
+    """
+
+    mock_transcript_id = mock_response.get("id", "mock_id")
+
+    # Mock initial submission response (transcript is processing)
+    mock_processing_response = factories.generate_dict_factory(
+        factories.TranscriptProcessingResponseFactory
+    )()
+
+    httpx_mock.add_response(
+        url=f"{aai.settings.base_url}/transcript",
+        status_code=httpx.codes.OK,
+        method="POST",
+        json={
+            **mock_processing_response,
+            "id": mock_transcript_id,  # inject ID from main mock response
+        },
+    )
+
+    # Mock polling-for-completeness response, with completed transcript
+    httpx_mock.add_response(
+        url=f"{aai.settings.base_url}/transcript/{mock_transcript_id}",
+        status_code=httpx.codes.OK,
+        method="GET",
+        json=mock_response,
+    )
+
+    # == Make API request via SDK ==
+    transcript = aai.Transcriber().transcribe(
+        data="https://example.org/audio.wav",
+        config=config,
+    )
+
+    # Check that submission and polling requests were made
+    assert len(httpx_mock.get_requests()) == 2
+
+    # Extract body of initial submission request
+    request = httpx_mock.get_requests()[0]
+    request_body = json.loads(request.content.decode())
+
+    return request_body, transcript
+
+
+def test_entity_detection_disabled_by_default(httpx_mock: HTTPXMock):
+    """
+    Tests that excluding `entity_detection` from the `TranscriptionConfig` will
+    result in the default behavior of it being excluded from the request body
+    """
+    request_body, transcript = __submit_mock_request(
+        httpx_mock,
+        mock_response=factories.generate_dict_factory(
+            factories.TranscriptCompletedResponseFactory
+        )(),
+        config=aai.TranscriptionConfig(),
+    )
+    assert request_body.get("entity_detection") is None
+    assert transcript.entities is None
+
+
+def test_entity_detection_enabled(httpx_mock: HTTPXMock):
+    """
+    Tests that including `entity_detection=True` in the `TranscriptionConfig`
+    will result in `entity_detection=True` in the request body, and that the
+    response is properly parsed into a `Transcript` object
+    """
+    mock_response = factories.generate_dict_factory(EntityDetectionResponseFactory)()
+    request_body, transcript = __submit_mock_request(
+        httpx_mock,
+        mock_response=mock_response,
+        config=aai.TranscriptionConfig(entity_detection=True),
+    )
+
+    # Check that request body was properly defined
+    assert request_body.get("entity_detection") == True
+
+    # Check that transcript was properly parsed from JSON response
+    assert transcript.error is None
+    assert transcript.entities is not None
+    assert len(transcript.entities) > 0
+    assert len(transcript.entities) == len(mock_response["entities"])
+
+    for entity in transcript.entities:
+        assert len(entity.text.strip()) > 0