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/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}