fn.py

import os
import glob
import json
from yt_dlp import YoutubeDL

def load_model():
    pass
    
def run(url, mode = None):
    if not url.startswith('https://'):
        url = f'https://www.youtube.com/watch?v={url}'
    video_id = url[-11:]

    # cleanup
    for file in glob.glob(f'{video_id}.*'):
        os.remove(file)

    if mode == 'info':
        with YoutubeDL() as ydl:
            info = ydl.extract_info(url, download=False)
            return json.dumps(info)
    elif mode == 'chat':
        args = {
            'outtmpl': '%(id)s'+'.%(ext)s',
            'writesubtitles': True,
            'skip_download': True,
        }
        with YoutubeDL(args) as ydl:
            result = ydl.download([url])
            print(result)
            for file in glob.glob(f'{video_id}.*'):
                return file
    elif mode == 'audio':
        args = {
            'outtmpl': '%(id)s'+'.%(ext)s',
            'format': 'bestaudio',
        }
        with YoutubeDL(args) as ydl:
            result = ydl.download([url])
            print(result)
            for file in glob.glob(f'{video_id}.*'):
                return file
    else:
        args = {
            'outtmpl': '%(id)s'+'.%(ext)s',
            'format': 'best',
        }
        with YoutubeDL(args) as ydl:
            result = ydl.download([url])
            print(result)
            for file in glob.glob(f'{video_id}.*'):
                return file