Merge pull request #29 from jkawamoto/structured-output

Refactor YouTube transcript retrieval to structured model

Author: jkawamoto

src/mcp_youtube_transcript/__init__.py

@@ -8,15 +8,15 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from functools import lru_cache, partial
-from typing import AsyncIterator
+from typing import AsyncIterator, Tuple
 from typing import Final
 from urllib.parse import urlparse, parse_qs
 
 import requests
 from bs4 import BeautifulSoup
 from mcp.server import FastMCP
 from mcp.server.fastmcp import Context
-from pydantic import Field
+from pydantic import Field, BaseModel
 from youtube_transcript_api import YouTubeTranscriptApi
 from youtube_transcript_api.proxies import WebshareProxyConfig, GenericProxyConfig, ProxyConfig
 
@@ -35,7 +35,7 @@ async def _app_lifespan(_server: FastMCP, proxy_config: ProxyConfig | None) -> A
 
 
 @lru_cache
-def _get_transcript(ctx: AppContext, video_id: str, lang: str) -> str:
+def _get_transcript(ctx: AppContext, video_id: str, lang: str) -> Tuple[str, str]:
     if lang == "en":
         languages = ["en"]
     else:
@@ -46,11 +46,17 @@ def _get_transcript(ctx: AppContext, video_id: str, lang: str) -> str:
     )
     page.raise_for_status()
     soup = BeautifulSoup(page.text, "html.parser")
-    title = soup.title.string if soup.title else "Transcript"
+    title = soup.title.string if soup.title and soup.title.string else "Transcript"
 
     transcripts = ctx.ytt_api.fetch(video_id, languages=languages)
+    return title, "\n".join((item.text for item in transcripts))
 
-    return f"# {title}\n" + "\n".join((item.text for item in transcripts))
+
+class Transcript(BaseModel):
+    """Transcript of a YouTube video."""
+
+    title: str = Field(description="Title of the video")
+    transcript: str = Field(description="Transcript of the video")
 
 
 def server(
@@ -74,7 +80,7 @@ async def get_transcript(
         ctx: Context,
         url: str = Field(description="The URL of the YouTube video"),
         lang: str = Field(description="The preferred language for the transcript", default="en"),
-    ) -> str:
+    ) -> Transcript:
         """Retrieves the transcript of a YouTube video."""
         parsed_url = urlparse(url)
         if parsed_url.hostname == "youtu.be":
@@ -86,9 +92,10 @@ async def get_transcript(
             video_id = q[0]
 
         app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
-        return _get_transcript(app_ctx, video_id, lang)
+        title, transcript = _get_transcript(app_ctx, video_id, lang)
+        return Transcript(title=title, transcript=transcript)
 
     return mcp
 
 
-__all__: Final = ["server"]
+__all__: Final = ["server", "Transcript"]

tests/test_mcp.py

@@ -15,6 +15,8 @@
 from mcp.types import TextContent
 from youtube_transcript_api import YouTubeTranscriptApi
 
+from mcp_youtube_transcript import Transcript
+
 params = StdioServerParameters(command="uv", args=["run", "mcp-youtube-transcript"])
 
 
@@ -44,14 +46,18 @@ async def test_get_transcript(mcp_client_session: ClientSession) -> None:
     video_id = "LPZh9BOjkQs"
 
     title = fetch_title(video_id, "en")
-    expect = f"# {title}\n" + "\n".join((item.text for item in YouTubeTranscriptApi().fetch(video_id)))
+    expect = Transcript(
+        title=title, transcript="\n".join((item.text for item in YouTubeTranscriptApi().fetch(video_id)))
+    )
 
     res = await mcp_client_session.call_tool(
         "get_transcript",
         arguments={"url": f"https//www.youtube.com/watch?v={video_id}"},
     )
     assert isinstance(res.content[0], TextContent)
-    assert res.content[0].text == expect
+
+    transcript = Transcript.model_validate_json(res.content[0].text)
+    assert transcript == expect
     assert not res.isError
 
 
@@ -61,14 +67,18 @@ async def test_get_transcript_with_language(mcp_client_session: ClientSession) -
     video_id = "WjAXZkQSE2U"
 
     title = fetch_title(video_id, "ja")
-    expect = f"# {title}\n" + "\n".join((item.text for item in YouTubeTranscriptApi().fetch(video_id, ["ja"])))
+    expect = Transcript(
+        title=title, transcript="\n".join((item.text for item in YouTubeTranscriptApi().fetch(video_id, ["ja"])))
+    )
 
     res = await mcp_client_session.call_tool(
         "get_transcript",
         arguments={"url": f"https//www.youtube.com/watch?v={video_id}", "lang": "ja"},
     )
     assert isinstance(res.content[0], TextContent)
-    assert res.content[0].text == expect
+
+    transcript = Transcript.model_validate_json(res.content[0].text)
+    assert transcript == expect
     assert not res.isError
 
 
@@ -80,7 +90,9 @@ async def test_get_transcript_fallback_language(
     video_id = "LPZh9BOjkQs"
 
     title = fetch_title(video_id, "en")
-    expect = f"# {title}\n" + "\n".join((item.text for item in YouTubeTranscriptApi().fetch(video_id)))
+    expect = Transcript(
+        title=title, transcript="\n".join((item.text for item in YouTubeTranscriptApi().fetch(video_id)))
+    )
 
     res = await mcp_client_session.call_tool(
         "get_transcript",
@@ -90,7 +102,9 @@ async def test_get_transcript_fallback_language(
         },
     )
     assert isinstance(res.content[0], TextContent)
-    assert res.content[0].text == expect
+
+    transcript = Transcript.model_validate_json(res.content[0].text)
+    assert transcript == expect
     assert not res.isError
 
 
@@ -115,12 +129,16 @@ async def test_get_transcript_with_short_url(mcp_client_session: ClientSession)
     video_id = "LPZh9BOjkQs"
 
     title = fetch_title(video_id, "en")
-    expect = f"# {title}\n" + "\n".join((item.text for item in YouTubeTranscriptApi().fetch(video_id)))
+    expect = Transcript(
+        title=title, transcript="\n".join((item.text for item in YouTubeTranscriptApi().fetch(video_id)))
+    )
 
     res = await mcp_client_session.call_tool(
         "get_transcript",
         arguments={"url": f"https://youtu.be/{video_id}"},
     )
     assert isinstance(res.content[0], TextContent)
-    assert res.content[0].text == expect
+
+    transcript = Transcript.model_validate_json(res.content[0].text)
+    assert transcript == expect
     assert not res.isError