From ad669920e7a22fc0cdff238cf49e3101e1355112 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Wed, 11 Sep 2024 22:22:14 +0200
Subject: [PATCH] feat: added youtube tool (#116)

* feat: added youtube tool

* fix: fixed excessive whitespace in patch example prompt
---
 gptme/tools/__init__.py |  2 ++
 gptme/tools/youtube.py  | 42 +++++++++++++++++++++++++++++++++++++++++
 poetry.lock             | 16 +++++++++++++++-
 pyproject.toml          |  3 +++
 4 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 gptme/tools/youtube.py

diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py
index 9dabf8d1..12285715 100644
--- a/gptme/tools/__init__.py
+++ b/gptme/tools/__init__.py
@@ -15,6 +15,7 @@
 from .shell import tool as shell_tool
 from .subagent import tool as subagent_tool
 from .tmux import tool as tmux_tool
+from .youtube import tool as youtube_tool
 
 logger = logging.getLogger(__name__)
 
@@ -37,6 +38,7 @@
     browser_tool,
     gh_tool,
     chats_tool,
+    youtube_tool,
     # python tool is loaded last to ensure all functions are registered
     get_python_tool,
 ]
diff --git a/gptme/tools/youtube.py b/gptme/tools/youtube.py
new file mode 100644
index 00000000..4ca7bddb
--- /dev/null
+++ b/gptme/tools/youtube.py
@@ -0,0 +1,42 @@
+import logging
+
+from ..llm import summarize
+from .base import ToolSpec
+
+logger = logging.getLogger(__name__)
+
+try:
+    # noreorder
+    from youtube_transcript_api import YouTubeTranscriptApi  # fmt: skip
+except ImportError:
+    YouTubeTranscriptApi = None
+    logger.warning(
+        "youtube_transcript_api not available. YouTube tool will be disabled."
+    )
+
+
+def get_transcript(video_id: str) -> str:
+    if not YouTubeTranscriptApi:
+        return "Error: youtube_transcript_api is not installed."
+    try:
+        transcript = YouTubeTranscriptApi.get_transcript(video_id)
+        return " ".join([entry["text"] for entry in transcript])
+    except Exception as e:
+        logger.error(f"Error fetching transcript: {e}")
+        return f"Error fetching transcript: {e}"
+
+
+def summarize_transcript(transcript: str) -> str:
+    return summarize(transcript).content
+
+
+tool: ToolSpec = ToolSpec(
+    name="youtube",
+    desc="Fetch and summarize YouTube video transcripts",
+    instructions="""
+    To use this tool, provide a YouTube video ID and specify whether you want the transcript or a summary.
+    """,
+    functions=[get_transcript, summarize_transcript],
+    block_types=["youtube"],
+    available=bool(YouTubeTranscriptApi),
+)
diff --git a/poetry.lock b/poetry.lock
index f405ae77..ca5b0ed3 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3586,6 +3586,20 @@ files = [
     {file = "windows_curses-2.3.3-cp39-cp39-win_amd64.whl", hash = "sha256:23b71f3a44cc5a5b92e12d3a519e840e54f8252ff3ef867b630963d5fe7124dd"},
 ]
 
+[[package]]
+name = "youtube-transcript-api"
+version = "0.6.2"
+description = "This is an python API which allows you to get the transcripts/subtitles for a given YouTube video. It also works for automatically generated subtitles, supports translating subtitles and it does not require a headless browser, like other selenium based solutions do!"
+optional = true
+python-versions = "*"
+files = [
+    {file = "youtube_transcript_api-0.6.2-py3-none-any.whl", hash = "sha256:019dbf265c6a68a0591c513fff25ed5a116ce6525832aefdfb34d4df5567121c"},
+    {file = "youtube_transcript_api-0.6.2.tar.gz", hash = "sha256:cad223d7620633cec44f657646bffc8bbc5598bd8e70b1ad2fa8277dec305eb7"},
+]
+
+[package.dependencies]
+requests = "*"
+
 [extras]
 all = ["flask", "matplotlib", "numpy", "pandas", "pillow", "playwright", "torch", "transformers"]
 browser = ["playwright"]
@@ -3596,4 +3610,4 @@ training = ["torch", "transformers"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "df5a8e446b099649784af5be88df0a7531cb0421899cabb9e13ee623b98c7f54"
+content-hash = "a47dbcddd50026b7f57fd6c03990d1b2ccb98fb57b0bc9db623316d3356799a3"
diff --git a/pyproject.toml b/pyproject.toml
index 051abd89..6bf3d1ae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,9 @@ playwright = {version = "^1.37.0", optional=true}
 openai = "^1.0"
 anthropic = "^0.34.0"
 
+# tools
+youtube_transcript_api = {version = "^0.6.1", optional = true}
+
 # datascience essentials
 matplotlib = {version = "*", optional=true}
 pandas = {version = "*", optional=true}