Skip to main content

codesynapse_core/
transcribe.rs

1use crate::error::CodeSynapseError;
2use sha2::{Digest, Sha256};
3use std::collections::HashSet;
4use std::env;
5use std::path::{Path, PathBuf};
6use std::process::Command;
7use std::sync::LazyLock;
8
9pub static VIDEO_EXTENSIONS: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
10    [
11        ".mp4", ".mov", ".webm", ".mkv", ".avi", ".m4v", ".mp3", ".wav", ".m4a", ".ogg",
12    ]
13    .into_iter()
14    .collect()
15});
16
17static URL_PREFIXES: &[&str] = &["http://", "https://", "www."];
18static DEFAULT_MODEL: &str = "base";
19static TRANSCRIPTS_DIR: &str = "codesynapse-out/transcripts";
20static FALLBACK_PROMPT: &str = "Use proper punctuation and paragraph breaks.";
21
22fn model_name() -> String {
23    env::var("CODESYNAPSE_WHISPER_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.to_string())
24}
25
26pub fn is_url(path: &str) -> bool {
27    URL_PREFIXES.iter().any(|p| path.starts_with(p))
28}
29
30fn url_hash(url: &str) -> String {
31    let mut hasher = Sha256::new();
32    hasher.update(url.as_bytes());
33    let result = hasher.finalize();
34    format!("{:x}", result)[..12].to_string()
35}
36
37fn find_exe(name: &str) -> Option<PathBuf> {
38    let path_var = env::var_os("PATH")?;
39    for dir in env::split_paths(&path_var) {
40        let candidate = dir.join(name);
41        if candidate.is_file() {
42            return Some(candidate);
43        }
44    }
45    None
46}
47
48pub fn download_audio(url: &str, output_dir: &Path) -> Result<PathBuf, CodeSynapseError> {
49    let hash = url_hash(url);
50
51    for ext in [".m4a", ".opus", ".mp3", ".ogg", ".wav", ".webm"] {
52        let candidate = output_dir.join(format!("yt_{}{}", hash, ext));
53        if candidate.exists() {
54            eprintln!(
55                "  cached audio: {}",
56                candidate.file_name().unwrap_or_default().to_string_lossy()
57            );
58            return Ok(candidate);
59        }
60    }
61
62    std::fs::create_dir_all(output_dir).map_err(CodeSynapseError::Io)?;
63
64    let exe = find_exe("yt-dlp").ok_or_else(|| {
65        CodeSynapseError::Validation(
66            "yt-dlp is required for URL download. Install: pip install yt-dlp".to_string(),
67        )
68    })?;
69
70    let out_template = output_dir.join(format!("yt_{}.%(ext)s", hash));
71    let out_template_str = out_template.to_string_lossy().to_string();
72    let display_url = &url[..url.len().min(80)];
73    eprintln!("  downloading audio: {} ...", display_url);
74
75    let result = Command::new(&exe)
76        .args([
77            "-f",
78            "bestaudio[ext=m4a]/bestaudio/best",
79            "-o",
80            &out_template_str,
81            "--quiet",
82            "--no-warnings",
83            "--no-playlist",
84            url,
85        ])
86        .output()
87        .map_err(CodeSynapseError::Io)?;
88
89    if !result.status.success() {
90        let stderr = String::from_utf8_lossy(&result.stderr);
91        let msg = &stderr[..stderr.len().min(800)];
92        return Err(CodeSynapseError::Validation(format!(
93            "yt-dlp failed: {}",
94            msg
95        )));
96    }
97
98    for ext in [".m4a", ".opus", ".mp3", ".ogg", ".wav", ".webm"] {
99        let candidate = output_dir.join(format!("yt_{}{}", hash, ext));
100        if candidate.exists() {
101            return Ok(candidate);
102        }
103    }
104
105    if let Ok(entries) = std::fs::read_dir(output_dir) {
106        let prefix = format!("yt_{}", hash);
107        for entry in entries.flatten() {
108            let name = entry.file_name().to_string_lossy().to_string();
109            if name.starts_with(&prefix) {
110                return Ok(entry.path());
111            }
112        }
113    }
114
115    Err(CodeSynapseError::Validation(format!(
116        "yt-dlp ran but no output file found for {}",
117        url
118    )))
119}
120
121fn build_whisper_prompt_inner(
122    god_nodes: &[serde_json::Value],
123    env_override: Option<&str>,
124) -> String {
125    if god_nodes.is_empty() {
126        return FALLBACK_PROMPT.to_string();
127    }
128
129    if let Some(ov) = env_override {
130        return ov.to_string();
131    }
132
133    let labels: Vec<&str> = god_nodes
134        .iter()
135        .filter_map(|n| n.get("label").and_then(|v| v.as_str()))
136        .filter(|s| !s.is_empty())
137        .take(5)
138        .collect();
139
140    if labels.is_empty() {
141        return FALLBACK_PROMPT.to_string();
142    }
143
144    let topics = labels.join(", ");
145    format!(
146        "Technical discussion about {}. Use proper punctuation and paragraph breaks.",
147        topics
148    )
149}
150
151pub fn build_whisper_prompt(god_nodes: &[serde_json::Value]) -> String {
152    let env_override = env::var("CODESYNAPSE_WHISPER_PROMPT").ok();
153    build_whisper_prompt_inner(god_nodes, env_override.as_deref())
154}
155
156fn run_whisper(
157    audio_path: &Path,
158    model: &str,
159    initial_prompt: Option<&str>,
160) -> Result<String, CodeSynapseError> {
161    let exe = find_exe("faster-whisper")
162        .or_else(|| find_exe("whisper"))
163        .ok_or_else(|| {
164            CodeSynapseError::Validation(
165                "faster-whisper or whisper CLI required. \
166                 Install: pip install faster-whisper"
167                    .to_string(),
168            )
169        })?;
170
171    let mut cmd = Command::new(&exe);
172    cmd.arg(audio_path.to_string_lossy().as_ref());
173    cmd.args(["--model", model]);
174    if let Some(prompt) = initial_prompt {
175        cmd.args(["--initial_prompt", prompt]);
176    }
177    cmd.args(["--output_format", "txt"]);
178    let out_dir = audio_path.parent().unwrap_or(Path::new("."));
179    cmd.arg("--output_dir").arg(out_dir);
180
181    let result = cmd.output().map_err(CodeSynapseError::Io)?;
182    if !result.status.success() {
183        let stderr = String::from_utf8_lossy(&result.stderr);
184        let msg = &stderr[..stderr.len().min(800)];
185        return Err(CodeSynapseError::Validation(format!(
186            "whisper failed: {}",
187            msg
188        )));
189    }
190
191    let txt_path = out_dir.join(format!(
192        "{}.txt",
193        audio_path.file_stem().unwrap_or_default().to_string_lossy()
194    ));
195    if txt_path.exists() {
196        return std::fs::read_to_string(&txt_path).map_err(CodeSynapseError::Io);
197    }
198
199    Ok(String::from_utf8_lossy(&result.stdout).to_string())
200}
201
202pub fn transcribe_with(
203    video_path: &Path,
204    output_dir: Option<&Path>,
205    initial_prompt: Option<&str>,
206    force: bool,
207    whisper_fn: impl Fn(&Path, &str, Option<&str>) -> Result<String, CodeSynapseError>,
208) -> Result<PathBuf, CodeSynapseError> {
209    let out_dir = output_dir
210        .map(|p| p.to_path_buf())
211        .unwrap_or_else(|| PathBuf::from(TRANSCRIPTS_DIR));
212    std::fs::create_dir_all(&out_dir).map_err(CodeSynapseError::Io)?;
213
214    let audio_path = if is_url(&video_path.to_string_lossy()) {
215        download_audio(&video_path.to_string_lossy(), &out_dir.join("downloads"))?
216    } else {
217        video_path.to_path_buf()
218    };
219
220    let stem = audio_path.file_stem().unwrap_or_default().to_string_lossy();
221    let transcript_path = out_dir.join(format!("{}.txt", stem));
222
223    if transcript_path.exists() && !force {
224        return Ok(transcript_path);
225    }
226
227    let model = model_name();
228    let prompt = initial_prompt.unwrap_or(FALLBACK_PROMPT);
229    eprintln!(
230        "  transcribing {} (model={}) ...",
231        audio_path.file_name().unwrap_or_default().to_string_lossy(),
232        model
233    );
234
235    let text = whisper_fn(&audio_path, &model, Some(prompt))?;
236    let lines: Vec<&str> = text
237        .lines()
238        .map(|l| l.trim())
239        .filter(|l| !l.is_empty())
240        .collect();
241    let transcript = lines.join("\n");
242
243    std::fs::write(&transcript_path, &transcript).map_err(CodeSynapseError::Io)?;
244    eprintln!(
245        "  transcript saved -> {} ({} segments)",
246        transcript_path.display(),
247        lines.len()
248    );
249
250    Ok(transcript_path)
251}
252
253pub fn transcribe(
254    video_path: &Path,
255    output_dir: Option<&Path>,
256    initial_prompt: Option<&str>,
257    force: bool,
258) -> Result<PathBuf, CodeSynapseError> {
259    transcribe_with(video_path, output_dir, initial_prompt, force, run_whisper)
260}
261
262pub fn transcribe_all(
263    video_files: &[String],
264    output_dir: Option<&Path>,
265    initial_prompt: Option<&str>,
266) -> Vec<String> {
267    if video_files.is_empty() {
268        return vec![];
269    }
270
271    video_files
272        .iter()
273        .filter_map(
274            |vf| match transcribe(Path::new(vf), output_dir, initial_prompt, false) {
275                Ok(p) => Some(p.to_string_lossy().to_string()),
276                Err(e) => {
277                    eprintln!("  warning: could not transcribe {}: {}", vf, e);
278                    None
279                }
280            },
281        )
282        .collect()
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288    use std::fs;
289    use tempfile::TempDir;
290
291    fn tmp() -> TempDir {
292        tempfile::tempdir().unwrap()
293    }
294
295    // VIDEO_EXTENSIONS
296
297    #[test]
298    fn test_video_extensions_set() {
299        assert!(VIDEO_EXTENSIONS.contains(".mp4"));
300        assert!(VIDEO_EXTENSIONS.contains(".mp3"));
301        assert!(VIDEO_EXTENSIONS.contains(".wav"));
302        assert!(VIDEO_EXTENSIONS.contains(".mov"));
303        assert!(!VIDEO_EXTENSIONS.contains(".py"));
304        assert!(!VIDEO_EXTENSIONS.contains(".rs"));
305    }
306
307    // build_whisper_prompt
308
309    #[test]
310    fn test_build_whisper_prompt_no_nodes() {
311        let prompt = build_whisper_prompt(&[]);
312        assert!(!prompt.is_empty());
313        assert!(prompt.to_lowercase().contains("punctuation"));
314    }
315
316    #[test]
317    fn test_build_whisper_prompt_env_override() {
318        let nodes = vec![
319            serde_json::json!({"label": "Python"}),
320            serde_json::json!({"label": "FastAPI"}),
321        ];
322        let prompt = build_whisper_prompt_inner(&nodes, Some("Custom domain hint."));
323        assert_eq!(prompt, "Custom domain hint.");
324    }
325
326    #[test]
327    fn test_build_whisper_prompt_returns_topic_string() {
328        let nodes = vec![
329            serde_json::json!({"label": "neural networks"}),
330            serde_json::json!({"label": "transformers"}),
331            serde_json::json!({"label": "attention"}),
332        ];
333        let prompt = build_whisper_prompt_inner(&nodes, None);
334        let lower = prompt.to_lowercase();
335        assert!(lower.contains("neural networks") || lower.contains("transformers"));
336        assert!(lower.contains("punctuation"));
337    }
338
339    #[test]
340    fn test_build_whisper_prompt_nodes_without_labels() {
341        let nodes = vec![
342            serde_json::json!({"id": "1"}),
343            serde_json::json!({"id": "2", "label": ""}),
344        ];
345        let prompt = build_whisper_prompt_inner(&nodes, None);
346        assert!(!prompt.is_empty());
347    }
348
349    // transcribe (via transcribe_with with injectable fn)
350
351    #[test]
352    fn test_transcribe_uses_cache() {
353        let dir = tmp();
354        let video = dir.path().join("lecture.mp4");
355        fs::write(&video, b"fake").unwrap();
356        let out_dir = dir.path().join("transcripts");
357        fs::create_dir_all(&out_dir).unwrap();
358        let cached = out_dir.join("lecture.txt");
359        fs::write(&cached, "Cached transcript content.").unwrap();
360
361        let called = std::cell::Cell::new(false);
362        let result = transcribe_with(&video, Some(&out_dir), None, false, |_, _, _| {
363            called.set(true);
364            Ok("should not be called".to_string())
365        })
366        .unwrap();
367
368        assert_eq!(result, cached);
369        assert!(!called.get(), "whisper should not run when cache hit");
370    }
371
372    #[test]
373    fn test_transcribe_force_reruns() {
374        let dir = tmp();
375        let video = dir.path().join("talk.mp4");
376        fs::write(&video, b"fake").unwrap();
377        let out_dir = dir.path().join("transcripts");
378        fs::create_dir_all(&out_dir).unwrap();
379        fs::write(out_dir.join("talk.txt"), "Old transcript.").unwrap();
380
381        let result = transcribe_with(&video, Some(&out_dir), None, true, |_, _, _| {
382            Ok("New transcript segment.".to_string())
383        })
384        .unwrap();
385
386        assert_eq!(
387            fs::read_to_string(&result).unwrap(),
388            "New transcript segment."
389        );
390    }
391
392    #[test]
393    fn test_transcribe_missing_faster_whisper() {
394        let dir = tmp();
395        let video = dir.path().join("clip.mp4");
396        fs::write(&video, b"fake").unwrap();
397        let out_dir = dir.path().join("out");
398
399        let result = transcribe_with(&video, Some(&out_dir), None, false, |_, _, _| {
400            Err(CodeSynapseError::Validation(
401                "faster-whisper not installed".to_string(),
402            ))
403        });
404
405        assert!(result.is_err());
406    }
407
408    // transcribe_all
409
410    #[test]
411    fn test_transcribe_all_empty() {
412        let results = transcribe_all(&[], None, None);
413        assert!(results.is_empty());
414    }
415
416    #[test]
417    fn test_transcribe_all_uses_cache() {
418        let dir = tmp();
419        let video = dir.path().join("lecture.mp4");
420        fs::write(&video, b"fake").unwrap();
421        let out_dir = dir.path().join("transcripts");
422        fs::create_dir_all(&out_dir).unwrap();
423        let cached = out_dir.join("lecture.txt");
424        fs::write(&cached, "Cached.").unwrap();
425
426        let results = transcribe_all(&[video.to_string_lossy().to_string()], Some(&out_dir), None);
427        assert_eq!(results.len(), 1);
428        assert!(results[0].contains("lecture.txt"));
429    }
430
431    #[test]
432    fn test_transcribe_all_skips_failed() {
433        let dir = tmp();
434        let video = dir.path().join("broken.mp4");
435        fs::write(&video, b"fake").unwrap();
436        let out_dir = dir.path().join("out");
437
438        // No cache exists and transcribe will try real whisper (which won't be found in test env)
439        // transcribe_all catches errors and returns empty
440        let results = transcribe_all(&[video.to_string_lossy().to_string()], Some(&out_dir), None);
441        assert!(results.is_empty());
442    }
443}