返回 douyin-downloader
transcript_manager.py
根目录 / core / transcript_manager.py
1 import json
2 import os
3 import tempfile
4 from pathlib import Path
5 from typing import Any, Dict, List, Optional, Tuple
6
7 import aiofiles
8 import aiohttp
9
10 from config import ConfigLoader
11 from core.audio_extraction import AudioExtractError, extract_audio
12 from storage import Database, FileManager
13 from utils.logger import setup_logger
14
15 logger = setup_logger("TranscriptManager")
16
17
18 # File extensions that the transcription endpoint already accepts as audio.
19 # When the source download is one of these we skip ``extract_audio`` and
20 # upload the file as-is. Lower-case keys.
21 _SOURCE_AUDIO_MIME = {
22 ".m4a": "audio/mp4",
23 ".mp3": "audio/mpeg",
24 ".wav": "audio/wav",
25 ".aac": "audio/aac",
26 ".opus": "audio/ogg",
27 ".flac": "audio/flac",
28 ".ogg": "audio/ogg",
29 }
30
31
32 def _mask_api_key_local(value: str) -> str:
33 """Pure mirror of ``server.app._mask_api_key`` for use inside the
34 shared transcript pipeline (which can't import from desktop-only
35 code). Same boundary semantics: empty → ``""``, 1-7 → all ``*``,
36 >=8 → ``"<first 4>...<last 4>"``.
37
38 Used to redact bearer tokens that might be echoed back in upstream
39 error responses before they land in ``transcript_jobs.error_message``
40 (Property 1 / 2).
41 """
42 if not value:
43 return ""
44 n = len(value)
45 if n >= 8:
46 return f"{value[:4]}...{value[-4:]}"
47 return "*" * n
48
49
50 def resolve_api_key_with_source(
51 transcript_cfg: Dict[str, Any],
52 ) -> Tuple[str, str]:
53 """Pure helper that resolves a transcription API key and reports
54 where it came from.
55
56 Used by both :class:`TranscriptManager` (during a real
57 ``process_video`` call) and the desktop sidecar's
58 ``POST /api/v1/transcript/test-connectivity`` endpoint, so the two
59 code paths can never disagree on which credential they're using.
60
61 Priority (first non-empty after strip wins):
62 1. The environment variable named by ``api_key_env``
63 (default ``OPENAI_API_KEY``).
64 2. The ``api_key`` field persisted in ``settings.yml``.
65
66 Returns:
67 Tuple of (api_key, source) where ``source`` is one of
68 ``"env"``, ``"settings"``, or ``"none"``.
69 """
70 api_key_env = str(transcript_cfg.get("api_key_env", "OPENAI_API_KEY") or "").strip()
71 if api_key_env:
72 env_value = os.getenv(api_key_env, "").strip()
73 if env_value:
74 return env_value, "env"
75
76 settings_value = str(transcript_cfg.get("api_key", "") or "").strip()
77 if settings_value:
78 return settings_value, "settings"
79 return "", "none"
80
81
82 class TranscriptManager:
83 def __init__(
84 self,
85 config: ConfigLoader,
86 file_manager: FileManager,
87 database: Optional[Database] = None,
88 ):
89 self.config = config
90 self.file_manager = file_manager
91 self.database = database
92
93 def _cfg(self) -> Dict[str, Any]:
94 return self.config.get("transcript", {}) or {}
95
96 def _enabled(self) -> bool:
97 return bool(self._cfg().get("enabled", False))
98
99 def _model(self) -> str:
100 return str(self._cfg().get("model", "gpt-4o-mini-transcribe")).strip()
101
102 def _upload_audio_only(self) -> bool:
103 """``transcript.upload_audio_only`` flag (R1.14, default ``True``).
104
105 Hidden from the Settings UI by design (R1.18); editable only via
106 ``settings.yml`` or a direct ``PATCH /api/v1/settings`` call so a
107 user wandering through the UI can't accidentally disable the
108 bandwidth-saving path.
109 """
110 v = self._cfg().get("upload_audio_only", True)
111 if v is None:
112 return True
113 return bool(v)
114
115 def _response_formats(self) -> List[str]:
116 formats = self._cfg().get("response_formats", ["txt", "json"])
117 if not isinstance(formats, list):
118 return ["txt", "json"]
119 normalized = [str(item).strip().lower() for item in formats if str(item).strip()]
120 return normalized or ["txt", "json"]
121
122 def _resolve_api_key(self) -> str:
123 """Resolve the API key per Requirement 5.6.
124
125 Priority (first non-empty after strip wins):
126 1. The environment variable named by ``transcript.api_key_env``
127 (default ``OPENAI_API_KEY``).
128 2. The ``transcript.api_key`` field persisted in ``settings.yml``.
129 Falling through both returns ``""`` and the caller goes through the
130 existing ``skip_reason="missing_api_key"`` branch.
131 """
132 api_key, _source = resolve_api_key_with_source(self._cfg())
133 return api_key
134
135 def _api_url(self) -> str:
136 api_url = str(
137 self._cfg().get("api_url", "https://api.openai.com/v1/audio/transcriptions")
138 ).strip()
139 return api_url or "https://api.openai.com/v1/audio/transcriptions"
140
141 def resolve_output_dir(self, video_path: Path) -> Path:
142 video_path = Path(video_path)
143 video_dir = video_path.parent
144 output_dir = str(self._cfg().get("output_dir", "")).strip()
145 if not output_dir:
146 return video_dir
147
148 output_root = Path(output_dir)
149 try:
150 relative_dir = video_dir.resolve().relative_to(self.file_manager.base_path.resolve())
151 return output_root / relative_dir
152 except Exception:
153 logger.warning(
154 "Failed to mirror transcript path for video %s, fallback to video dir",
155 video_path,
156 )
157 return video_dir
158
159 def build_output_paths(self, video_path: Path) -> Tuple[Path, Path]:
160 video_path = Path(video_path)
161 output_dir = self.resolve_output_dir(video_path)
162 output_dir.mkdir(parents=True, exist_ok=True)
163 stem = video_path.stem
164 return (
165 output_dir / f"{stem}.transcript.txt",
166 output_dir / f"{stem}.transcript.json",
167 )
168
169 async def process_video(self, video_path: Path, aweme_id: str) -> Dict[str, Any]:
170 video_path = Path(video_path)
171
172 if not self._enabled():
173 return {"status": "skipped", "reason": "disabled"}
174
175 api_key = self._resolve_api_key()
176 text_path, json_path = self.build_output_paths(video_path)
177 model = self._model()
178
179 if not api_key:
180 await self._record_job(
181 aweme_id=aweme_id,
182 video_path=video_path,
183 transcript_dir=text_path.parent,
184 text_path=text_path,
185 json_path=json_path,
186 model=model,
187 status="skipped",
188 skip_reason="missing_api_key",
189 error_message=None,
190 )
191 logger.warning("Transcript skipped for aweme %s: missing_api_key", aweme_id)
192 return {"status": "skipped", "reason": "missing_api_key"}
193
194 # ------------------------------------------------------------------
195 # Pick what to upload:
196 # 1. Source already audio (m4a/mp3/...): pass through (R1.8).
197 # 2. upload_audio_only=true (default): extract audio first (R1.1).
198 # 3. upload_audio_only=false: legacy behaviour, upload the video
199 # file itself (R1.16 / R6.5).
200 # ------------------------------------------------------------------
201 source_ext = video_path.suffix.lower()
202 is_source_audio = source_ext in _SOURCE_AUDIO_MIME
203 tmp_audio_dir: Optional[tempfile.TemporaryDirectory] = None
204
205 upload_path = video_path
206 upload_filename = video_path.name
207 upload_content_type = self._guess_video_content_type(video_path)
208
209 try:
210 if not is_source_audio and self._upload_audio_only():
211 tmp_audio_dir = tempfile.TemporaryDirectory(
212 prefix="transcript_audio_"
213 )
214 try:
215 upload_path = await extract_audio(
216 video_path, Path(tmp_audio_dir.name)
217 )
218 except AudioExtractError as exc:
219 error_message = str(exc)
220 await self._record_job(
221 aweme_id=aweme_id,
222 video_path=video_path,
223 transcript_dir=text_path.parent,
224 text_path=text_path,
225 json_path=json_path,
226 model=model,
227 status="failed",
228 skip_reason=None,
229 error_message=error_message,
230 )
231 logger.error(
232 "Transcript audio extraction failed for aweme %s: %s",
233 aweme_id,
234 error_message,
235 )
236 return {
237 "status": "failed",
238 "reason": "audio_extract_failed",
239 "error": error_message,
240 }
241 upload_filename = f"{video_path.stem}.mp3"
242 upload_content_type = "audio/mpeg"
243 elif is_source_audio:
244 upload_filename = video_path.name
245 upload_content_type = _SOURCE_AUDIO_MIME[source_ext]
246
247 try:
248 payload = await self._call_openai_transcription(
249 api_key=api_key,
250 file_path=upload_path,
251 filename=upload_filename,
252 content_type=upload_content_type,
253 model=model,
254 )
255 # ``_write_outputs`` re-derives the text from ``payload`` —
256 # no need to pre-extract it here.
257 await self._write_outputs(payload, text_path, json_path)
258 await self._record_job(
259 aweme_id=aweme_id,
260 video_path=video_path,
261 transcript_dir=text_path.parent,
262 text_path=text_path,
263 json_path=json_path,
264 model=model,
265 status="success",
266 skip_reason=None,
267 error_message=None,
268 )
269 return {
270 "status": "success",
271 "text_path": str(text_path),
272 "json_path": str(json_path),
273 }
274 except Exception as exc:
275 error_message = str(exc)
276 await self._record_job(
277 aweme_id=aweme_id,
278 video_path=video_path,
279 transcript_dir=text_path.parent,
280 text_path=text_path,
281 json_path=json_path,
282 model=model,
283 status="failed",
284 skip_reason=None,
285 error_message=error_message,
286 )
287 logger.error(
288 "Transcript failed for aweme %s: %s", aweme_id, error_message
289 )
290 return {
291 "status": "failed",
292 "reason": "transcription_error",
293 "error": error_message,
294 }
295 finally:
296 if tmp_audio_dir is not None:
297 # Cleanup is best-effort. R6.7: a cleanup error must not
298 # surface as a transcript task failure — log a WARNING and
299 # let the surrounding return path run.
300 try:
301 tmp_audio_dir.cleanup()
302 except Exception as exc: # noqa: BLE001 — broad is correct here
303 logger.warning(
304 "Failed to clean up transcript audio temp dir %s: %r",
305 tmp_audio_dir.name,
306 exc,
307 )
308
309 async def _write_outputs(
310 self, payload: Dict[str, Any], text_path: Path, json_path: Path
311 ) -> None:
312 formats = set(self._response_formats())
313
314 if "txt" in formats:
315 text = str(payload.get("text", "")).strip()
316 async with aiofiles.open(text_path, "w", encoding="utf-8") as f:
317 await f.write(text)
318
319 if "json" in formats:
320 async with aiofiles.open(json_path, "w", encoding="utf-8") as f:
321 await f.write(json.dumps(payload, ensure_ascii=False, indent=2))
322
323 async def _call_openai_transcription(
324 self,
325 *,
326 api_key: str,
327 file_path: Path,
328 filename: str,
329 content_type: str,
330 model: str,
331 ) -> Dict[str, Any]:
332 """POST a multipart transcription request.
333
334 ``file_path`` is whatever the caller decided to upload — could be
335 the original video, the source audio file (passthrough), or the
336 ffmpeg-extracted mp3. The caller passes the appropriate
337 ``filename`` + ``content_type`` so the multipart body advertises
338 the right MIME.
339 """
340 if not file_path.exists():
341 raise FileNotFoundError(f"Upload file not found: {file_path}")
342
343 transcript_cfg = self._cfg()
344 language_hint = str(transcript_cfg.get("language_hint", "")).strip()
345 api_url = self._api_url()
346
347 form = aiohttp.FormData()
348 form.add_field("model", model)
349 form.add_field("response_format", "json")
350 if language_hint:
351 form.add_field("language", language_hint)
352
353 with file_path.open("rb") as f:
354 form.add_field(
355 "file",
356 f,
357 filename=filename,
358 content_type=content_type,
359 )
360 timeout = aiohttp.ClientTimeout(total=600)
361 async with aiohttp.ClientSession(timeout=timeout) as session:
362 async with session.post(
363 api_url,
364 data=form,
365 headers={"Authorization": f"Bearer {api_key}"},
366 ) as response:
367 if response.status != 200:
368 body = await response.text()
369 # Some misbehaving proxies echo the bearer token
370 # into 4xx error pages; redact before the body
371 # ends up in ``transcript_jobs.error_message``
372 # (Property 1 / 2).
373 if api_key and api_key in body:
374 body = body.replace(api_key, _mask_api_key_local(api_key))
375 raise RuntimeError(
376 f"OpenAI transcription failed: status={response.status}, body={body}"
377 )
378
379 payload = await response.json(content_type=None)
380 if not isinstance(payload, dict):
381 raise RuntimeError("OpenAI transcription returned invalid payload")
382 return payload
383
384 @staticmethod
385 def _guess_video_content_type(video_path: Path) -> str:
386 suffix = video_path.suffix.lower()
387 if suffix == ".mp4":
388 return "video/mp4"
389 if suffix == ".m4a":
390 return "audio/mp4"
391 if suffix == ".wav":
392 return "audio/wav"
393 if suffix == ".mp3":
394 return "audio/mpeg"
395 return "application/octet-stream"
396
397 async def _record_job(
398 self,
399 *,
400 aweme_id: str,
401 video_path: Path,
402 transcript_dir: Path,
403 text_path: Path,
404 json_path: Path,
405 model: str,
406 status: str,
407 skip_reason: Optional[str],
408 error_message: Optional[str],
409 ) -> None:
410 if not self.database:
411 return
412
413 await self.database.upsert_transcript_job(
414 {
415 "aweme_id": aweme_id,
416 "video_path": str(video_path),
417 "transcript_dir": str(transcript_dir),
418 "text_path": str(text_path),
419 "json_path": str(json_path),
420 "model": model,
421 "status": status,
422 "skip_reason": skip_reason,
423 "error_message": error_message,
424 }
425 )
426
426 lines PYTHON