feat: Add CLI module for YouTube summarization tool
This commit is contained in:
311
summarize_yt/cli.py
Normal file
311
summarize_yt/cli.py
Normal file
@@ -0,0 +1,311 @@
|
||||
# Standard library imports remain the same
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
# Third-party packages
|
||||
import anthropic
|
||||
|
||||
# Local modules
|
||||
from . import __version__
|
||||
|
||||
# Add cost estimation constants
|
||||
CLAUDE_COST_PER_1K_INPUT = 0.015 # Cost per 1K tokens for input
|
||||
CLAUDE_COST_PER_1K_OUTPUT = 0.075 # Cost per 1K tokens for output
|
||||
ESTIMATED_TOKENS_PER_CHAR = 0.25 # Rough estimate of tokens per character
|
||||
|
||||
|
||||
def estimate_api_cost(text: str) -> float:
|
||||
"""Estimate the cost of sending text to Claude API."""
|
||||
estimated_input_tokens = len(text) * ESTIMATED_TOKENS_PER_CHAR
|
||||
estimated_output_tokens = 1024 # max_tokens setting
|
||||
|
||||
input_cost = (estimated_input_tokens / 1000) * CLAUDE_COST_PER_1K_INPUT
|
||||
output_cost = (estimated_output_tokens / 1000) * CLAUDE_COST_PER_1K_OUTPUT
|
||||
|
||||
return input_cost + output_cost
|
||||
|
||||
|
||||
def sanitize_filename(url: str) -> str:
|
||||
"""Convert URL to safe filename, keeping video ID."""
|
||||
# Extract video ID if it's a YouTube URL
|
||||
video_id = None
|
||||
if "youtube.com" in url or "youtu.be" in url:
|
||||
if "v=" in url:
|
||||
video_id = url.split("v=")[1].split("&")[0]
|
||||
else:
|
||||
video_id = url.split("/")[-1].split("?")[0]
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
return f"{timestamp}_{video_id if video_id else 'video'}"
|
||||
|
||||
|
||||
def ensure_transcript_dir() -> Path:
|
||||
"""Create and return path to transcript directory."""
|
||||
transcript_dir = Path.home() / ".yts" / "transcripts"
|
||||
transcript_dir.mkdir(parents=True, exist_ok=True)
|
||||
return transcript_dir
|
||||
|
||||
|
||||
def download_subtitles(url: str, quiet: bool = False) -> Tuple[bool, Optional[Path]]:
|
||||
"""Download subtitles from YouTube using yt-dlp."""
|
||||
try:
|
||||
cmd = ["yt-dlp", "--skip-download", "--write-auto-sub", "--sub-lang", "en"]
|
||||
if quiet:
|
||||
cmd.append("--quiet")
|
||||
cmd.append(url)
|
||||
|
||||
print(f"Debug: Running command: {cmd}") # Debug line
|
||||
|
||||
process = subprocess.Popen(
|
||||
cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True # This might help with encoding issues
|
||||
)
|
||||
|
||||
try:
|
||||
stdout, stderr = process.communicate(timeout=30)
|
||||
print(f"Debug: stdout: {stdout}") # Debug line
|
||||
print(f"Debug: stderr: {stderr}") # Debug line
|
||||
if process.returncode != 0:
|
||||
print(f"Error downloading subtitles: {stderr}", file=sys.stderr)
|
||||
return False, None
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
print("Download timed out after 30 seconds", file=sys.stderr)
|
||||
return False, None
|
||||
|
||||
# Find the downloaded VTT file
|
||||
current_dir = Path(".")
|
||||
vtt_files = list(current_dir.glob("*.en.vtt"))
|
||||
print(f"Debug: Found VTT files: {vtt_files}") # Debug line
|
||||
if not vtt_files:
|
||||
print("No VTT file found after download", file=sys.stderr)
|
||||
return False, None
|
||||
|
||||
return True, vtt_files[0]
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error during subtitle download: {str(e)}", file=sys.stderr)
|
||||
print(f"Debug: Full exception info: {type(e).__name__}: {str(e)}") # Debug line
|
||||
import traceback
|
||||
|
||||
traceback.print_exc() # This will print the full traceback
|
||||
return False, None
|
||||
|
||||
|
||||
def clean_vtt_text(text: str) -> str:
|
||||
"""Clean WebVTT text by removing timestamps, formatting tags, and metadata."""
|
||||
lines = text.split("\n")
|
||||
|
||||
# Remove header lines
|
||||
while lines and (
|
||||
lines[0].startswith("WEBVTT")
|
||||
or lines[0].startswith("Kind:")
|
||||
or lines[0].startswith("Language:")
|
||||
or not lines[0].strip()
|
||||
):
|
||||
lines.pop(0)
|
||||
|
||||
# Process remaining lines
|
||||
cleaned_lines = []
|
||||
for line in lines:
|
||||
if "-->" in line or not line.strip():
|
||||
continue
|
||||
line = re.sub(r"<\d{2}:\d{2}:\d{2}\.\d{3}>", "", line)
|
||||
line = re.sub(r"</?c>", "", line)
|
||||
if line.strip():
|
||||
cleaned_lines.append(line.strip())
|
||||
|
||||
seen = set()
|
||||
unique_lines = [x for x in cleaned_lines if not (x in seen or seen.add(x))]
|
||||
return "\n".join(unique_lines)
|
||||
|
||||
|
||||
def process_file(input_path: Path, output_path: Path | None = None, quiet: bool = False) -> str:
|
||||
"""Process a single VTT file and output cleaned text."""
|
||||
try:
|
||||
with open(input_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
cleaned_text = clean_vtt_text(content)
|
||||
|
||||
if output_path:
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
f.write(cleaned_text)
|
||||
if not quiet: # Changed from args.quiet to quiet
|
||||
print(f"Processed {input_path} -> {output_path}")
|
||||
return cleaned_text
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {input_path}: {str(e)}", file=sys.stderr)
|
||||
return ""
|
||||
|
||||
|
||||
def save_transcript(text: str, url: str, prompt: str) -> Path:
|
||||
"""Save transcript with metadata to ~/.yts/transcripts/."""
|
||||
transcript_dir = ensure_transcript_dir()
|
||||
filename = sanitize_filename(url) + ".txt"
|
||||
filepath = transcript_dir / filename
|
||||
|
||||
metadata = f"""URL: {url}
|
||||
Script Version: {__version__}
|
||||
Timestamp: {datetime.now().isoformat()}
|
||||
Claude Prompt: {prompt}
|
||||
---
|
||||
"""
|
||||
|
||||
with open(filepath, "w", encoding="utf-8") as f:
|
||||
f.write(metadata + text)
|
||||
|
||||
return filepath
|
||||
|
||||
|
||||
def cleanup_files(vtt_path: Optional[Path]):
|
||||
"""Remove downloaded files after processing."""
|
||||
try:
|
||||
if vtt_path and vtt_path.exists():
|
||||
vtt_path.unlink()
|
||||
|
||||
# Also cleanup any other VTT files in current directory
|
||||
current_dir = Path(".")
|
||||
for vtt_file in current_dir.glob("*.vtt"):
|
||||
vtt_file.unlink()
|
||||
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not clean up temporary files: {str(e)}", file=sys.stderr)
|
||||
|
||||
|
||||
def get_summary_from_claude(text: str, prompt: str = "Please summarize this transcript in 500 words or less") -> str:
|
||||
"""Send text to Claude API for summarization."""
|
||||
try:
|
||||
api_key = os.environ.get("ANTHROPIC_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("ANTHROPIC_API_KEY environment variable not set")
|
||||
|
||||
client = anthropic.Anthropic()
|
||||
message = client.messages.create(
|
||||
model="claude-3-sonnet-20240229",
|
||||
max_tokens=1024,
|
||||
temperature=0,
|
||||
system="You are a helpful assistant that summarizes transcripts accurately and concisely.",
|
||||
messages=[{"role": "user", "content": f"{prompt}:\n\n{text}"}],
|
||||
)
|
||||
return message.content[0].text
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error getting summary from Claude: {str(e)}", file=sys.stderr)
|
||||
return ""
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Download YouTube subtitles and get Claude summary")
|
||||
parser.add_argument("url", help="YouTube video URL", type=str)
|
||||
parser.add_argument("-o", "--output", help="Output file for summary")
|
||||
parser.add_argument("-q", "--quiet", action="store_true", help="Suppress status messages")
|
||||
parser.add_argument("--keep-files", action="store_true", help="Don't delete downloaded VTT files")
|
||||
parser.add_argument("--transcript", action="store_true", help="Include full transcript in output")
|
||||
parser.add_argument(
|
||||
"--prompt", default="Please summarize this transcript in 500 words or less", help="Custom prompt for Claude"
|
||||
)
|
||||
parser.add_argument("-y", "--yes", action="store_true", help="Skip cost confirmation")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up signal handler for clean exit
|
||||
def signal_handler(sig, frame):
|
||||
print("\nCleaning up and exiting...")
|
||||
cleanup_files(None) # Clean any VTT files
|
||||
print(f"Debug: Script terminated by signal at {datetime.now().isoformat()}")
|
||||
sys.exit(0)
|
||||
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
# Download subtitles
|
||||
if not args.quiet:
|
||||
print("Downloading subtitles...")
|
||||
success, vtt_path = download_subtitles(args.url, args.quiet)
|
||||
if not success:
|
||||
cleanup_files(None)
|
||||
sys.exit(1)
|
||||
|
||||
# Process the VTT file
|
||||
cleaned_text = process_file(vtt_path, None, args.quiet)
|
||||
|
||||
# Save transcript
|
||||
transcript_path = save_transcript(cleaned_text, args.url, args.prompt)
|
||||
print(f"\nTranscript saved to: {transcript_path}")
|
||||
|
||||
# Estimate and display cost
|
||||
estimated_cost = estimate_api_cost(cleaned_text)
|
||||
print(f"\nEstimated API cost: ${estimated_cost:.4f}")
|
||||
|
||||
if not args.yes:
|
||||
try:
|
||||
import tty
|
||||
import termios
|
||||
|
||||
# Save the terminal settings
|
||||
fd = sys.stdin.fileno()
|
||||
old_settings = termios.tcgetattr(fd)
|
||||
try:
|
||||
# Set the terminal to raw mode
|
||||
tty.setraw(sys.stdin.fileno())
|
||||
sys.stdout.write("\nDo you want to proceed with getting the summary? (y/N): ")
|
||||
sys.stdout.flush()
|
||||
# Read a single character
|
||||
char = sys.stdin.read(1)
|
||||
# Print a newline since we're in raw mode
|
||||
sys.stdout.write("\n")
|
||||
sys.stdout.flush()
|
||||
finally:
|
||||
# Restore terminal settings
|
||||
termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
|
||||
|
||||
if char.lower() != "y":
|
||||
print("Operation cancelled by user.")
|
||||
cleanup_files(vtt_path)
|
||||
sys.exit(0)
|
||||
|
||||
except (EOFError, KeyboardInterrupt, termios.error):
|
||||
print("\nOperation cancelled by user.")
|
||||
cleanup_files(vtt_path)
|
||||
sys.exit(0)
|
||||
|
||||
# Get summary from Claude
|
||||
if not args.quiet:
|
||||
print("\nGetting summary from Claude...")
|
||||
summary = get_summary_from_claude(cleaned_text, args.prompt)
|
||||
|
||||
# Prepare output
|
||||
output = ""
|
||||
if args.transcript:
|
||||
output += "=== Full Transcript ===\n\n"
|
||||
output += cleaned_text
|
||||
output += "\n\n=== Summary ===\n\n"
|
||||
output += summary
|
||||
|
||||
# Output results
|
||||
output_path = args.output
|
||||
if output_path:
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
f.write(output)
|
||||
print(f"\nSummary saved to: {output_path}")
|
||||
else:
|
||||
print("\n=== Summary ===")
|
||||
print(output)
|
||||
|
||||
# Cleanup downloaded files unless --keep-files is specified
|
||||
if not args.keep_files:
|
||||
cleanup_files(vtt_path)
|
||||
|
||||
print(f"\nDebug: Script completed successfully at {datetime.now().isoformat()}")
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user