diff --git a/assemblyai/types.py b/assemblyai/types.py index 92377f6..111e0c7 100644 --- a/assemblyai/types.py +++ b/assemblyai/types.py @@ -2161,16 +2161,16 @@ class Word(BaseModel): channel: Optional[str] = None # This is a workaround to address an issue where sentiment_analysis_results - # may return contains sentiments where `start` is null. + # may return contains sentiments where `start` or `end` is null. if pydantic_v2: - @field_validator("start", mode="before") + @field_validator("start", "end", mode="before") def set_start_default(cls, v): return 0 if v is None else v else: - @validator("start", pre=True) + @validator("start", "end", pre=True) def set_start_default(cls, v): return 0 if v is None else v diff --git a/tests/unit/test_sentiment_analysis.py b/tests/unit/test_sentiment_analysis.py index 2d86c01..125eb12 100644 --- a/tests/unit/test_sentiment_analysis.py +++ b/tests/unit/test_sentiment_analysis.py @@ -106,3 +106,34 @@ def test_sentiment_analysis_null_start(httpx_mock: HTTPXMock): transcript.sentiment_analysis, ): assert transcript_sentiment_result.start == 0 + + +def test_sentiment_analysis_null_end(httpx_mock: HTTPXMock): + """ + Tests that `end` converts null values to 0. + """ + mock_response = { + "audio_url": "https://example/audio.mp3", + "status": "completed", + "sentiment_analysis_results": [ + { + "text": "hi", + "start": 100, + "end": None, + "confidence": 0.99, + "sentiment": "POSITIVE", + } + ], + } + request_body, transcript = unit_test_utils.submit_mock_transcription_request( + httpx_mock, + mock_response=mock_response, + config=aai.TranscriptionConfig(sentiment_analysis=True), + ) + + for response_sentiment_result, transcript_sentiment_result in zip( + mock_response["sentiment_analysis_results"], + transcript.sentiment_analysis, + ): + assert transcript_sentiment_result.start == 100 + assert transcript_sentiment_result.end == 0