synaptic-loaders 0.4.0

Document loaders: Text, JSON, CSV, Markdown, Directory, Web, Notion, GitHub, YouTube, arXiv
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use async_trait::async_trait;
use serde_json::Value;
use std::collections::HashMap;
use synaptic_core::{Document, Loader, SynapticError};

/// Loader for YouTube video transcripts.
pub struct YoutubeLoader {
    client: reqwest::Client,
    video_ids: Vec<String>,
    language: String,
}

impl YoutubeLoader {
    pub fn new(video_ids: Vec<String>) -> Self {
        Self {
            client: reqwest::Client::new(),
            video_ids,
            language: "en".to_string(),
        }
    }

    pub fn with_language(mut self, lang: impl Into<String>) -> Self {
        self.language = lang.into();
        self
    }

    async fn fetch_transcript(&self, video_id: &str) -> Result<String, SynapticError> {
        let url = format!(
            "https://www.youtube.com/api/timedtext?v={}&lang={}&fmt=json3",
            video_id, self.language
        );
        let resp = self
            .client
            .get(&url)
            .send()
            .await
            .map_err(|e| SynapticError::Loader(format!("YouTube fetch: {e}")))?;
        let body: Value = resp
            .json()
            .await
            .map_err(|e| SynapticError::Loader(format!("YouTube parse: {e}")))?;

        let text = body["events"]
            .as_array()
            .map(|events| {
                events
                    .iter()
                    .filter_map(|event| {
                        event["segs"].as_array().map(|segs| {
                            segs.iter()
                                .filter_map(|seg| seg["utf8"].as_str())
                                .collect::<Vec<_>>()
                                .join("")
                        })
                    })
                    .filter(|s| !s.trim().is_empty())
                    .collect::<Vec<_>>()
                    .join(" ")
            })
            .unwrap_or_default();

        Ok(text)
    }

    async fn fetch_title(&self, video_id: &str) -> Option<String> {
        let url = format!(
            "https://www.youtube.com/oembed?url=https://www.youtube.com/watch?v={}&format=json",
            video_id
        );
        self.client
            .get(&url)
            .send()
            .await
            .ok()?
            .json::<Value>()
            .await
            .ok()?["title"]
            .as_str()
            .map(|s| s.to_string())
    }
}

#[async_trait]
impl Loader for YoutubeLoader {
    async fn load(&self) -> Result<Vec<Document>, SynapticError> {
        let mut documents = Vec::new();
        for video_id in &self.video_ids {
            let content = self.fetch_transcript(video_id).await?;
            if content.is_empty() {
                continue;
            }
            let title = self.fetch_title(video_id).await;
            let mut metadata = HashMap::new();
            metadata.insert(
                "source".to_string(),
                Value::String(format!("youtube:{}", video_id)),
            );
            metadata.insert(
                "url".to_string(),
                Value::String(format!("https://www.youtube.com/watch?v={}", video_id)),
            );
            if let Some(t) = title {
                metadata.insert("title".to_string(), Value::String(t));
            }
            documents.push(Document {
                id: video_id.clone(),
                content,
                metadata,
            });
        }
        Ok(documents)
    }
}