feat: add search_studios tool

pengqun · pengqun · commit 033fda01aef4 · 2025-07-29T11:22:11.000+08:00
diff --git a/demo.py b/demo.py
@@ -124,6 +124,31 @@ async def demo_search_datasets(client: Client) -> None:
     print()
 
 
+async def demo_search_studios(client: Client) -> None:
+    """Demo searching studios."""
+    tool_name = "search_studios"
+    print_step_title(tool_name, "🔍 Search studios (keyword='TTS', sort='VisitsCount', limit 3 results)")
+
+    result = await client.call_tool(tool_name, {"query": "TTS", "sort": "VisitsCount", "limit": 3})
+    data = parse_tool_response(result)
+
+    if isinstance(data, list) and data:
+        summaries = []
+        for studio in data:
+            name = studio.get("name", "N/A")
+            chinese_name = studio.get("chinese_name", "N/A")
+            status = studio.get("status", "N/A")
+            stars = studio.get("stars", 0)
+            visits = studio.get("visits", 0)
+
+            summaries.append(f"{name} ({chinese_name}) - Status={status}, Stars={stars}, Visits={visits}")
+
+        print(f"   • Result: Found {len(data)} items - {' | '.join(summaries)}")
+    else:
+        print("   • Result: No studios found")
+    print()
+
+
 async def demo_search_papers(client: Client) -> None:
     """Demo searching papers."""
     tool_name = "search_papers"
@@ -242,6 +267,7 @@ async def main() -> None:
         await demo_environment_info(client)
         await demo_search_models(client)
         await demo_search_datasets(client)
+        await demo_search_studios(client)
         await demo_search_papers(client)
         await demo_search_mcp_servers(client)
 
diff --git a/src/modelscope_mcp_server/server.py b/src/modelscope_mcp_server/server.py
@@ -18,6 +18,7 @@
 from .tools.mcp import register_mcp_tools
 from .tools.model import register_model_tools
 from .tools.paper import register_paper_tools
+from .tools.studio import register_studio_tools
 from .utils.metadata import get_server_name_with_version
 
 logger = logging.get_logger(__name__)
@@ -42,6 +43,7 @@ def create_mcp_server() -> FastMCP:
     register_context_tools(mcp)
     register_model_tools(mcp)
     register_dataset_tools(mcp)
+    register_studio_tools(mcp)
     register_paper_tools(mcp)
     register_mcp_tools(mcp)
     register_aigc_tools(mcp)
diff --git a/src/modelscope_mcp_server/tools/studio.py b/src/modelscope_mcp_server/tools/studio.py
@@ -0,0 +1,120 @@
+"""ModelScope MCP Server Studio tools.
+
+Provides tools for studio-related operations in the ModelScope MCP Server,
+such as searching for studios and retrieving studio details.
+"""
+
+from typing import Annotated, Literal
+
+from fastmcp import FastMCP
+from fastmcp.utilities import logging
+from pydantic import Field
+
+from ..client import default_client
+from ..settings import settings
+from ..types import Studio
+
+logger = logging.get_logger(__name__)
+
+
+def register_studio_tools(mcp: FastMCP) -> None:
+    """Register all studio-related tools with the MCP server.
+
+    Args:
+        mcp (FastMCP): The MCP server instance
+
+    """
+
+    @mcp.tool(
+        annotations={
+            "title": "Search Studios (创空间 AI 应用)",
+        }
+    )
+    async def search_studios(
+        query: Annotated[
+            str,
+            Field(
+                description="Keyword to search for related studios. "
+                "Leave empty to get all studios based on other filters."
+            ),
+        ] = "",
+        domains: Annotated[
+            list[Literal["multi-modal", "cv", "nlp", "audio", "AutoML"]] | None,
+            Field(description="Domain categories to filter by"),
+        ] = None,
+        sort: Annotated[
+            Literal["Default", "gmt_modified", "VisitsCount", "StarsCount"],
+            Field(description="Sort order"),
+        ] = "Default",
+        limit: Annotated[int, Field(description="Maximum number of studios to return", ge=1, le=30)] = 10,
+    ) -> list[Studio]:
+        """Search for studios on ModelScope."""
+        url = f"{settings.main_domain}/api/v1/dolphin/studios"
+
+        # Build criterion for filters
+        criterion = []
+
+        # Add create_type filter (always include all types)
+        criterion.append(
+            {
+                "category": "create_type",
+                "predicate": "contains",
+                "values": ["interactive", "programmatic"],
+            }
+        )
+
+        # Add domains filter
+        if domains:
+            criterion.append(
+                {
+                    "category": "domains",
+                    "predicate": "contains",
+                    "values": domains,
+                }
+            )
+
+        request_data = {
+            "Name": query,
+            "Criterion": criterion,
+            "SortBy": sort,
+            "PageNumber": 1,
+            "PageSize": limit,
+        }
+
+        response = default_client.put(url, json_data=request_data)
+
+        studios_data = response.get("Data", {}).get("Studios", [])
+
+        studios = []
+        for studio_data in studios_data:
+            path = studio_data.get("Path", "")
+            name = studio_data.get("Name", "")
+            modelscope_url = f"{settings.main_domain}/studios/{path}/{name}"
+
+            if not path or not name:
+                logger.warning(f"Skipping studio with invalid path or name: {studio_data}")
+                continue
+
+            studio = Studio(
+                id=str(studio_data.get("Id", "")),
+                path=path,
+                name=name,
+                chinese_name=studio_data.get("ChineseName", ""),
+                description=studio_data.get("Description", ""),
+                created_by=studio_data.get("CreatedBy", ""),
+                license=studio_data.get("License", ""),
+                modelscope_url=modelscope_url,
+                independent_url=studio_data.get("IndependentUrl"),
+                cover_image=studio_data.get("CoverImage"),
+                type=studio_data.get("Type", ""),
+                status=studio_data.get("Status", ""),
+                domains=studio_data.get("Domain") or [],
+                stars=studio_data.get("Stars", 0),
+                visits=studio_data.get("Visits", 0),
+                created_at=studio_data.get("CreatedTime", 0),
+                updated_at=studio_data.get("LastUpdatedTime", 0),
+                deployed_at=studio_data.get("DeployedTime", 0),
+            )
+            studios.append(studio)
+
+        return studios
diff --git a/src/modelscope_mcp_server/types.py b/src/modelscope_mcp_server/types.py
@@ -79,6 +79,38 @@ class Dataset(BaseModel):
     updated_at: Annotated[int, Field(description="Last updated time (unix timestamp, seconds)")] = 0
 
 
+class Studio(BaseModel):
+    """Studio information."""
+
+    # Basic information
+    id: Annotated[str, Field(description="Unique studio ID")]
+    path: Annotated[str, Field(description="Studio path, for example 'ttwwwaa'")]
+    name: Annotated[str, Field(description="Studio name, for example 'ChatTTS_Speaker'")]
+    chinese_name: Annotated[str, Field(description="Chinese name")]
+    description: Annotated[str, Field(description="Studio description")]
+    created_by: Annotated[str, Field(description="User who created the studio")]
+    license: Annotated[str, Field(description="Open source license")]
+
+    # Links
+    modelscope_url: Annotated[str, Field(description="Detail page URL on ModelScope")]
+    independent_url: Annotated[str | None, Field(description="Independent access URL")] = None
+    cover_image: Annotated[str | None, Field(description="Cover image URL")] = None
+
+    # Classification
+    type: Annotated[str, Field(description="Studio type, for example 'programmatic' or 'interactive'")]
+    status: Annotated[str, Field(description="Current status, for example 'Running'")]
+    domains: Annotated[list[str], Field(description="Domain categories")] = []
+
+    # Metrics
+    stars: Annotated[int, Field(description="Number of stars")] = 0
+    visits: Annotated[int, Field(description="Number of visits")] = 0
+
+    # Timestamps
+    created_at: Annotated[int, Field(description="Created time (unix timestamp, seconds)")] = 0
+    updated_at: Annotated[int, Field(description="Last updated time (unix timestamp, seconds)")] = 0
+    deployed_at: Annotated[int, Field(description="Deployed time (unix timestamp, seconds)")] = 0
+
+
 class Paper(BaseModel):
     """Paper information."""
 
diff --git a/tests/tools/test_search_studios.py b/tests/tools/test_search_studios.py
@@ -0,0 +1,105 @@
+"""测试 search_studios 工具的功能。"""
+
+import pytest
+from fastmcp import Client
+
+
+# Helper functions
+async def search_studios_helper(client, params):
+    """Helper function to search studios and validate basic response structure."""
+    result = await client.call_tool("search_studios", params)
+    assert hasattr(result, "data"), "Result should have data attribute"
+    studios = result.data
+    assert isinstance(studios, list), "Studios should be a list"
+    return studios
+
+
+def print_studio_info(studio, extra_fields=None):
+    """Print studio information with optional extra fields."""
+    base_info = (
+        f"id: {studio.get('id', '')} | "
+        f"name: {studio.get('name', '')} | "
+        f"chinese_name: {studio.get('chinese_name', '')} | "
+        f"type: {studio.get('type', '')}"
+    )
+
+    if extra_fields:
+        for field in extra_fields:
+            base_info += f" | {field}: {studio.get(field, 0)}"
+
+    print(base_info)
+
+
+def print_studios_list(studios, description, extra_fields=None):
+    """Print a list of studios with description and optional extra fields."""
+    print(f"✅ Received {len(studios)} studios {description}:")
+    for studio in studios:
+        print_studio_info(studio, extra_fields)
+
+
+def validate_studio_fields(studio):
+    """Validate that studio has all required fields."""
+    required_fields = [
+        "id",
+        "path",
+        "name",
+        "chinese_name",
+        "description",
+        "created_by",
+        "license",
+        "modelscope_url",
+        "type",
+        "status",
+        "domains",
+        "stars",
+        "visits",
+        "created_at",
+        "updated_at",
+        "deployed_at",
+    ]
+
+    for field in required_fields:
+        assert field in studio, f"Studio should have {field}"
+
+
+@pytest.mark.integration
+async def test_search_studios(mcp_server):
+    async with Client(mcp_server) as client:
+        studios = await search_studios_helper(client, {"query": "ChatTTS", "limit": 5})
+
+        print_studios_list(studios, "", ["stars"])
+
+        assert len(studios) > 0, "Studios should not be empty"
+        validate_studio_fields(studios[0])
+
+
+@pytest.mark.integration
+async def test_search_studios_with_domain_filter(mcp_server):
+    async with Client(mcp_server) as client:
+        studios = await search_studios_helper(client, {"query": "音频", "domains": ["audio"], "limit": 3})
+
+        print_studios_list(studios, "with audio domain filter", ["visits"])
+
+        # Verify that all returned studios have audio domain
+        for studio in studios:
+            assert "audio" in studio.get("domains", []), f"Studio {studio.get('id', '')} should have audio domain"
+
+
+@pytest.mark.integration
+async def test_search_studios_with_domain_filter_cv(mcp_server):
+    async with Client(mcp_server) as client:
+        studios = await search_studios_helper(client, {"domains": ["cv"], "limit": 3})
+
+        print_studios_list(studios, "with cv domain filter", ["predicts"])
+
+        # Verify that all returned studios have cv domain
+        for studio in studios:
+            assert "cv" in studio.get("domains", []), f"Studio {studio.get('id', '')} should have cv domain"
+
+
+@pytest.mark.integration
+async def test_search_studios_sort_by_stars(mcp_server):
+    async with Client(mcp_server) as client:
+        studios = await search_studios_helper(client, {"query": "生成", "sort": "StarsCount", "limit": 3})
+
+        print_studios_list(studios, "sorted by stars", ["stars"])
