Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/multimodal/loaders/audio.py: 20%

1"""Audio loader for multi-modal RAG. 

2 

3This module provides loading capabilities for audio documents with support for: 

4- Multiple formats: MP3, WAV, FLAC, OGG, AAC, M4A, WMA 

5- Audio metadata extraction (ID3 tags, etc.) 

6- Optional speech-to-text transcription 

7- Audio validation and normalization 

8""" 

9 

10from __future__ import annotations 

11 

12import asyncio 

13from pathlib import Path 

14from typing import Any 

15 

16from lexigram.logging import ( 

17 get_logger, 

18) 

19 

20logger = get_logger(__name__) 

21 

22try: 

23 import librosa 

24 

25 # import soundfile as sf 

26 

27 LIBROSA_AVAILABLE = True 

28except ImportError: 

29 LIBROSA_AVAILABLE = False 

30 

31try: 

32 import mutagen 

33 

34 MUTAGEN_AVAILABLE = True 

35except ImportError: 

36 MUTAGEN_AVAILABLE = False 

37 

38import contextlib 

39 

40from lexigram.ai.rag.exceptions import AudioLoaderError 

41from lexigram.ai.rag.multimodal.types import ( 

42 AudioDocument, 

43 AudioFormat, 

44 AudioMetadata, 

45) 

46 

47 

48class AudioLoader: 

49 """Loader for audio documents. 

50 

51 Supports loading audio from files or bytes with optional: 

52 - Metadata extraction (ID3, etc.) 

53 - Speech-to-text transcription 

54 - Format validation 

55 - Audio normalization 

56 

57 Args: 

58 extract_metadata: Whether to extract audio metadata 

59 transcribe: Whether to transcribe audio to text 

60 sample_rate: Target sample rate (None to keep original) 

61 mono: Convert to mono audio 

62 normalize: Normalize audio volume 

63 

64 Example: 

65 >>> loader = AudioLoader(transcribe=True) 

66 >>> doc = await loader.load("/path/to/audio.mp3") 

67 >>> print(doc.transcript) 

68 "Transcribed speech" 

69 """ 

70 

71 def __init__( 

72 self, 

73 extract_metadata: bool = True, 

74 transcribe: bool = False, 

75 sample_rate: int | None = None, 

76 mono: bool = False, 

77 normalize: bool = False, 

78 ): 

79 """Initialize audio loader.""" 

80 if not LIBROSA_AVAILABLE: 

81 msg = "librosa and soundfile are required for audio loading. Install with: pip install librosa soundfile" 

82 raise ImportError(msg) 

83 

84 self.extract_metadata = extract_metadata 

85 self.transcribe = transcribe 

86 self.sample_rate = sample_rate 

87 self.mono = mono 

88 self.normalize = normalize 

89 

90 # Lazy load transcription model 

91 self._transcription_model = None 

92 

93 async def load( 

94 self, 

95 source: str | Path | bytes, 

96 title: str | None = None, 

97 artist: str | None = None, 

98 **metadata_kwargs: Any, 

99 ) -> AudioDocument: 

100 """Load an audio document. 

101 

102 Args: 

103 source: File path or raw bytes 

104 title: Optional audio title 

105 artist: Optional artist name 

106 **metadata_kwargs: Additional metadata fields 

107 

108 Returns: 

109 AudioDocument with loaded audio data 

110 

111 Raises: 

112 AudioLoaderError: If loading fails 

113 """ 

114 try: 

115 # Load audio 

116 content: bytes | str 

117 if isinstance(source, (str, Path)): 

118 file_path = Path(source) 

119 exists = await asyncio.to_thread(file_path.exists) 

120 if not exists: 

121 msg = f"Audio file not found: {source}" 

122 raise AudioLoaderError(msg) 

123 

124 # Load with librosa (Blocking I/O and CPU) 

125 def _load_librosa() -> Any: 

126 return librosa.load( 

127 str(file_path), 

128 sr=self.sample_rate, 

129 mono=self.mono, 

130 ) 

131 

132 y, sr = await asyncio.to_thread(_load_librosa) 

133 content = str(file_path) 

134 

135 # Detect format from file extension 

136 audio_format = self._detect_format(file_path) 

137 

138 elif isinstance(source, bytes): 

139 # Load from bytes 

140 import io 

141 

142 def _load_librosa_bytes() -> Any: 

143 return librosa.load( 

144 io.BytesIO(source), 

145 sr=self.sample_rate, 

146 mono=self.mono, 

147 ) 

148 

149 y, sr = await asyncio.to_thread(_load_librosa_bytes) 

150 content = source 

151 file_path = None 

152 audio_format = AudioFormat.WAV # Default for bytes 

153 

154 else: 

155 msg = f"Unsupported source type: {type(source)}" 

156 raise AudioLoaderError(msg) 

157 

158 # Normalize if requested (CPU intensive) 

159 if self.normalize: 

160 y = await asyncio.to_thread(librosa.util.normalize, y) 

161 

162 # Get audio properties (CPU intensive) 

163 duration = await asyncio.to_thread(librosa.get_duration, y=y, sr=sr) 

164 channels = 1 if self.mono or y.ndim == 1 else y.shape[0] 

165 

166 # Extract metadata (Blocking I/O) 

167 id3_data = {} 

168 if self.extract_metadata and file_path and MUTAGEN_AVAILABLE: 

169 id3_data = await asyncio.to_thread(self._extract_metadata, file_path) 

170 

171 # Build metadata 

172 metadata = AudioMetadata( 

173 title=title or id3_data.get("title"), 

174 artist=artist or id3_data.get("artist"), 

175 album=id3_data.get("album"), 

176 genre=id3_data.get("genre"), 

177 year=id3_data.get("year"), 

178 id3=id3_data, 

179 **metadata_kwargs, 

180 ) 

181 

182 # Transcribe if requested (Already async/uses threads internally) 

183 transcript = None 

184 if self.transcribe: 

185 transcript = await self._transcribe_audio(y, sr) 

186 

187 # Create document 

188 return AudioDocument( 

189 content=content, 

190 format=audio_format, 

191 duration=duration, 

192 sample_rate=sr, 

193 channels=channels, 

194 metadata=metadata, 

195 transcript=transcript, 

196 file_path=file_path, 

197 ) 

198 

199 except Exception as e: 

200 msg = f"Failed to load audio: {e}" 

201 raise AudioLoaderError(msg) from e 

202 

203 async def load_batch( 

204 self, 

205 sources: list[str | Path | bytes], 

206 **kwargs: Any, 

207 ) -> list[AudioDocument]: 

208 """Load multiple audio files. 

209 

210 Args: 

211 sources: List of audio sources 

212 **kwargs: Common metadata for all audio files 

213 

214 Returns: 

215 List of loaded AudioDocuments 

216 """ 

217 if not sources: 

218 return [] 

219 

220 async def _safe_load(source) -> Any: 

221 try: 

222 return await self.load(source, **kwargs) 

223 except AudioLoaderError as e: 

224 logger.warning("Failed to load audio %r: %s", source, e) 

225 return None 

226 

227 tasks = [_safe_load(source) for source in sources] 

228 results = await asyncio.gather(*tasks) 

229 return [doc for doc in results if doc is not None] 

230 

231 def _detect_format(self, file_path: Path) -> AudioFormat: 

232 """Detect audio format from file extension. 

233 

234 Args: 

235 file_path: Path to audio file 

236 

237 Returns: 

238 AudioFormat enum value 

239 """ 

240 suffix = file_path.suffix.lower().lstrip(".") 

241 

242 format_mapping = { 

243 "mp3": AudioFormat.MP3, 

244 "wav": AudioFormat.WAV, 

245 "flac": AudioFormat.FLAC, 

246 "ogg": AudioFormat.OGG, 

247 "aac": AudioFormat.AAC, 

248 "m4a": AudioFormat.M4A, 

249 "wma": AudioFormat.WMA, 

250 } 

251 

252 if suffix not in format_mapping: 

253 msg = f"Unsupported audio format: {suffix}" 

254 raise AudioLoaderError(msg) 

255 

256 return format_mapping[suffix] 

257 

258 def _extract_metadata(self, file_path: Path) -> dict[str, Any]: 

259 """Extract metadata from audio file. 

260 

261 Args: 

262 file_path: Path to audio file 

263 

264 Returns: 

265 Dictionary of metadata 

266 """ 

267 metadata: dict[str, Any] = {} 

268 

269 try: 

270 audio_file = mutagen.File(str(file_path)) 

271 if audio_file is not None: 

272 # Extract common tags 

273 if "title" in audio_file: 

274 metadata["title"] = str(audio_file["title"][0]) 

275 if "artist" in audio_file: 

276 metadata["artist"] = str(audio_file["artist"][0]) 

277 if "album" in audio_file: 

278 metadata["album"] = str(audio_file["album"][0]) 

279 if "genre" in audio_file: 

280 metadata["genre"] = str(audio_file["genre"][0]) 

281 if "date" in audio_file: 

282 with contextlib.suppress(ValueError, IndexError): 

283 metadata["year"] = int(str(audio_file["date"][0])[:4]) 

284 

285 # Store all tags 

286 metadata["all_tags"] = {k: str(v) for k, v in audio_file.items()} 

287 

288 except (ValueError, TypeError, OSError, KeyError) as e: 

289 logger.warning("Failed to extract metadata: %s", e) 

290 

291 return metadata 

292 

293 async def _transcribe_audio(self, audio: Any, sample_rate: int) -> str: 

294 """Transcribe audio to text using OpenAI Whisper. 

295 

296 Args: 

297 audio: Audio data array 

298 sample_rate: Sample rate 

299 

300 Returns: 

301 Transcribed text 

302 """ 

303 try: 

304 import whisper 

305 except ImportError: 

306 logger.warning("Whisper not installed. Skipping transcription.") 

307 return "" 

308 

309 if self._transcription_model is None: 

310 # Load small model by default 

311 def _load() -> Any: 

312 return whisper.load_model("base") 

313 

314 self._transcription_model = await asyncio.to_thread(_load) # type: ignore[func-returns-value] 

315 

316 # Transcribe 

317 model = self._transcription_model 

318 assert model is not None, "transcription model not loaded" # noqa: S101 

319 

320 def _transcribe() -> Any: 

321 result = model.transcribe(audio, fp16=False) 

322 return result.get("text", "").strip() 

323 

324 return await asyncio.to_thread(_transcribe) 

325 

326 def supports_transcription(self) -> bool: 

327 """Check if transcription is available. 

328 

329 Returns: 

330 True if transcription model is loaded 

331 """ 

332 return self._transcription_model is not None