diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/google/client.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/google/client.py index 0deba91f..1baabdf8 100644 --- a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/google/client.py +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/google/client.py @@ -10,6 +10,7 @@ from celeste.parameters import ParameterMapper from celeste_speech_generation.client import SpeechGenerationClient from celeste_speech_generation.io import ( + SpeechGenerationFinishReason, SpeechGenerationInput, SpeechGenerationUsage, ) @@ -41,6 +42,13 @@ def _parse_usage(self, response_data: dict[str, Any]) -> SpeechGenerationUsage: usage = super()._parse_usage(response_data) return SpeechGenerationUsage(**usage) + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> SpeechGenerationFinishReason: + """Parse finish reason from response.""" + finish_reason = super()._parse_finish_reason(response_data) + return SpeechGenerationFinishReason(reason=finish_reason.reason) + def _parse_content( self, response_data: dict[str, Any],