main.py

import gradio as gr
import sqlite3
import requests
import os

def fetch_data(db_path, limit_links=None):
    os.makedirs(os.path.dirname(db_path), exist_ok=True)
    if not os.path.exists(db_path):
        # We can either create it or return empty. I will let sqlite create it if it doesn't exist.
        pass
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    cursor.execute('''CREATE TABLE IF NOT EXISTS links (
                        id INTEGER PRIMARY KEY,
                        url TEXT,
                        status TEXT DEFAULT 'pending'
                    )''')
    try:
        cursor.execute("SELECT api_key FROM accounts")
        api_keys = [row[0] for row in cursor.fetchall() if row[0]]
    except Exception as e:
        api_keys = []
        
    try:
        query = "SELECT id, url FROM links WHERE status='pending'"
        if limit_links:
            query += f" LIMIT {limit_links}"
            
        cursor.execute(query)
        link_records = cursor.fetchall()
        links = [row[1] for row in link_records if row[1]]
        
        if link_records:
            ids = [str(row[0]) for row in link_records]
            cursor.execute(f"UPDATE links SET status='processing' WHERE id IN ({','.join(ids)})")
            conn.commit()
    except Exception as e:
        links = []
        
    conn.close()
    return api_keys, links

def trigger_workers(db_path, worker_urls_str, links_per_worker):
    worker_urls = [u.strip() for u in worker_urls_str.split(',') if u.strip()]
    if not worker_urls:
        return "Error: No worker URLs provided."
        
    api_keys, links = fetch_data(db_path, limit_links=int(links_per_worker) * len(worker_urls))
    if not api_keys:
        return "Error: No API keys found in DB."
    if not links:
        return "Error: No pending links found in DB."
        
    link_chunks = [links[i:i + int(links_per_worker)] for i in range(0, len(links), int(links_per_worker))]
    
    logs = []
    logs.append(f"Loaded {len(api_keys)} API keys and {len(links)} links.")
    logs.append(f"Divided into {len(link_chunks)} worker chunks.")
    
    for idx, chunk in enumerate(link_chunks):
        if idx >= len(worker_urls):
            logs.append("More chunks than workers. Remaining chunks won't be sent.")
            break
            
        worker_url = worker_urls[idx]
        api_key = api_keys[idx % len(api_keys)]
        
        payload = {
            "links": chunk,
            "api_key": api_key
        }
        
        endpoint = f"{worker_url.rstrip('/')}/process"
        try:
            resp = requests.post(endpoint, json=payload, timeout=10)
            logs.append(f"Sent {len(chunk)} links to {worker_url}: {resp.status_code} - {resp.text}")
        except Exception as e:
            logs.append(f"Failed to send to {worker_url}: {str(e)}")
            
    return "\n".join(logs)

import json
import uuid
import shutil
import subprocess
from pathlib import Path

def get_kaggle_username():
    kaggle_json = Path.home() / ".kaggle" / "kaggle.json"
    if not kaggle_json.exists():
        return "your-username"
    with open(kaggle_json, "r") as f:
        creds = json.load(f)
        return creds.get("username", "your-username")

def trigger_kaggle_test():
    username = get_kaggle_username()
    worker_title = f"yt-test-worker-{uuid.uuid4().hex[:6]}"
    worker_dir = Path("kaggle_test_worker")
    
    if worker_dir.exists():
        shutil.rmtree(worker_dir)
    worker_dir.mkdir(parents=True)
    
    notebook_content = {
        "cells": [
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "print('Hello from Kaggle Worker!')\n",
                    "print('Test successful.')\n"
                ]
            }
        ],
        "metadata": {
            "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
            "language_info": {"name": "python"}
        },
        "nbformat": 4,
        "nbformat_minor": 4
    }
    
    with open(worker_dir / "notebook.ipynb", "w") as f:
        json.dump(notebook_content, f, indent=2)
        
    metadata = {
        "id": f"{username}/{worker_title}",
        "title": worker_title,
        "code_file": "notebook.ipynb",
        "language": "python",
        "kernel_type": "notebook",
        "is_private": True,
        "enable_gpu": False,
        "enable_internet": True,
        "dataset_sources": [],
        "competition_sources": [],
        "kernel_sources": []
    }
    
    with open(worker_dir / "kernel-metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)
        
    logs = [f"Prepared Kaggle Test Worker: {worker_title}"]
    
    # Use kaggle from PATH (assuming it is installed via requirements.txt)
    command = ["kaggle", "kernels", "push", "-p", str(worker_dir)]
    logs.append(f"Pushing {worker_title} to Kaggle...")
    
    try:
        res = subprocess.run(command, capture_output=True, text=True)
        if res.returncode == 0:
            logs.append("Push successful.")
            logs.append(res.stdout)
        else:
            logs.append("Push failed.")
            logs.append(res.stderr)
    except Exception as e:
        logs.append(f"Error running kaggle command: {e}")
        logs.append("Make sure Kaggle API credentials are provided in ~/.kaggle/kaggle.json")
        
    return "\n".join(logs)

def create_worker_notebook(worker_dir, config_name):
    notebook_content = {
        "cells": [
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "!pip install openai huggingface_hub yt-dlp opencv-python-headless\n"
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "import json\n",
                    "import os\n",
                    "import subprocess\n",
                    "import sys\n",
                    "\n",
                    f"with open('{config_name}') as f:\n",
                    "    config = json.load(f)\n",
                    "\n",
                    "links = config['links']\n",
                    "api_key = config['api_key']\n",
                    "os.environ['FEATHERLESS_API_KEY'] = api_key\n",
                    "\n",
                    "print(f'Starting worker with {len(links)} links using API Key: {api_key[:5]}...')\n",
                    "\n",
                    "for link in links:\n",
                    "    print(f'\\n[+] Processing: {link}')\n",
                    "    # 1. Download Video\n",
                    "    dl_cmd = [sys.executable, 'worker/youtube_worker/downloader.py', link, '--resolution', 'HD', '-o', 'test_output']\n",
                    "    res = subprocess.run(dl_cmd, capture_output=True, text=True)\n",
                    "    if res.returncode != 0:\n",
                    "        print(f'Download failed:\\n{res.stderr}')\n",
                    "        continue\n",
                    "    \n",
                    "    # Find the downloaded video path from stdout (last line)\n",
                    "    lines = res.stdout.strip().split('\\n')\n",
                    "    video_path = lines[-1]\n",
                    "    if not os.path.exists(video_path):\n",
                    "        print(f'Could not find downloaded video at: {video_path}')\n",
                    "        continue\n",
                    "        \n",
                    "    print(f'[+] Video downloaded to: {video_path}')\n",
                    "    \n",
                    "    # 2. Curate Model (Video Worker)\n",
                    "    prompt = 'Analyze this video comprehensively. Focus on visuals, objects, and text.'\n",
                    "    vw_cmd = [sys.executable, 'worker/models_worker/video_worker.py', '--video-url', video_path, '--prompt', prompt, '--api-key', api_key]\n",
                    "    print('[+] Running model curation...')\n",
                    "    subprocess.run(vw_cmd)\n",
                    "    \n",
                    "    # 3. Optional: Upload to HF\n",
                    "    # hf_cmd = [sys.executable, 'worker/hf_uploader.py', 'test_output']\n",
                    "    # subprocess.run(hf_cmd)\n",
                    "\n",
                    "print('\\nAll tasks completed!')\n"
                ]
            }
        ],
        "metadata": {
            "kernelspec": {
                "display_name": "Python 3",
                "language": "python",
                "name": "python3"
            },
            "language_info": {
                "name": "python"
            }
        },
        "nbformat": 4,
        "nbformat_minor": 4
    }
    
    with open(worker_dir / "notebook.ipynb", "w") as f:
        json.dump(notebook_content, f, indent=2)

def trigger_kaggle_workers(db_path, links_per_worker, max_workers):
    api_keys, links = fetch_data(db_path, limit_links=int(links_per_worker) * int(max_workers))
    if not api_keys:
        return "Error: No API keys found in DB."
    if not links:
        return "Error: No pending links found in DB."
        
    username = get_kaggle_username()
    workers_dir = Path("kaggle_workers_out")
    if workers_dir.exists():
        shutil.rmtree(workers_dir)
    workers_dir.mkdir(parents=True)
    
    logs = []
    logs.append(f"Loaded {len(api_keys)} API keys and {len(links)} links.")
    
    link_chunks = [links[i:i + int(links_per_worker)] for i in range(0, len(links), int(links_per_worker))]
    logs.append(f"Divided into {len(link_chunks)} worker chunks.")
    
    for idx, chunk in enumerate(link_chunks):
        worker_id = f"worker-{uuid.uuid4().hex[:8]}"
        worker_title = f"yt-worker-{idx+1}-{uuid.uuid4().hex[:4]}"
        
        api_key = api_keys[idx % len(api_keys)]
        worker_path = workers_dir / worker_id
        worker_path.mkdir(parents=True)
        
        # Copy worker if exists locally
        if Path("worker").exists():
            shutil.copytree("worker", worker_path / "worker")
        else:
            # We copy from parent directory if this is running locally, otherwise it's an error on HF unless packed
            parent_worker = Path("../hf_space_worker/worker")
            if parent_worker.exists():
                shutil.copytree(parent_worker, worker_path / "worker")
            else:
                logs.append("Warning: 'worker' directory not found. Kaggle notebook may fail.")
                
        if Path("cookies.txt").exists():
            shutil.copy("cookies.txt", worker_path / "cookies.txt")
        elif Path("../cookies.txt").exists():
            shutil.copy("../cookies.txt", worker_path / "cookies.txt")
            
        config_name = "worker_config.json"
        config_data = {
            "links": chunk,
            "api_key": api_key,
        }
        with open(worker_path / config_name, "w") as f:
            json.dump(config_data, f, indent=2)
            
        create_worker_notebook(worker_path, config_name)
        
        metadata = {
            "id": f"{username}/{worker_title}",
            "title": worker_title,
            "code_file": "notebook.ipynb",
            "language": "python",
            "kernel_type": "notebook",
            "is_private": True,
            "enable_gpu": False,
            "enable_internet": True,
            "dataset_sources": [],
            "competition_sources": [],
            "kernel_sources": []
        }
        with open(worker_path / "kernel-metadata.json", "w") as f:
            json.dump(metadata, f, indent=2)
            
        logs.append(f"Prepared Worker {idx+1}: {worker_title} with {len(chunk)} links.")
        
        command = ["kaggle", "kernels", "push", "-p", str(worker_path)]
        logs.append(f"Pushing {worker_title} to Kaggle...")
        try:
            res = subprocess.run(command, capture_output=True, text=True)
            if res.returncode == 0:
                logs.append("Push successful.")
            else:
                logs.append("Push failed.")
                logs.append(res.stderr)
        except Exception as e:
            logs.append(f"Error: {e}")
            
    return "\n".join(logs)

with gr.Blocks() as app:
    gr.Markdown("# YouTube Video Orchestrator Host")
    gr.Markdown("Deploy this to Hugging Face Spaces or run locally. It fetches pending links from your database and delegates them to FastAPI workers or Kaggle.")
    
    with gr.Tabs():
        with gr.Tab("FastAPI Workers"):
            with gr.Row():
                db_input = gr.Textbox(label="DB Path", value="/data/db/accounts.db")
                workers_input = gr.Textbox(label="Worker URLs (comma separated)", placeholder="http://localhost:7860, https://your-worker.hf.space")
                links_per_worker = gr.Number(label="Links per worker", value=2)
                
            run_btn = gr.Button("Run Orchestration", variant="primary")
            output_log = gr.Textbox(label="Logs", lines=10)
            
            run_btn.click(fn=trigger_workers, inputs=[db_input, workers_input, links_per_worker], outputs=output_log)

        with gr.Tab("Kaggle Workers"):
            gr.Markdown("Trigger Kaggle notebook workers. Note: You must mount `kaggle.json` to `~/.kaggle/kaggle.json` in the container for this to work.")
            
            with gr.Row():
                k_db_input = gr.Textbox(label="DB Path", value="/data/db/accounts.db")
                k_links_per_worker = gr.Number(label="Links per worker", value=2)
                k_max_workers = gr.Number(label="Max Workers", value=5)
                
            k_run_btn = gr.Button("Run Kaggle Orchestration", variant="primary")
            k_test_run_btn = gr.Button("Run Simple Kaggle Test Worker", variant="secondary")
            kaggle_log = gr.Textbox(label="Kaggle Logs", lines=15)
            
            k_run_btn.click(fn=trigger_kaggle_workers, inputs=[k_db_input, k_links_per_worker, k_max_workers], outputs=kaggle_log)
            k_test_run_btn.click(fn=trigger_kaggle_test, inputs=[], outputs=kaggle_log)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)