add search, yt api v3, pydantic

firkin-gadabout · firkin-gadabout · commit 83a5bdf8ca7f · 2025-05-17T22:47:38.000-04:00
diff --git a/README.md b/README.md
@@ -61,8 +61,9 @@ SMART enhances the problem-solving capabilities of language models by breaking d
 - Supports various search focuses including web, news, Wikipedia, academia, Reddit, images, and videos.
 
 ### 4. YouTube Tool (youtube.py)
-- Retrieves information about YouTube videos, including metadata and transcriptions.
-- Supports multiple languages for video transcripts.
+- Retrieves information about YouTube videos, including metadata and transcriptions. Supports multiple languages for video transcripts.
+- Searches YouTube by keyword (1–50 results), returning for each video its id and metadata.
+- Uses Youtube API v3 and needs a free api key from Google Cloud. https://developers.google.com/youtube/v3/getting-started
 
 **Installation**
 
diff --git a/tools/youtube.py b/tools/youtube.py
@@ -1,90 +1,208 @@
 """
-title: WolframAlpha API
+title: YouTube Utility Tools
 author: MartainInGreen
 author_url: https://github.com/MartianInGreen/OpenWebUI-Tools
-version: 0.1.0
-requirements: pytube, youtube_transcript_api
+version: 0.2.0
+requirements: youtube_transcript_api, google-api-python-client, pydantic
+
+This module provides tools for interacting with YouTube via Data API v3
+and fetching transcripts via the youtube_transcript_api.
+
+Available tools:
+  - Tools.transcript_download(video_id)  : Download video metadata and full transcript.
+  - Tools.search(query, max_results)    : Search YouTube for videos matching a keyword.
+
+All tools return only JSON-serializable types and support optional streaming
+via an `__event_emitter__` callback for incremental output. Search results
+are cached with LRU (maxsize=128) and enforce max_results bounds.
 """
 
 import json
-from pytubefix import YouTube #type: ignore
-from youtube_transcript_api import YouTubeTranscriptApi #type: ignore
-from pydantic import BaseModel, Field #type: ignore
-from typing import Callable, Awaitable
+from youtube_transcript_api import YouTubeTranscriptApi  # type: ignore
+from pydantic import BaseModel, Field  # type: ignore
+from typing import Callable, Awaitable, List, Dict
+from googleapiclient.discovery import build  # type: ignore
+from functools import lru_cache
 
-def youtube_func(video_id: str):
-    #print(f"Getting video with id: {video_id}")
-    yt = YouTube(f'https://www.youtube.com/watch?v={video_id}', use_po_token=True)
 
-    try:
-        # Get avalilbe languages
-        languages_raw = YouTubeTranscriptApi.list_transcripts(video_id)
-        languages = []
+class TranscriptDownloadResult(BaseModel):
+    video_id: str = Field(..., description="YouTube video ID")
+    title: str = Field(..., description="Video title")
+    channel: str = Field(..., description="Channel or author name")
+    description: str = Field(..., description="Full video description")
+    duration: str = Field(..., description="ISO8601 duration (e.g. 'PT5M33S')")
+    view_count: str = Field(..., description="Total view count as string")
+    transcription: List[str] = Field(
+        ..., description="List of transcript text segments"
+    )
+
+
+class SearchItem(BaseModel):
+    video_id: str = Field(..., description="Unique YouTube video ID")
+    title: str = Field(..., description="Video title")
+    channel: str = Field(..., description="Channel or uploader name")
+    published_at: str = Field(..., description="ISO8601 publication timestamp")
+    description: str = Field(..., description="Snippet description")
+    view_count: str = Field(..., description="Total view count as string")
+    like_count: str = Field(..., description="Total like count as string")
+    comment_count: str = Field(..., description="Total comment count as string")
+    length: str = Field(..., description="ISO8601 duration of the video")
+
+
+class SearchResult(BaseModel):
+    results: List[SearchItem] = Field(..., description="List of search results")
 
-        for lang in languages_raw:
-            languages.append(lang.language_code)
-        #print(f"Available languages: {languages}")
 
-        # Get transcript (get english by default, if not available get first language)
-        try:
-            transcript = YouTubeTranscriptApi.get_transcript(video_id, languages=['en'])
-        except:
-            transcript = YouTubeTranscriptApi.get_transcript(video_id, languages=[languages[0]])
-   
-        text = []
-        for part in transcript:
-            text.append(part['text'])
-
-        print(text)
-    except: 
-        text = "Could not get video transcript :("
-
-    try:
-        i: int = yt.watch_html.find('"shortDescription":"')
-        desc: str = '"'
-        i += 20  # excluding the `"shortDescription":"`
-        while True:
-            letter = yt.watch_html[i]
-            desc += letter  # letter can be added in any case
-            i += 1
-            if letter == '\\':
-                desc += yt.watch_html[i]
-                i += 1
-            elif letter == '"':
-                break
-        
-        return {
-            'title': yt.title,
-            'channel': yt.author,
-            'description': desc,
-            'length': str(yt.length) + " s",
-            'views': yt.views,
-            'transcription': text
-        }
-    except Exception as e:
-        print(str(e))
-        return {"Something went wrong :("}
-    
 class Tools:
     class Valves(BaseModel):
-        WOLFRAMALPHA_APP_ID: str = Field(
-            default="",
-            description="The App ID (api key) to authorize WolframAlpha",
+        YOUTUBE_API_KEY: str = Field(
+            "", description="YouTube Data API v3 key for all API calls"
         )
 
     def __init__(self):
+        """
+        Initialize the Tools container.
+
+        - Set `valves.YOUTUBE_API_KEY` before calling any methods.
+        - `self.citation = True` indicates support for citing results.
+        """
         self.valves = self.Valves()
         self.citation = True
 
-    def youtube(
-        self, video_id: str, __event_emitter__: Callable[[dict], Awaitable[None]]
-    ) -> str:
+    def transcript_download(
+        self,
+        video_id: str,
+        __event_emitter__: Callable[[dict], Awaitable[None]] = None,
+    ) -> Dict:
         """
-        This function lets you get information about YouTube videos. Including Metadata and Transcription.
-        :param video_id: Video ID of the YouTube video
-        :return: A short answer or explanation of the result of the query_string
+        Download metadata and full transcript for a given YouTube video.
+
+        Returns a dict matching TranscriptDownloadResult schema.
         """
+        # 1) Fetch transcript
+        try:
+            transcripts = YouTubeTranscriptApi.get_transcript(video_id)
+            transcription = [seg.get("text", "") for seg in transcripts]
+        except Exception:
+            transcription = []
+
+        # 2) Fetch metadata
+        youtube = build("youtube", "v3", developerKey=self.valves.YOUTUBE_API_KEY)
+        try:
+            resp = (
+                youtube.videos()
+                .list(part="snippet,contentDetails,statistics", id=video_id)
+                .execute()
+            )
+            items = resp.get("items", [])
+            if items:
+                item = items[0]
+                sn = item.get("snippet", {})
+                cd = item.get("contentDetails", {})
+                st = item.get("statistics", {})
+                data = {
+                    "video_id": video_id,
+                    "title": sn.get("title", ""),
+                    "channel": sn.get("channelTitle", ""),
+                    "description": sn.get("description", ""),
+                    "duration": cd.get("duration", ""),
+                    "view_count": st.get("viewCount", "0"),
+                    "transcription": transcription,
+                }
+            else:
+                data = {
+                    "video_id": video_id,
+                    "title": "",
+                    "channel": "",
+                    "description": "",
+                    "duration": "",
+                    "view_count": "0",
+                    "transcription": transcription,
+                }
+        except Exception:
+            data = {
+                "video_id": video_id,
+                "title": "",
+                "channel": "",
+                "description": "",
+                "duration": "",
+                "view_count": "0",
+                "transcription": transcription,
+            }
 
-        data = youtube_func(video_id)
+        # Validate & serialize via Pydantic
+        result = TranscriptDownloadResult(**data).dict()
+        if __event_emitter__:
+            __event_emitter__(
+                {"type": "message", "data": {"content": json.dumps(result)}}
+            )
+        return result
 
-        return json.dumps(data)
+    @lru_cache(maxsize=128)
+    def _search_logic(self, query: str, max_results: int) -> List[Dict]:
+        # Validate bounds
+        if not 1 <= max_results <= 50:
+            raise ValueError("max_results must be between 1 and 50")
+        youtube = build("youtube", "v3", developerKey=self.valves.YOUTUBE_API_KEY)
+        search_resp = (
+            youtube.search()
+            .list(part="snippet", q=query, type="video", maxResults=max_results)
+            .execute()
+        )
+        results: List[Dict] = []
+        video_ids: List[str] = []
+        for item in search_resp.get("items", []):
+            vid = item["id"]["videoId"]
+            snip = item["snippet"]
+            video_ids.append(vid)
+            entry = {
+                "video_id": vid,
+                "title": snip.get("title", ""),
+                "channel": snip.get("channelTitle", ""),
+                "published_at": snip.get("publishedAt", ""),
+                "description": snip.get("description", ""),
+            }
+            results.append(entry)
+        if video_ids:
+            detail_resp = (
+                youtube.videos()
+                .list(part="statistics,contentDetails", id=",".join(video_ids))
+                .execute()
+            )
+            detail_map = {item["id"]: item for item in detail_resp.get("items", [])}
+            for entry in results:
+                det = detail_map.get(entry["video_id"], {})
+                stats = det.get("statistics", {})
+                cd = det.get("contentDetails", {})
+                entry.update(
+                    {
+                        "view_count": stats.get("viewCount", "0"),
+                        "like_count": stats.get("likeCount", "0"),
+                        "comment_count": stats.get("commentCount", "0"),
+                        "length": cd.get("duration", ""),
+                    }
+                )
+        return results
+
+    def search(
+        self,
+        query: str,
+        max_results: int = 10,
+        __event_emitter__: Callable[[dict], Awaitable[None]] = None,
+    ) -> Dict:
+        """
+        Search YouTube for videos matching a keyword, with LRU caching
+        and max_results validation.
+
+        Returns a dict matching SearchResult schema.
+        """
+        entries = self._search_logic(query, max_results)
+        # Validate & serialize via Pydantic
+        items = [SearchItem(**e) for e in entries]
+        result = SearchResult(results=items).dict()
+        if __event_emitter__:
+            for entry in items:
+                __event_emitter__(
+                    {"type": "message", "data": {"content": json.dumps(entry.dict())}}
+                )
+        return result
