Coverage for baibu/chat/speech.py: 100%

1"""Speech to text for voice messages. 

2 

3Voice input is off unless ``CHAT_STT_PROVIDER`` names a subclass of 

4:class:`SpeechToText`. Two are included: 

5 

6- :class:`LiteLLMSpeechToText` sends the audio to any transcription model 

7 LiteLLM supports, including a self-hosted server that speaks the 

8 OpenAI-compatible transcription API. The deployment chooses the model. 

9- :class:`MockSpeechToText` returns made-up transcripts without any model, 

10 for development and tests. 

11 

12A provider returns a :class:`Transcript` or raises :class:`TranscriptionError`. 

13""" 

14 

15import os 

16from dataclasses import dataclass 

17from functools import cache 

18 

19from django.conf import settings 

20from django.utils.module_loading import import_string 

21 

22 

23class TranscriptionError(Exception): 

24 """Transcription failed. ``code`` is safe to show and store.""" 

25 

26 def __init__(self, message: str, *, code: str = "transcription_error"): 

27 super().__init__(message) 

28 self.code = code 

29 

30 

31@dataclass(frozen=True) 

32class Transcript: 

33 text: str 

34 # Language the provider detected or was told, if it says. 

35 language: str = "" 

36 # Length of the audio in seconds, if the provider says. 

37 duration: float | None = None 

38 # Which provider and model produced the text, for the record. 

39 provider: str = "" 

40 

41 

42class SpeechToText: 

43 name = "provider" 

44 

45 def transcribe(self, audio: bytes, *, content_type: str, language_code: str = "") -> Transcript: 

46 raise NotImplementedError 

47 

48 

49class MockSpeechToText(SpeechToText): 

50 """Deterministic transcripts for development and tests. 

51 

52 - Audio containing ``[stt-fail]`` fails. 

53 - Audio containing ``[stt-empty]`` gives an empty transcript. 

54 - Audio containing ``[say:some text]`` is transcribed as ``some text``. 

55 - Anything else is transcribed as :attr:`default_text`. 

56 """ 

57 

58 name = "mock" 

59 default_text = "This is a mock transcript." 

60 

61 def transcribe(self, audio: bytes, *, content_type: str, language_code: str = "") -> Transcript: 

62 if b"[stt-fail]" in audio: 

63 msg = "Mock transcription failure." 

64 raise TranscriptionError(msg, code="mock_failure") 

65 text = self.default_text 

66 if b"[stt-empty]" in audio: 

67 text = "" 

68 elif (start := audio.find(b"[say:")) >= 0 and (end := audio.find(b"]", start)) > start: 

69 text = audio[start + 5 : end].decode("utf-8", errors="replace") 

70 return Transcript(text=text, language=language_code, duration=round(len(audio) / 1000, 3), provider="mock") 

71 

72 

73# File name extensions for the content types browsers record in. Some 

74# transcription servers look at the name to decide how to decode the audio. 

75EXTENSIONS = { 

76 "audio/webm": "webm", 

77 "audio/ogg": "ogg", 

78 "audio/mp4": "m4a", 

79 "audio/x-m4a": "m4a", 

80 "audio/aac": "aac", 

81 "audio/mpeg": "mp3", 

82 "audio/wav": "wav", 

83 "audio/x-wav": "wav", 

84 "audio/wave": "wav", 

85 "audio/flac": "flac", 

86} 

87 

88 

89def extension_for(content_type: str) -> str: 

90 return EXTENSIONS.get(content_type, "bin") 

91 

92 

93class LiteLLMSpeechToText(SpeechToText): 

94 """Transcribes through ``litellm.transcription``. 

95 

96 Settings: ``CHAT_STT_MODEL`` (any transcription model name LiteLLM 

97 accepts), ``CHAT_STT_API_BASE``, ``CHAT_STT_API_KEY_ENV`` (the *name* of 

98 the environment variable holding the key), ``CHAT_STT_TIMEOUT`` and 

99 ``CHAT_STT_PARAMETERS`` (extra arguments for every call, for example 

100 ``{"language": "sw"}``). 

101 """ 

102 

103 name = "litellm" 

104 

105 def transcribe(self, audio: bytes, *, content_type: str, language_code: str = "") -> Transcript: 

106 import litellm 

107 

108 model = settings.CHAT_STT_MODEL 

109 if not model: 

110 msg = "CHAT_STT_MODEL is not set." 

111 raise TranscriptionError(msg, code="not_configured") 

112 kwargs = { 

113 **settings.CHAT_STT_PARAMETERS, 

114 "model": model, 

115 "file": (f"audio.{extension_for(content_type)}", audio, content_type), 

116 "timeout": settings.CHAT_STT_TIMEOUT, 

117 } 

118 if settings.CHAT_STT_API_BASE: 

119 kwargs["api_base"] = settings.CHAT_STT_API_BASE 

120 if settings.CHAT_STT_API_KEY_ENV and (key := os.environ.get(settings.CHAT_STT_API_KEY_ENV, "")): 

121 kwargs["api_key"] = key 

122 try: 

123 response = litellm.transcription(**kwargs) 

124 except Exception as exc: 

125 code = "timeout" if "timeout" in type(exc).__name__.lower() else "transcription_error" 

126 msg = f"Transcription failed: {type(exc).__name__}" 

127 raise TranscriptionError(msg, code=code) from exc 

128 text = getattr(response, "text", None) 

129 if not isinstance(text, str): 

130 msg = "Transcription returned no text." 

131 raise TranscriptionError(msg, code="empty_response") 

132 duration = getattr(response, "duration", None) 

133 language = getattr(response, "language", None) 

134 return Transcript( 

135 text=text, 

136 language=language if isinstance(language, str) else "", 

137 duration=float(duration) if isinstance(duration, int | float) else None, 

138 provider=f"{self.name}:{model}"[:128], 

139 ) 

140 

141 

142@cache 

143def _provider_class(path: str) -> type[SpeechToText]: 

144 return import_string(path) 

145 

146 

147def voice_enabled() -> bool: 

148 return bool(settings.CHAT_STT_PROVIDER) 

149 

150 

151def get_provider() -> SpeechToText | None: 

152 path = settings.CHAT_STT_PROVIDER 

153 return _provider_class(path)() if path else None