ingestion.py

import os
import re
import json
import sqlite3
from typing import List, Dict, Any
import numpy as np
from youtube_transcript_api import YouTubeTranscriptApi
import googleapiclient.discovery
import openai
import argparse
from dotenv import load_dotenv

load_dotenv()

class YouTubeEmbeddingIngestion:
    def __init__(self, db_path: str = "youtube_embeddings.db"):
        youtube_api_key = os.getenv('YOUTUBE_API_KEY')
        openai_api_key = os.getenv('OPENAI_API_KEY')
        
        if not youtube_api_key:
            raise ValueError("YOUTUBE_API_KEY not found in environment variables")
        if not openai_api_key:
            raise ValueError("OPENAI_API_KEY not found in environment variables")
        
        self.youtube = googleapiclient.discovery.build("youtube", "v3", developerKey=youtube_api_key)
        self.client = openai.OpenAI(api_key=openai_api_key)
        self.ytt_api = YouTubeTranscriptApi()
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS video_chunks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                video_id TEXT NOT NULL,
                video_title TEXT NOT NULL,
                chunk_text TEXT NOT NULL,
                chunk_index INTEGER NOT NULL,
                start_time REAL,
                end_time REAL,
                embedding BLOB NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_video_id ON video_chunks(video_id);
        ''')
        
        conn.commit()
        conn.close()
    
    def extract_video_id(self, url: str) -> str:
        patterns = [
            r'(?:youtube\.com\/watch\?v=|youtu\.be\/|youtube\.com\/embed\/)([^&\n?#]+)',
            r'youtube\.com\/playlist\?list=([^&\n?#]+)'
        ]
        
        for pattern in patterns:
            match = re.search(pattern, url)
            if match:
                return match.group(1)
        
        return url.strip()
    
    def get_playlist_videos(self, playlist_id: str) -> List[Dict[str, str]]:
        videos = []
        next_page_token = None
        
        while True:
            request = self.youtube.playlistItems().list(
                part="snippet",
                playlistId=playlist_id,
                maxResults=50,
                pageToken=next_page_token
            )
            response = request.execute()
            
            for item in response['items']:
                videos.append({
                    'id': item['snippet']['resourceId']['videoId'],
                    'title': item['snippet']['title']
                })
            
            next_page_token = response.get('nextPageToken')
            if not next_page_token:
                break
        
        return videos
    
    def get_video_info(self, video_id: str) -> Dict[str, str]:
        request = self.youtube.videos().list(
            part="snippet",
            id=video_id
        )
        response = request.execute()
        
        if response['items']:
            return {
                'id': video_id,
                'title': response['items'][0]['snippet']['title']
            }
        return {'id': video_id, 'title': 'Unknown Title'}
    
    def get_transcript(self, video_id: str) -> List[Dict[str, Any]]:
        try:
            transcript_list = self.ytt_api.list(video_id)
            
            try:
                transcript = transcript_list.find_manually_created_transcript(['en'])
                fetched_transcript = transcript.fetch()
                print(f"Found manual English transcript for {video_id}")
                return fetched_transcript.to_raw_data()
            except:
                pass
            
            try:
                transcript = transcript_list.find_generated_transcript(['en'])
                fetched_transcript = transcript.fetch()
                print(f"Found auto-generated English transcript for {video_id}")
                return fetched_transcript.to_raw_data()
            except:
                pass
            
            try:
                transcript = transcript_list.find_transcript(['en', 'es'])
                fetched_transcript = transcript.fetch()
                print(f"Found transcript in {transcript.language_code} for {video_id}")
                return fetched_transcript.to_raw_data()
            except:
                pass
            
        except Exception as e:
            print(f"Error getting transcript for {video_id}: {e}")
        
        return []
    
    def chunk_transcript(self, transcript: List[Dict[str, Any]], chunk_size: int = 500) -> List[Dict[str, Any]]:
        chunks = []
        current_chunk = ""
        current_start = None
        current_end = None
        
        for entry in transcript:
            if current_start is None:
                current_start = entry['start']
            
            if len(current_chunk) + len(entry['text']) > chunk_size and current_chunk:
                chunks.append({
                    'text': current_chunk.strip(),
                    'start_time': current_start,
                    'end_time': current_end
                })
                current_chunk = entry['text'] + " "
                current_start = entry['start']
            else:
                current_chunk += entry['text'] + " "
            
            current_end = entry['start'] + entry['duration']
        
        if current_chunk.strip():
            chunks.append({
                'text': current_chunk.strip(),
                'start_time': current_start,
                'end_time': current_end
            })
        
        return chunks
    
    def store_embeddings(self, video_info: Dict[str, str], chunks: List[Dict[str, Any]]):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("DELETE FROM video_chunks WHERE video_id = ?", (video_info['id'],))
        
        texts = [chunk['text'] for chunk in chunks]
        
        try:
            response = self.client.embeddings.create(
                model="text-embedding-3-small",
                input=texts
            )
            embeddings = [np.array(embedding.embedding, dtype=np.float32) for embedding in response.data]
        except Exception as e:
            print(f"Error generating embeddings: {e}")
            return
        
        for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
            embedding_blob = embedding.tobytes()
            
            cursor.execute('''
                INSERT INTO video_chunks 
                (video_id, video_title, chunk_text, chunk_index, start_time, end_time, embedding)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', (
                video_info['id'],
                video_info['title'],
                chunk['text'],
                i,
                chunk['start_time'],
                chunk['end_time'],
                embedding_blob
            ))
        
        conn.commit()
        conn.close()
        print(f"Stored {len(chunks)} chunks for video: {video_info['title']}")
    
    def process_url(self, url: str):
        video_id = self.extract_video_id(url)
        
        if 'playlist' in url.lower() or len(video_id) > 20:
            print(f"Processing playlist: {video_id}")
            videos = self.get_playlist_videos(video_id)
            
            for video in videos:
                print(f"Processing video: {video['title']}")
                transcript = self.get_transcript(video['id'])
                
                if transcript:
                    chunks = self.chunk_transcript(transcript)
                    self.store_embeddings(video, chunks)
                else:
                    print(f"No transcript available for: {video['title']}")
        
        else:
            print(f"Processing single video: {video_id}")
            video_info = self.get_video_info(video_id)
            transcript = self.get_transcript(video_id)
            
            if transcript:
                chunks = self.chunk_transcript(transcript)
                self.store_embeddings(video_info, chunks)
            else:
                print(f"No transcript available for: {video_info['title']}")

def main():
    parser = argparse.ArgumentParser(description='Ingest YouTube videos into embeddings database')
    parser.add_argument('url', help='YouTube video or playlist URL')
    parser.add_argument('--db-path', default='youtube_embeddings.db', help='Database path')
    
    args = parser.parse_args()
    
    ingestion = YouTubeEmbeddingIngestion(args.db_path)
    ingestion.process_url(args.url)

if __name__ == "__main__":
    main()