Rewrite for direct use with different LLM providersHEADmaster

1 files changed, 23 insertions, 6 deletions

diff --git a/src/tldw/main.py b/src/tldw/main.py
index efda992..b2652e7 100644
--- a/src/tldw/main.py
+++ b/src/tldw/main.py
@@ -1,20 +1,37 @@
+import anyio
 from argparse import ArgumentParser
 from urllib.parse import urlparse, parse_qs
 from youtube_transcript_api import YouTubeTranscriptApi
-from subprocess import run
+from multillm import Client
+
+DEFAULT_MODEL = "claude/sonnet"
+DEFAULT_SYSTEM_PROMPT = "Summarize the following video transcript:"
+
+
+async def summarize(model: str, prompt: str):
+    client = Client()
+    result = await client.single(model, prompt)
+    print(result)
+
 
 def main():
     parser = ArgumentParser()
     parser.add_argument("url")
     parser.add_argument("-l", "--language", default="en", help="Language code for captions (e.g., en, es, fr)")
+    parser.add_argument("-r", "--raw", action="store_true", help="Output raw transcript without summarizing")
+    parser.add_argument("-m", "--model", default=DEFAULT_MODEL, help="Model (e.g., openai/gpt-4o, anthropic/claude-sonnet-4-20250514, claude/sonnet)")
+    parser.add_argument("-s", "--system-prompt", default=DEFAULT_SYSTEM_PROMPT, help="Custom prompt for the LLM")
     args = parser.parse_args()
     api = YouTubeTranscriptApi()
     url = urlparse(args.url)
-    query = parse_qs(url.query)
-    transcript = api.fetch(query["v"][0], languages=[args.language])
-    text = "Summarize the following video transcript:\n"
-    text += " ".join([s.text for s in transcript.snippets])
-    run(["hey", text])
+    qs = parse_qs(url.query)
+    transcript = api.fetch(qs["v"][0], languages=[args.language])
+    text = " ".join([s.text for s in transcript.snippets])
+    if args.raw:
+        print(text)
+        return
+    prompt = args.system_prompt + "\n" + text
+    anyio.run(summarize, args.model, prompt)
 
 
 if __name__ == "__main__":