try local version

app.py

@@ -1,50 +1,47 @@
 #!/usr/bin/env python3
 """
-MCP Server for Hugging Face Dataset and Model Search API using Gradio
+MCP Server for Hugging Face Dataset and Model Search API
 """
 
-import os
+import asyncio
+import json
 import logging
-from typing import Optional
+import os
+from typing import Any, Dict, List, Optional
+from urllib.parse import urlencode
 
-import gradio as gr
 import httpx
+from mcp.server import Server
+from mcp.server.stdio import stdio_server
+from mcp.types import (
+    Tool,
+    TextContent,
+    CallToolResult,
+    CallToolRequest,
+    ListToolsResult,
+)
 
-# Set up logging
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize HTTP client with longer timeout for MCP usage
-client = httpx.Client(timeout=60.0)  # Increased timeout
-base_url = os.getenv("HF_SEARCH_API_URL", "https://davanstrien-huggingface-datasets-search-v2.hf.space")
-
-
-def search_datasets(
-    query: str,
-    k: int = 5,
-    sort_by: str = "similarity",
-    min_likes: int = 0,
-    min_downloads: int = 0
-) -> str:
-    """
-    Search for datasets using semantic/similarity search based on a text query.
-    
-    This uses AI-powered semantic search to find datasets whose descriptions 
-    are semantically similar to your query, not just keyword matching.
-    
-    Args:
-        query: Search query text (natural language description of what you're looking for)
-        k: Number of results to return (1-100)
-        sort_by: Sort method for results (similarity, likes, downloads, trending)
-        min_likes: Minimum likes filter
-        min_downloads: Minimum downloads filter
-    
-    Returns:
-        Formatted search results with dataset IDs, summaries, and metadata
-    """
-    try:
-        logger.info(f"Searching datasets: query='{query}', k={k}, sort_by='{sort_by}'")
+class HFSearchServer:
+    def __init__(self, base_url: str = "https://davanstrien-huggingface-datasets-search-v2.hf.space"):
+        self.base_url = base_url
+        self.client = httpx.AsyncClient(timeout=60.0)
         
+    async def close(self):
+        await self.client.aclose()
+        
+    async def search_datasets(
+        self,
+        query: str,
+        k: int = 5,
+        sort_by: str = "similarity",
+        min_likes: int = 0,
+        min_downloads: int = 0
+    ) -> Dict[str, Any]:
+        """Search for datasets based on a text query"""
         params = {
             "query": query,
             "k": k,
@@ -53,488 +50,595 @@ def search_datasets(
             "min_downloads": min_downloads
         }
         
-        logger.info(f"Making request to: {base_url}/search/datasets")
-        response = client.get(f"{base_url}/search/datasets", params=params)
+        response = await self.client.get(
+            f"{self.base_url}/search/datasets",
+            params=params
+        )
         response.raise_for_status()
-        data = response.json()
+        return response.json()
+    
+    async def find_similar_datasets(
+        self,
+        dataset_id: str,
+        k: int = 5,
+        sort_by: str = "similarity",
+        min_likes: int = 0,
+        min_downloads: int = 0
+    ) -> Dict[str, Any]:
+        """Find similar datasets to a specified dataset"""
+        params = {
+            "dataset_id": dataset_id,
+            "k": k,
+            "sort_by": sort_by,
+            "min_likes": min_likes,
+            "min_downloads": min_downloads
+        }
         
-        logger.info(f"Successfully retrieved {len(data.get('results', []))} results")
-    except httpx.TimeoutException:
-        logger.error(f"Request timed out for query: {query}")
-        return "Request timed out. The search service may be slow or unavailable."
-    except httpx.HTTPStatusError as e:
-        logger.error(f"HTTP error {e.response.status_code}: {e.response.text}")
-        return f"Search failed with HTTP error {e.response.status_code}: {e.response.text}"
-    except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}")
-        return f"Search failed: {str(e)}"
-    
-    results = data.get("results", [])
-    if not results:
-        return "No datasets found."
-    
-    output = []
-    for i, result in enumerate(results, 1):
-        output.append(f"{i}. **{result['dataset_id']}**")
-        output.append(f"   - Summary: {result['summary']}")
-        output.append(f"   - Similarity: {result['similarity']:.3f}")
-        output.append(f"   - Likes: {result['likes']:,} | Downloads: {result['downloads']:,}")
-        output.append("")
-    
-    return "\n".join(output)
-
-
-def find_similar_datasets(
-    dataset_id: str,
-    k: int = 5,
-    sort_by: str = "similarity",
-    min_likes: int = 0,
-    min_downloads: int = 0
-) -> str:
-    """
-    Find datasets similar to a specified dataset.
-    
-    Args:
-        dataset_id: Dataset ID to find similar datasets for
-        k: Number of results to return (1-100)
-        sort_by: Sort method for results (similarity, likes, downloads, trending)
-        min_likes: Minimum likes filter
-        min_downloads: Minimum downloads filter
-    
-    Returns:
-        Formatted list of similar datasets with metadata
-    """
-    params = {
-        "dataset_id": dataset_id,
-        "k": k,
-        "sort_by": sort_by,
-        "min_likes": min_likes,
-        "min_downloads": min_downloads
-    }
-    
-    response = client.get(f"{base_url}/similarity/datasets", params=params)
-    response.raise_for_status()
-    data = response.json()
-    
-    results = data.get("results", [])
-    if not results:
-        return "No similar datasets found."
-    
-    output = []
-    for i, result in enumerate(results, 1):
-        output.append(f"{i}. **{result['dataset_id']}**")
-        output.append(f"   - Summary: {result['summary']}")
-        output.append(f"   - Similarity: {result['similarity']:.3f}")
-        output.append(f"   - Likes: {result['likes']:,} | Downloads: {result['downloads']:,}")
-        output.append("")
-    
-    return "\n".join(output)
-
-
-def search_models(
-    query: str,
-    k: int = 5,
-    sort_by: str = "similarity",
-    min_likes: int = 0,
-    min_downloads: int = 0,
-    min_param_count: int = 0,
-    max_param_count: Optional[int] = None
-) -> str:
-    """
-    Search for models using semantic/similarity search based on a text query with optional parameter count filtering.
-    
-    This uses AI-powered semantic search to find models whose descriptions
-    are semantically similar to your query, not just keyword matching.
-    
-    Args:
-        query: Search query text (natural language description of what you're looking for)
-        k: Number of results to return (1-100)
-        sort_by: Sort method for results (similarity, likes, downloads, trending)
-        min_likes: Minimum likes filter
-        min_downloads: Minimum downloads filter
-        min_param_count: Minimum parameter count (excludes models with unknown params)
-        max_param_count: Maximum parameter count (None for no limit)
-    
-    Returns:
-        Formatted search results with model IDs, summaries, and metadata
-    """
-    params = {
-        "query": query,
-        "k": k,
-        "sort_by": sort_by,
-        "min_likes": min_likes,
-        "min_downloads": min_downloads,
-        "min_param_count": min_param_count
-    }
-    if max_param_count is not None:
-        params["max_param_count"] = max_param_count
-    
-    response = client.get(f"{base_url}/search/models", params=params)
-    response.raise_for_status()
-    data = response.json()
-    
-    results = data.get("results", [])
-    if not results:
-        return "No models found."
-    
-    output = []
-    for i, result in enumerate(results, 1):
-        output.append(f"{i}. **{result['model_id']}**")
-        output.append(f"   - Summary: {result['summary']}")
-        output.append(f"   - Similarity: {result['similarity']:.3f}")
-        output.append(f"   - Likes: {result['likes']:,} | Downloads: {result['downloads']:,}")
-        if result.get('param_count') is not None and result['param_count'] > 0:
-            # Format parameter count nicely
-            param_count = result['param_count']
-            if param_count >= 1_000_000_000:
-                param_str = f"{param_count / 1_000_000_000:.1f}B"
-            elif param_count >= 1_000_000:
-                param_str = f"{param_count / 1_000_000:.1f}M"
-            elif param_count >= 1_000:
-                param_str = f"{param_count / 1_000:.1f}K"
-            else:
-                param_str = str(param_count)
-            output.append(f"   - Parameters: {param_str}")
-        output.append("")
-    
-    return "\n".join(output)
-
-
-def find_similar_models(
-    model_id: str,
-    k: int = 5,
-    sort_by: str = "similarity",
-    min_likes: int = 0,
-    min_downloads: int = 0,
-    min_param_count: int = 0,
-    max_param_count: Optional[int] = None
-) -> str:
-    """
-    Find models similar to a specified model.
-    
-    Args:
-        model_id: Model ID to find similar models for
-        k: Number of results to return (1-100)
-        sort_by: Sort method for results (similarity, likes, downloads, trending)
-        min_likes: Minimum likes filter
-        min_downloads: Minimum downloads filter
-        min_param_count: Minimum parameter count (excludes models with unknown params)
-        max_param_count: Maximum parameter count (None for no limit)
-    
-    Returns:
-        Formatted list of similar models with metadata
-    """
-    params = {
-        "model_id": model_id,
-        "k": k,
-        "sort_by": sort_by,
-        "min_likes": min_likes,
-        "min_downloads": min_downloads,
-        "min_param_count": min_param_count
-    }
-    if max_param_count is not None:
-        params["max_param_count"] = max_param_count
-    
-    response = client.get(f"{base_url}/similarity/models", params=params)
-    response.raise_for_status()
-    data = response.json()
-    
-    results = data.get("results", [])
-    if not results:
-        return "No similar models found."
-    
-    output = []
-    for i, result in enumerate(results, 1):
-        output.append(f"{i}. **{result['model_id']}**")
-        output.append(f"   - Summary: {result['summary']}")
-        output.append(f"   - Similarity: {result['similarity']:.3f}")
-        output.append(f"   - Likes: {result['likes']:,} | Downloads: {result['downloads']:,}")
-        if result.get('param_count') is not None and result['param_count'] > 0:
-            # Format parameter count nicely
-            param_count = result['param_count']
-            if param_count >= 1_000_000_000:
-                param_str = f"{param_count / 1_000_000_000:.1f}B"
-            elif param_count >= 1_000_000:
-                param_str = f"{param_count / 1_000_000:.1f}M"
-            elif param_count >= 1_000:
-                param_str = f"{param_count / 1_000:.1f}K"
-            else:
-                param_str = str(param_count)
-            output.append(f"   - Parameters: {param_str}")
-        output.append("")
-    
-    return "\n".join(output)
-
-
-def get_trending_models(
-    limit: int = 10,
-    min_likes: int = 0,
-    min_downloads: int = 0,
-    min_param_count: int = 0,
-    max_param_count: Optional[int] = None
-) -> str:
-    """
-    Get trending models with their summaries and optional filtering.
-    
-    Args:
-        limit: Number of results to return (1-100)
-        min_likes: Minimum likes filter
-        min_downloads: Minimum downloads filter
-        min_param_count: Minimum parameter count (excludes models with unknown params)
-        max_param_count: Maximum parameter count (None for no limit)
-    
-    Returns:
-        Formatted list of trending models with metadata
-    """
-    params = {
-        "limit": limit,
-        "min_likes": min_likes,
-        "min_downloads": min_downloads,
-        "min_param_count": min_param_count
-    }
-    if max_param_count is not None:
-        params["max_param_count"] = max_param_count
-    
-    response = client.get(f"{base_url}/trending/models", params=params)
-    response.raise_for_status()
-    data = response.json()
-    
-    results = data.get("results", [])
-    if not results:
-        return "No trending models found."
-    
-    output = []
-    for i, result in enumerate(results, 1):
-        output.append(f"{i}. **{result['model_id']}**")
-        output.append(f"   - Summary: {result['summary']}")
-        output.append(f"   - Similarity: {result['similarity']:.3f}")
-        output.append(f"   - Likes: {result['likes']:,} | Downloads: {result['downloads']:,}")
-        if result.get('param_count') is not None and result['param_count'] > 0:
-            # Format parameter count nicely
-            param_count = result['param_count']
-            if param_count >= 1_000_000_000:
-                param_str = f"{param_count / 1_000_000_000:.1f}B"
-            elif param_count >= 1_000_000:
-                param_str = f"{param_count / 1_000_000:.1f}M"
-            elif param_count >= 1_000:
-                param_str = f"{param_count / 1_000:.1f}K"
-            else:
-                param_str = str(param_count)
-            output.append(f"   - Parameters: {param_str}")
-        output.append("")
-    
-    return "\n".join(output)
-
-
-def get_trending_datasets(
-    limit: int = 10,
-    min_likes: int = 0,
-    min_downloads: int = 0
-) -> str:
-    """
-    Get trending datasets with their summaries.
-    
-    Args:
-        limit: Number of results to return (1-100)
-        min_likes: Minimum likes filter
-        min_downloads: Minimum downloads filter
-    
-    Returns:
-        Formatted list of trending datasets with metadata
-    """
-    params = {
-        "limit": limit,
-        "min_likes": min_likes,
-        "min_downloads": min_downloads
-    }
-    
-    response = client.get(f"{base_url}/trending/datasets", params=params)
-    response.raise_for_status()
-    data = response.json()
-    
-    results = data.get("results", [])
-    if not results:
-        return "No trending datasets found."
-    
-    output = []
-    for i, result in enumerate(results, 1):
-        output.append(f"{i}. **{result['dataset_id']}**")
-        output.append(f"   - Summary: {result['summary']}")
-        output.append(f"   - Similarity: {result['similarity']:.3f}")
-        output.append(f"   - Likes: {result['likes']:,} | Downloads: {result['downloads']:,}")
-        output.append("")
-    
-    return "\n".join(output)
-
-
-def download_model_card(model_id: str) -> str:
-    """
-    Download the README card for a HuggingFace model.
-    
-    Args:
-        model_id: The model ID (e.g., 'username/model-name')
-    
-    Returns:
-        The content of the model card (README.md)
-    """
-    url = f"https://huggingface.co/{model_id}/raw/main/README.md"
-    response = client.get(url)
-    response.raise_for_status()
-    return response.text
-
-
-def download_dataset_card(dataset_id: str) -> str:
-    """
-    Download the README card for a HuggingFace dataset.
-    
-    Args:
-        dataset_id: The dataset ID (e.g., 'username/dataset-name')
-    
-    Returns:
-        The content of the dataset card (README.md)
-    """
-    url = f"https://huggingface.co/datasets/{dataset_id}/raw/main/README.md"
-    response = client.get(url)
-    response.raise_for_status()
-    return response.text
-
-
-# Create Gradio interface
-with gr.Blocks(title="HuggingFace Search MCP Server") as demo:
-    gr.Markdown("# HuggingFace Search MCP Server")
-    gr.Markdown("This server provides semantic search capabilities for HuggingFace models and datasets.")
-    gr.Markdown(f"**Backend API:** {base_url}")
-    
-    with gr.Tab("Search Datasets"):
-        gr.Interface(
-            fn=search_datasets,
-            inputs=[
-                gr.Textbox(label="Query", placeholder="Enter search query"),
-                gr.Slider(1, 100, value=5, step=1, label="Number of results"),
-                gr.Dropdown(["similarity", "likes", "downloads", "trending"], value="similarity", label="Sort by"),
-                gr.Number(value=0, label="Minimum likes"),
-                gr.Number(value=0, label="Minimum downloads")
-            ],
-            outputs=gr.Markdown(label="Results"),
-            title="Search Datasets",
-            description="Search for datasets based on a text query"
-        )
-    
-    with gr.Tab("Find Similar Datasets"):
-        gr.Interface(
-            fn=find_similar_datasets,
-            inputs=[
-                gr.Textbox(label="Dataset ID", placeholder="username/dataset-name"),
-                gr.Slider(1, 100, value=5, step=1, label="Number of results"),
-                gr.Dropdown(["similarity", "likes", "downloads", "trending"], value="similarity", label="Sort by"),
-                gr.Number(value=0, label="Minimum likes"),
-                gr.Number(value=0, label="Minimum downloads")
-            ],
-            outputs=gr.Markdown(label="Results"),
-            title="Find Similar Datasets",
-            description="Find datasets similar to a specified dataset"
+        response = await self.client.get(
+            f"{self.base_url}/similarity/datasets",
+            params=params
         )
-    
-    with gr.Tab("Search Models"):
-        gr.Interface(
-            fn=search_models,
-            inputs=[
-                gr.Textbox(label="Query", placeholder="Enter search query"),
-                gr.Slider(1, 100, value=5, step=1, label="Number of results"),
-                gr.Dropdown(["similarity", "likes", "downloads", "trending"], value="similarity", label="Sort by"),
-                gr.Number(value=0, label="Minimum likes"),
-                gr.Number(value=0, label="Minimum downloads"),
-                gr.Number(value=0, label="Minimum parameter count"),
-                gr.Number(value=None, label="Maximum parameter count (leave empty for no limit)")
-            ],
-            outputs=gr.Markdown(label="Results"),
-            title="Search Models",
-            description="Search for models based on a text query with optional parameter count filtering"
+        response.raise_for_status()
+        return response.json()
+    
+    async def search_models(
+        self,
+        query: str,
+        k: int = 5,
+        sort_by: str = "similarity",
+        min_likes: int = 0,
+        min_downloads: int = 0,
+        min_param_count: int = 0,
+        max_param_count: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Search for models based on a text query"""
+        params = {
+            "query": query,
+            "k": k,
+            "sort_by": sort_by,
+            "min_likes": min_likes,
+            "min_downloads": min_downloads,
+            "min_param_count": min_param_count
+        }
+        if max_param_count is not None:
+            params["max_param_count"] = max_param_count
+            
+        response = await self.client.get(
+            f"{self.base_url}/search/models",
+            params=params
         )
-    
-    with gr.Tab("Find Similar Models"):
-        gr.Interface(
-            fn=find_similar_models,
-            inputs=[
-                gr.Textbox(label="Model ID", placeholder="username/model-name"),
-                gr.Slider(1, 100, value=5, step=1, label="Number of results"),
-                gr.Dropdown(["similarity", "likes", "downloads", "trending"], value="similarity", label="Sort by"),
-                gr.Number(value=0, label="Minimum likes"),
-                gr.Number(value=0, label="Minimum downloads"),
-                gr.Number(value=0, label="Minimum parameter count"),
-                gr.Number(value=None, label="Maximum parameter count (leave empty for no limit)")
-            ],
-            outputs=gr.Markdown(label="Results"),
-            title="Find Similar Models",
-            description="Find models similar to a specified model"
+        response.raise_for_status()
+        return response.json()
+    
+    async def find_similar_models(
+        self,
+        model_id: str,
+        k: int = 5,
+        sort_by: str = "similarity",
+        min_likes: int = 0,
+        min_downloads: int = 0,
+        min_param_count: int = 0,
+        max_param_count: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Find similar models to a specified model"""
+        params = {
+            "model_id": model_id,
+            "k": k,
+            "sort_by": sort_by,
+            "min_likes": min_likes,
+            "min_downloads": min_downloads,
+            "min_param_count": min_param_count
+        }
+        if max_param_count is not None:
+            params["max_param_count"] = max_param_count
+            
+        response = await self.client.get(
+            f"{self.base_url}/similarity/models",
+            params=params
         )
-    
-    with gr.Tab("Trending Models"):
-        gr.Interface(
-            fn=get_trending_models,
-            inputs=[
-                gr.Slider(1, 100, value=10, step=1, label="Number of results"),
-                gr.Number(value=0, label="Minimum likes"),
-                gr.Number(value=0, label="Minimum downloads"),
-                gr.Number(value=0, label="Minimum parameter count"),
-                gr.Number(value=None, label="Maximum parameter count (leave empty for no limit)")
-            ],
-            outputs=gr.Markdown(label="Results"),
-            title="Get Trending Models",
-            description="Get trending models with their summaries and optional filtering"
+        response.raise_for_status()
+        return response.json()
+    
+    async def get_trending_models(
+        self,
+        limit: int = 10,
+        min_likes: int = 0,
+        min_downloads: int = 0,
+        min_param_count: int = 0,
+        max_param_count: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Get trending models with their summaries"""
+        params = {
+            "limit": limit,
+            "min_likes": min_likes,
+            "min_downloads": min_downloads,
+            "min_param_count": min_param_count
+        }
+        if max_param_count is not None:
+            params["max_param_count"] = max_param_count
+            
+        response = await self.client.get(
+            f"{self.base_url}/trending/models",
+            params=params
         )
-    
-    with gr.Tab("Trending Datasets"):
-        gr.Interface(
-            fn=get_trending_datasets,
-            inputs=[
-                gr.Slider(1, 100, value=10, step=1, label="Number of results"),
-                gr.Number(value=0, label="Minimum likes"),
-                gr.Number(value=0, label="Minimum downloads")
-            ],
-            outputs=gr.Markdown(label="Results"),
-            title="Get Trending Datasets",
-            description="Get trending datasets with their summaries"
+        response.raise_for_status()
+        return response.json()
+    
+    async def get_trending_datasets(
+        self,
+        limit: int = 10,
+        min_likes: int = 0,
+        min_downloads: int = 0
+    ) -> Dict[str, Any]:
+        """Get trending datasets with their summaries"""
+        params = {
+            "limit": limit,
+            "min_likes": min_likes,
+            "min_downloads": min_downloads
+        }
+        
+        response = await self.client.get(
+            f"{self.base_url}/trending/datasets",
+            params=params
         )
+        response.raise_for_status()
+        return response.json()
     
-    with gr.Tab("Download Model Card"):
-        gr.Interface(
-            fn=download_model_card,
-            inputs=gr.Textbox(label="Model ID", placeholder="username/model-name"),
-            outputs=gr.Textbox(label="Model Card Content", lines=20),
-            title="Download Model Card",
-            description="Download the README card for a HuggingFace model"
-        )
+    async def download_model_card(self, model_id: str) -> str:
+        """Download the README card for a HuggingFace model"""
+        url = f"https://huggingface.co/{model_id}/raw/main/README.md"
+        response = await self.client.get(url)
+        response.raise_for_status()
+        return response.text
     
-    with gr.Tab("Download Dataset Card"):
-        gr.Interface(
-            fn=download_dataset_card,
-            inputs=gr.Textbox(label="Dataset ID", placeholder="username/dataset-name"),
-            outputs=gr.Textbox(label="Dataset Card Content", lines=20),
-            title="Download Dataset Card",
-            description="Download the README card for a HuggingFace dataset"
-        )
+    async def download_dataset_card(self, dataset_id: str) -> str:
+        """Download the README card for a HuggingFace dataset"""
+        url = f"https://huggingface.co/datasets/{dataset_id}/raw/main/README.md"
+        response = await self.client.get(url)
+        response.raise_for_status()
+        return response.text
 
+# Initialize server and API client
+server = Server("hf-search")
+api_client: Optional[HFSearchServer] = None
 
-if __name__ == "__main__":
-    # Launch with MCP server enabled
-    # Try environment variable approach as workaround for initialization race condition
-    import os
-    os.environ["GRADIO_MCP_SERVER"] = "true"
+@server.list_tools()
+async def list_tools() -> ListToolsResult:
+    """List available tools"""
+    return ListToolsResult(
+        tools=[
+            Tool(
+                name="search_datasets",
+                description="Search for datasets using semantic/similarity search based on a text query",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "Search query text (natural language description of what you're looking for)"
+                        },
+                        "k": {
+                            "type": "integer",
+                            "description": "Number of results to return (1-100)",
+                            "minimum": 1,
+                            "maximum": 100,
+                            "default": 5
+                        },
+                        "sort_by": {
+                            "type": "string",
+                            "description": "Sort method for results",
+                            "enum": ["similarity", "likes", "downloads", "trending"],
+                            "default": "similarity"
+                        },
+                        "min_likes": {
+                            "type": "integer",
+                            "description": "Minimum likes filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_downloads": {
+                            "type": "integer",
+                            "description": "Minimum downloads filter",
+                            "minimum": 0,
+                            "default": 0
+                        }
+                    },
+                    "required": ["query"]
+                }
+            ),
+            Tool(
+                name="find_similar_datasets",
+                description="Find datasets similar to a specified dataset",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "dataset_id": {
+                            "type": "string",
+                            "description": "Dataset ID to find similar datasets for"
+                        },
+                        "k": {
+                            "type": "integer",
+                            "description": "Number of results to return (1-100)",
+                            "minimum": 1,
+                            "maximum": 100,
+                            "default": 5
+                        },
+                        "sort_by": {
+                            "type": "string",
+                            "description": "Sort method for results",
+                            "enum": ["similarity", "likes", "downloads", "trending"],
+                            "default": "similarity"
+                        },
+                        "min_likes": {
+                            "type": "integer",
+                            "description": "Minimum likes filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_downloads": {
+                            "type": "integer",
+                            "description": "Minimum downloads filter",
+                            "minimum": 0,
+                            "default": 0
+                        }
+                    },
+                    "required": ["dataset_id"]
+                }
+            ),
+            Tool(
+                name="search_models",
+                description="Search for models using semantic/similarity search based on a text query with optional parameter count filtering",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "Search query text (natural language description of what you're looking for)"
+                        },
+                        "k": {
+                            "type": "integer",
+                            "description": "Number of results to return (1-100)",
+                            "minimum": 1,
+                            "maximum": 100,
+                            "default": 5
+                        },
+                        "sort_by": {
+                            "type": "string",
+                            "description": "Sort method for results",
+                            "enum": ["similarity", "likes", "downloads", "trending"],
+                            "default": "similarity"
+                        },
+                        "min_likes": {
+                            "type": "integer",
+                            "description": "Minimum likes filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_downloads": {
+                            "type": "integer",
+                            "description": "Minimum downloads filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_param_count": {
+                            "type": "integer",
+                            "description": "Minimum parameter count (excludes models with unknown params)",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "max_param_count": {
+                            "type": ["integer", "null"],
+                            "description": "Maximum parameter count (null for no limit)",
+                            "minimum": 0,
+                            "default": None
+                        }
+                    },
+                    "required": ["query"]
+                }
+            ),
+            Tool(
+                name="find_similar_models",
+                description="Find models similar to a specified model",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "model_id": {
+                            "type": "string",
+                            "description": "Model ID to find similar models for"
+                        },
+                        "k": {
+                            "type": "integer",
+                            "description": "Number of results to return (1-100)",
+                            "minimum": 1,
+                            "maximum": 100,
+                            "default": 5
+                        },
+                        "sort_by": {
+                            "type": "string",
+                            "description": "Sort method for results",
+                            "enum": ["similarity", "likes", "downloads", "trending"],
+                            "default": "similarity"
+                        },
+                        "min_likes": {
+                            "type": "integer",
+                            "description": "Minimum likes filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_downloads": {
+                            "type": "integer",
+                            "description": "Minimum downloads filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_param_count": {
+                            "type": "integer",
+                            "description": "Minimum parameter count (excludes models with unknown params)",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "max_param_count": {
+                            "type": ["integer", "null"],
+                            "description": "Maximum parameter count (null for no limit)",
+                            "minimum": 0,
+                            "default": None
+                        }
+                    },
+                    "required": ["model_id"]
+                }
+            ),
+            Tool(
+                name="get_trending_models",
+                description="Get trending models with their summaries and optional filtering",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "limit": {
+                            "type": "integer",
+                            "description": "Number of results to return (1-100)",
+                            "minimum": 1,
+                            "maximum": 100,
+                            "default": 10
+                        },
+                        "min_likes": {
+                            "type": "integer",
+                            "description": "Minimum likes filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_downloads": {
+                            "type": "integer",
+                            "description": "Minimum downloads filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_param_count": {
+                            "type": "integer",
+                            "description": "Minimum parameter count (excludes models with unknown params)",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "max_param_count": {
+                            "type": ["integer", "null"],
+                            "description": "Maximum parameter count (null for no limit)",
+                            "minimum": 0,
+                            "default": None
+                        }
+                    }
+                }
+            ),
+            Tool(
+                name="get_trending_datasets",
+                description="Get trending datasets with their summaries",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "limit": {
+                            "type": "integer",
+                            "description": "Number of results to return (1-100)",
+                            "minimum": 1,
+                            "maximum": 100,
+                            "default": 10
+                        },
+                        "min_likes": {
+                            "type": "integer",
+                            "description": "Minimum likes filter",
+                            "minimum": 0,
+                            "default": 0
+                        },
+                        "min_downloads": {
+                            "type": "integer",
+                            "description": "Minimum downloads filter",
+                            "minimum": 0,
+                            "default": 0
+                        }
+                    }
+                }
+            ),
+            Tool(
+                name="download_model_card",
+                description="Download the README card for a HuggingFace model",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "model_id": {
+                            "type": "string",
+                            "description": "The model ID (e.g., 'username/model-name')"
+                        }
+                    },
+                    "required": ["model_id"]
+                }
+            ),
+            Tool(
+                name="download_dataset_card",
+                description="Download the README card for a HuggingFace dataset",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "dataset_id": {
+                            "type": "string",
+                            "description": "The dataset ID (e.g., 'username/dataset-name')"
+                        }
+                    },
+                    "required": ["dataset_id"]
+                }
+            )
+        ]
+    )
+
+@server.call_tool()
+async def call_tool(request: CallToolRequest) -> CallToolResult:
+    """Handle tool calls"""
+    global api_client
+    
+    if api_client is None:
+        # Initialize API client with base URL from environment or default
+        base_url = os.getenv("HF_SEARCH_API_URL", "https://davanstrien-huggingface-datasets-search-v2.hf.space")
+        api_client = HFSearchServer(base_url)
     
     try:
-        # First try with environment variable only
-        demo.launch(
-            server_name="0.0.0.0", 
-            server_port=7860,
-            show_error=True
+        # Parse arguments
+        args = request.params.arguments if hasattr(request.params, 'arguments') else {}
+        
+        # Format results helper
+        def format_dataset_results(data: Dict[str, Any]) -> str:
+            results = data.get("results", [])
+            if not results:
+                return "No datasets found."
+            
+            output = []
+            for i, result in enumerate(results, 1):
+                output.append(f"{i}. **{result['dataset_id']}**")
+                output.append(f"   - Summary: {result['summary']}")
+                output.append(f"   - Similarity: {result['similarity']:.3f}")
+                output.append(f"   - Likes: {result['likes']:,} | Downloads: {result['downloads']:,}")
+                output.append("")
+            
+            return "\n".join(output)
+        
+        def format_model_results(data: Dict[str, Any]) -> str:
+            results = data.get("results", [])
+            if not results:
+                return "No models found."
+            
+            output = []
+            for i, result in enumerate(results, 1):
+                output.append(f"{i}. **{result['model_id']}**")
+                output.append(f"   - Summary: {result['summary']}")
+                output.append(f"   - Similarity: {result['similarity']:.3f}")
+                output.append(f"   - Likes: {result['likes']:,} | Downloads: {result['downloads']:,}")
+                if result.get('param_count') is not None and result['param_count'] > 0:
+                    # Format parameter count nicely
+                    param_count = result['param_count']
+                    if param_count >= 1_000_000_000:
+                        param_str = f"{param_count / 1_000_000_000:.1f}B"
+                    elif param_count >= 1_000_000:
+                        param_str = f"{param_count / 1_000_000:.1f}M"
+                    elif param_count >= 1_000:
+                        param_str = f"{param_count / 1_000:.1f}K"
+                    else:
+                        param_str = str(param_count)
+                    output.append(f"   - Parameters: {param_str}")
+                output.append("")
+            
+            return "\n".join(output)
+        
+        # Route to appropriate method
+        if request.params.name == "search_datasets":
+            result = await api_client.search_datasets(**args)
+            formatted = format_dataset_results(result)
+            return CallToolResult(
+                content=[TextContent(text=formatted)],
+                isError=False
+            )
+            
+        elif request.params.name == "find_similar_datasets":
+            result = await api_client.find_similar_datasets(**args)
+            formatted = format_dataset_results(result)
+            return CallToolResult(
+                content=[TextContent(text=formatted)],
+                isError=False
+            )
+            
+        elif request.params.name == "search_models":
+            result = await api_client.search_models(**args)
+            formatted = format_model_results(result)
+            return CallToolResult(
+                content=[TextContent(text=formatted)],
+                isError=False
+            )
+            
+        elif request.params.name == "find_similar_models":
+            result = await api_client.find_similar_models(**args)
+            formatted = format_model_results(result)
+            return CallToolResult(
+                content=[TextContent(text=formatted)],
+                isError=False
+            )
+            
+        elif request.params.name == "get_trending_models":
+            result = await api_client.get_trending_models(**args)
+            formatted = format_model_results(result)
+            return CallToolResult(
+                content=[TextContent(text=formatted)],
+                isError=False
+            )
+            
+        elif request.params.name == "get_trending_datasets":
+            result = await api_client.get_trending_datasets(**args)
+            formatted = format_dataset_results(result)
+            return CallToolResult(
+                content=[TextContent(text=formatted)],
+                isError=False
+            )
+            
+        elif request.params.name == "download_model_card":
+            result = await api_client.download_model_card(**args)
+            return CallToolResult(
+                content=[TextContent(text=result)],
+                isError=False
+            )
+            
+        elif request.params.name == "download_dataset_card":
+            result = await api_client.download_dataset_card(**args)
+            return CallToolResult(
+                content=[TextContent(text=result)],
+                isError=False
+            )
+            
+        else:
+            return CallToolResult(
+                content=[TextContent(text=f"Unknown tool: {request.params.name}")],
+                isError=True
+            )
+            
+    except httpx.HTTPStatusError as e:
+        error_msg = f"API request failed with status {e.response.status_code}: {e.response.text}"
+        logger.error(error_msg)
+        return CallToolResult(
+            content=[TextContent(text=error_msg)],
+            isError=True
         )
     except Exception as e:
-        logger.error(f"Environment variable approach failed: {e}")
-        # Fallback to parameter approach
-        demo.launch(
-            mcp_server=True,
-            server_name="0.0.0.0", 
-            server_port=7860,
-            show_error=True
-        )
+        error_msg = f"Error calling tool {request.params.name}: {str(e)}"
+        logger.error(error_msg, exc_info=True)
+        return CallToolResult(
+            content=[TextContent(text=error_msg)],
+            isError=True
+        )
+
+async def main():
+    """Main entry point"""
+    async with stdio_server() as (read_stream, write_stream):
+        await server.run(
+            read_stream, 
+            write_stream,
+            server.create_initialization_options()
+        )
+        
+        # Cleanup
+        if api_client:
+            await api_client.close()
+
+if __name__ == "__main__":
+    asyncio.run(main())