diff --git a/summarize.py b/summarize.py index 73f1432..9ae3a45 100755 --- a/summarize.py +++ b/summarize.py @@ -6,11 +6,14 @@ from typing import cast import click import pytgpt.phind as phind +debug = False + def summarize_text_file(content: str, prompt: str | None) -> str: prompt = prompt or "Please summarize the following transcript:" bot = phind.PHIND() - return bot.chat(f"{prompt} {content}") + res = bot.chat(f"{prompt} {content}") + return res def extract_transcript_contents(content: str, keep_newlines: bool = False) -> str: @@ -43,7 +46,7 @@ def grab_subtitles(url: str | Path) -> Path: with yt_dlp.YoutubeDL(ydl_opts) as ydl: info = ydl.extract_info(url, download=True) filename = ydl.prepare_filename(info) - print(f"Subtitle file saved as: {filename}") + dprint(f"Subtitle file saved as: {filename}") for root, _, files in Path(temp_dir).walk(): for file in files: @@ -53,21 +56,6 @@ def grab_subtitles(url: str | Path) -> Path: raise ValueError("No correct json3 transcript object found.") -def get_temp_dir() -> str: - # Create a temporary directory - temp_dir = tempfile.mkdtemp() - print(f"Creating temp dir {temp_dir}") - return temp_dir - - -def rm_dir(dir: Path | str) -> None: - # Remove the temporary directory - import shutil - - if Path(dir).is_dir(): - shutil.rmtree(dir) - - @click.command() # TODO: Can I set it so it checks existence *only* when no youtube flag exists? @click.argument("file_path", type=click.Path(exists=False)) @@ -86,13 +74,29 @@ def rm_dir(dir: Path | str) -> None: @click.option( "--prompt", "-p", - default="Please provide a detailed but concise summary for the following transcript:", + default="Please summarize the following transcript:", type=str, help="Use custom prompt.", ) -def cli(file_path: Path | str, json_transcript: bool, youtube: bool, prompt: str): +@click.option( + "--log-level", + "-l", + default=0, + help="Set log level to 1 for debug.", +) +def cli( + file_path: Path | str, + json_transcript: bool, + youtube: bool, + prompt: str, + log_level: int, +): """Provide summary for a file at the specified path or a youtube video at the specified url.""" + if log_level: + global debug + debug = True + content = "" # youtube link, dl transcript @@ -100,7 +104,7 @@ def cli(file_path: Path | str, json_transcript: bool, youtube: bool, prompt: str file_path = grab_subtitles(file_path) file_path = cast(Path, file_path) - print(f"DEBUG: file path = {file_path}") + dprint(f"file path = {file_path}") # load local file with Path(file_path).open() as f: content = f.read() @@ -108,12 +112,28 @@ def cli(file_path: Path | str, json_transcript: bool, youtube: bool, prompt: str if json_transcript or youtube: content = extract_transcript_contents(content) - print(f"DEBUG: content = {content}") + dprint(f"content = {content}") if not content: print("Please provide a file with valid content.") print(summarize_text_file(content, prompt)) +cached_dir: Path | None = None + + +def get_temp_dir() -> Path: + global cached_dir + if cached_dir is None: + cached_dir = Path(tempfile.mkdtemp()) + dprint(f"Created and cached temp dir {cached_dir}") + return cached_dir + + +def dprint(content: str | None): + if debug: + print(f"[DEBUG] {content}") + + if __name__ == "__main__": cli()