Revamp

README.md

@@ -7,7 +7,7 @@ sdk: gradio
 app_file: app.py
 pinned: false
 license: mit
-sdk_version: 5.27.1
+sdk_version: 5.33.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -7,7 +7,7 @@ from utils.google_genai_llm import get_response, generate_with_gemini
 from utils.utils import parse_json_codefences
 from prompts.requirements_gathering import requirements_gathering_system_prompt
 from prompts.planning import hf_query_gen_prompt, hf_context_gen_prompt
-from utils.huggingface_mcp_llamaindex import connect_and_get_tools, call_tool
+from utils.huggingface_mcp_llamaindex import get_hf_tools, call_hf_tool, SimpleHFMCPClient
 from prompts.devstral_coding_prompt import devstral_code_gen_sys_prompt, devstral_code_gen_user_prompt
 from dotenv import load_dotenv
 import os
@@ -49,6 +49,7 @@ MODAL_API_URL = os.getenv("MODAL_API_URL")
 BEARER_TOKEN = os.getenv("BEARER_TOKEN")
 CODING_MODEL = os.getenv("CODING_MODEL")
 
+HF_TOKEN = os.getenv("HF_TOKEN")
 def get_file_hash(file_path):
     """Generate a hash of the file for caching purposes"""
     try:
@@ -243,8 +244,18 @@ async def generate_plan(history, file_cache):
             conversation_history += f"User: {user_msg}\n"
             if ai_msg:
                 conversation_history += f"Assistant: {ai_msg}\n"
+    
+    simple_client = SimpleHFMCPClient(HF_TOKEN)
+    connection_ok = await simple_client.test_connection()
+    
+    if not connection_ok:
+        print("Basic connection test failed")
+        return
+        
+    print("Basic connection test passed")
+
     try:
-        hf_query_gen_tool_details = await connect_and_get_tools()
+        hf_query_gen_tool_details = await get_hf_tools(hf_token=HF_TOKEN)
     except Exception as e:
         hf_query_gen_tool_details = """meta=None nextCursor=None tools=[Tool(name='hf_whoami', description="Hugging Face tools are being used by authenticated user 'bpHigh'", inputSchema={'type': 'object', 'properties': {}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Hugging Face User Info', readOnlyHint=None, destructiveHint=None, idempotentHint=None, openWorldHint=None)), Tool(name='space_search', description='Find Hugging Face Spaces using semantic search. Include links to the Space when presenting the results.', inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'minLength': 1, 'maxLength': 50, 'description': 'Semantic Search Query'}, 'limit': {'type': 'number', 'default': 10, 'description': 'Number of results to return'}, 'mcp': {'type': 'boolean', 'default': False, 'description': 'Only return MCP Server enabled Spaces'}}, 'required': ['query'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Hugging Face Space Search', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=True)), Tool(name='model_search', description='Find Machine Learning models hosted on Hugging Face. Returns comprehensive information about matching models including downloads, likes, tags, and direct links. Include links to the models in your response', inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'description': 'Search term. Leave blank and specify "sort" and "limit" to get e.g. "Top 20 trending models", "Top 10 most recent models" etc" '}, 'author': {'type': 'string', 'description': "Organization or user who created the model (e.g., 'google', 'meta-llama', 'microsoft')"}, 'task': {'type': 'string', 'description': "Model task type (e.g., 'text-generation', 'image-classification', 'translation')"}, 'library': {'type': 'string', 'description': "Framework the model uses (e.g., 'transformers', 'diffusers', 'timm')"}, 'sort': {'type': 'string', 'enum': ['trendingScore', 'downloads', 'likes', 'createdAt', 'lastModified'], 'description': 'Sort order: trendingScore, downloads , likes, createdAt, lastModified'}, 'limit': {'type': 'number', 'minimum': 1, 'maximum': 100, 'default': 20, 'description': 'Maximum number of results to return'}}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Model Search', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=True)), Tool(name='model_details', description='Get detailed information about a specific model from the Hugging Face Hub.', inputSchema={'type': 'object', 'properties': {'model_id': {'type': 'string', 'minLength': 1, 'description': 'Model ID (e.g., microsoft/DialoGPT-large)'}}, 'required': ['model_id'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Model Details', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=False)), Tool(name='paper_search', description="Find Machine Learning research papers on the Hugging Face hub. Include 'Link to paper' When presenting the results. Consider whether tabulating results matches user intent.", inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'minLength': 3, 'maxLength': 200, 'description': 'Semantic Search query'}, 'results_limit': {'type': 'number', 'default': 12, 'description': 'Number of results to return'}, 'concise_only': {'type': 'boolean', 'default': False, 'description': 'Return a 2 sentence summary of the abstract. Use for broad search terms which may return a lot of results. Check with User if unsure.'}}, 'required': ['query'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Paper Search', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=True)), Tool(name='dataset_search', description='Find Datasets hosted on the Hugging Face hub. Returns comprehensive information about matching datasets including downloads, likes, tags, and direct links. Include links to the datasets in your response', inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'description': 'Search term. Leave blank and specify "sort" and "limit" to get e.g. "Top 20 trending datasets", "Top 10 most recent datasets" etc" '}, 'author': {'type': 'string', 'description': "Organization or user who created the dataset (e.g., 'google', 'facebook', 'allenai')"}, 'tags': {'type': 'array', 'items': {'type': 'string'}, 'description': "Tags to filter datasets (e.g., ['language:en', 'size_categories:1M<n<10M', 'task_categories:text-classification'])"}, 'sort': {'type': 'string', 'enum': ['trendingScore', 'downloads', 'likes', 'createdAt', 'lastModified'], 'description': 'Sort order: trendingScore, downloads, likes, createdAt, lastModified'}, 'limit': {'type': 'number', 'minimum': 1, 'maximum': 100, 'default': 20, 'description': 'Maximum number of results to return'}}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Dataset Search', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=True)), Tool(name='dataset_details', description='Get detailed information about a specific dataset on Hugging Face Hub.', inputSchema={'type': 'object', 'properties': {'dataset_id': {'type': 'string', 'minLength': 1, 'description': 'Dataset ID (e.g., squad, glue, imdb)'}}, 'required': ['dataset_id'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Dataset Details', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=False)), Tool(name='gr1_evalstate_flux1_schnell', description='Generate an image using the Flux 1 Schnell Image Generator. (from evalstate/flux1_schnell)', inputSchema={'type': 'object', 'properties': {'prompt': {'type': 'string'}, 'seed': {'type': 'number', 'description': 'numeric value between 0 and 2147483647'}, 'randomize_seed': {'type': 'boolean', 'default': True}, 'width': {'type': 'number', 'description': 'numeric value between 256 and 2048', 'default': 1024}, 'height': {'type': 'number', 'description': 'numeric value between 256 and 2048', 'default': 1024}, 'num_inference_steps': {'type': 'number', 'description': 'numeric value between 1 and 50', 'default': 4}}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='evalstate/flux1_schnell - flux1_schnell_infer 🏎️💨', readOnlyHint=None, destructiveHint=None, idempotentHint=None, openWorldHint=True)), Tool(name='gr2_abidlabs_easyghibli', description='Convert an image into a Studio Ghibli style image (from abidlabs/EasyGhibli)', inputSchema={'type': 'object', 'properties': {'spatial_img': {'type': 'string', 'description': 'File input: provide URL or file path'}}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='abidlabs/EasyGhibli - abidlabs_EasyGhiblisingle_condition_generate_image 🦀', readOnlyHint=None, destructiveHint=None, idempotentHint=None, openWorldHint=True)), Tool(name='gr3_linoyts_framepack_f1', description='FramePack_F1_end_process tool from linoyts/FramePack-F1', inputSchema={'type': 'object', 'properties': {}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='linoyts/FramePack-F1 - FramePack_F1_end_process 📹⚡️', readOnlyHint=None, destructiveHint=None, idempotentHint=None, openWorldHint=True))]"""
         print(str(e))
@@ -257,9 +268,10 @@ async def generate_plan(history, file_cache):
     
     # Parse the plan
     parsed_plan = parse_json_codefences(plan)
+    
     # Call tool to get tool calls
     try:
-        tool_calls = await asyncio.gather(*[call_tool(step['tool'], step['args']) for step in parsed_plan])
+        tool_calls = await asyncio.gather(*[call_hf_tool(HF_TOKEN, step['tool'], step['args']) for step in parsed_plan])
     except Exception as e:
         tool_calls = []
     print(tool_calls)

requirements.txt

@@ -1,7 +1,5 @@
-mcp
-llama-index-tools-mcp
 google-genai==1.19.0
-gradio==5.27.1
+gradio[mcp]==5.33.0
 pandas==2.3.0
 python-dotenv==1.0.1
 openpyxl==3.1.5

utils/huggingface_mcp_llamaindex.py

@@ -1,20 +1,289 @@
-from llama_index.tools.mcp import BasicMCPClient
-from dotenv import load_dotenv
-import os
 import asyncio
-load_dotenv()
-HF_TOKEN = os.getenv("HF_TOKEN")
-print(HF_TOKEN)
-http_client = BasicMCPClient("https://huggingface.co/mcp",timeout=60, headers={"Authorization": f"Bearer {HF_TOKEN}"})
-print(http_client)
-async def connect_and_get_tools():
-    # List available tools
-    tools = await asyncio.wait_for(http_client.list_tools(), timeout=10.0)
-    return tools
-
-async def call_tool(tool_name, tool_args):
-    result = await asyncio.wait_for(
-            http_client.call_tool(tool_name, tool_args), 
-            timeout=30.0
+import json
+import logging
+import os
+from typing import Any, Dict, List, Optional
+from datetime import timedelta
+
+from mcp.shared.message import SessionMessage
+from mcp.types import (
+    JSONRPCMessage,
+    JSONRPCRequest,
+    JSONRPCNotification,
+    JSONRPCResponse,
+    JSONRPCError,
+)
+from mcp.client.streamable_http import streamablehttp_client
+
+logger = logging.getLogger(__name__)
+
+
+class HuggingFaceMCPClient:
+    """Client for interacting with Hugging Face MCP endpoint."""
+    
+    def __init__(self, hf_token: str, timeout: int = 30):
+        """
+        Initialize the Hugging Face MCP client.
+        
+        Args:
+            hf_token: Hugging Face API token
+            timeout: Timeout in seconds for HTTP requests
+        """
+        self.hf_token = hf_token
+        self.url = "https://huggingface.co/mcp"
+        self.headers = {"Authorization": f"Bearer {hf_token}"}
+        self.timeout = timedelta(seconds=timeout)
+        self.request_id_counter = 0
+        
+    def _get_next_request_id(self) -> int:
+        """Get the next request ID."""
+        self.request_id_counter += 1
+        return self.request_id_counter
+    
+    async def _send_request_and_get_response(
+        self, 
+        method: str, 
+        params: Optional[Dict[str, Any]] = None
+    ) -> Any:
+        """
+        Send a JSON-RPC request and wait for the response.
+        
+        Args:
+            method: The JSON-RPC method name
+            params: Optional parameters for the method
+            
+        Returns:
+            The response result or raises an exception
+        """
+        request_id = self._get_next_request_id()
+        
+        # Create JSON-RPC request
+        jsonrpc_request = JSONRPCRequest(
+            jsonrpc="2.0",
+            id=request_id,
+            method=method,
+            params=params
         )
-    return result
+        
+        message = JSONRPCMessage(jsonrpc_request)
+        session_message = SessionMessage(message)
+        
+        async with streamablehttp_client(
+            url=self.url,
+            headers=self.headers,
+            timeout=self.timeout,
+            terminate_on_close=True
+        ) as (read_stream, write_stream, get_session_id):
+            
+            try:
+                # Send initialization request first
+                init_request = JSONRPCRequest(
+                    jsonrpc="2.0",
+                    id=self._get_next_request_id(),
+                    method="initialize",
+                    params={
+                        "protocolVersion": "2024-11-05",
+                        "capabilities": {
+                            "tools": {}
+                        },
+                        "clientInfo": {
+                            "name": "hf-mcp-client",
+                            "version": "1.0.0"
+                        }
+                    }
+                )
+                
+                init_message = JSONRPCMessage(init_request)
+                init_session_message = SessionMessage(init_message)
+                
+                await write_stream.send(init_session_message)
+                
+                # Wait for initialization response
+                init_response_received = False
+                timeout_counter = 0
+                max_iterations = 100  # Prevent infinite loops
+                
+                while not init_response_received and timeout_counter < max_iterations:
+                    try:
+                        response = await read_stream.receive()
+                        timeout_counter += 1
+                        
+                        if isinstance(response, Exception):
+                            raise response
+                        
+                        if isinstance(response, SessionMessage):
+                            msg = response.message.root
+                            if isinstance(msg, JSONRPCResponse) and msg.id == init_request.id:
+                                logger.info("MCP client initialized successfully")
+                                init_response_received = True
+                            elif isinstance(msg, JSONRPCError) and msg.id == init_request.id:
+                                raise Exception(f"Initialization failed: {msg.error}")
+                    except Exception as e:
+                        if "ClosedResourceError" in str(type(e)):
+                            logger.error("Stream closed during initialization")
+                            raise Exception("Connection closed during initialization")
+                        raise
+                
+                if not init_response_received:
+                    raise Exception("Initialization timeout")
+                
+                # Send initialized notification
+                initialized_notification = JSONRPCNotification(
+                    jsonrpc="2.0",
+                    method="notifications/initialized"
+                )
+                
+                init_notif_message = JSONRPCMessage(initialized_notification)
+                init_notif_session_message = SessionMessage(init_notif_message)
+                
+                await write_stream.send(init_notif_session_message)
+                
+                # Small delay to let the notification process
+                await asyncio.sleep(0.1)
+                
+                # Now send our actual request
+                await write_stream.send(session_message)
+                
+                # Wait for the response to our request
+                response_received = False
+                timeout_counter = 0
+                
+                while not response_received and timeout_counter < max_iterations:
+                    try:
+                        response = await read_stream.receive()
+                        timeout_counter += 1
+                        
+                        if isinstance(response, Exception):
+                            raise response
+                        
+                        if isinstance(response, SessionMessage):
+                            msg = response.message.root
+                            if isinstance(msg, JSONRPCResponse) and msg.id == request_id:
+                                return msg.result
+                            elif isinstance(msg, JSONRPCError) and msg.id == request_id:
+                                raise Exception(f"Request failed: {msg.error}")
+                    except Exception as e:
+                        if "ClosedResourceError" in str(type(e)):
+                            logger.error("Stream closed during request processing")
+                            raise Exception("Connection closed during request processing")
+                        raise
+                
+                if not response_received:
+                    raise Exception("Request timeout")
+                    
+            except Exception as e:
+                logger.error(f"Error during MCP communication: {e}")
+                raise
+            finally:
+                # Ensure streams are properly closed
+                try:
+                    await write_stream.aclose()
+                except:
+                    pass
+    
+    async def get_all_tools(self) -> List[Dict[str, Any]]:
+        """
+        Get all available tools from the Hugging Face MCP endpoint.
+        
+        Returns:
+            List of tool definitions
+        """
+        try:
+            logger.info("Fetching all available tools from Hugging Face MCP")
+            result = await self._send_request_and_get_response("tools/list")
+            
+            if isinstance(result, dict) and "tools" in result:
+                tools = result["tools"]
+                logger.info(f"Found {len(tools)} available tools")
+                return tools
+            else:
+                logger.warning(f"Unexpected response format: {result}")
+                return []
+                
+        except Exception as e:
+            logger.error(f"Failed to get tools: {e}")
+            raise
+    
+    async def call_tool(self, tool_name: str, args: Dict[str, Any]) -> Any:
+        """
+        Call a specific tool with the given arguments.
+        
+        Args:
+            tool_name: Name of the tool to call
+            args: Arguments to pass to the tool
+            
+        Returns:
+            The tool's response
+        """
+        try:
+            logger.info(f"Calling tool '{tool_name}' with args: {args}")
+            
+            params = {
+                "name": tool_name,
+                "arguments": args
+            }
+            
+            result = await self._send_request_and_get_response("tools/call", params)
+            logger.info(f"Tool '{tool_name}' executed successfully")
+            return result
+            
+        except Exception as e:
+            logger.error(f"Failed to call tool '{tool_name}': {e}")
+            raise
+
+
+# Convenience functions for easier usage
+async def get_hf_tools(hf_token: str) -> List[Dict[str, Any]]:
+    """
+    Get all available tools from Hugging Face MCP.
+    
+    Args:
+        hf_token: Hugging Face API token
+        
+    Returns:
+        List of tool definitions
+    """
+    client = HuggingFaceMCPClient(hf_token)
+    return await client.get_all_tools()
+
+
+async def call_hf_tool(hf_token: str, tool_name: str, args: Dict[str, Any]) -> Any:
+    """
+    Call a specific Hugging Face MCP tool.
+    
+    Args:
+        hf_token: Hugging Face API token
+        tool_name: Name of the tool to call
+        args: Arguments to pass to the tool
+        
+    Returns:
+        The tool's response
+    """
+    client = HuggingFaceMCPClient(hf_token)
+    return await client.call_tool(tool_name, args)
+
+
+# Alternative simpler implementation for debugging
+class SimpleHFMCPClient:
+    """Simplified version for debugging connection issues."""
+    
+    def __init__(self, hf_token: str):
+        self.hf_token = hf_token
+        self.url = "https://huggingface.co/mcp"
+        self.headers = {"Authorization": f"Bearer {hf_token}"}
+        
+    async def test_connection(self):
+        """Test basic connection to HF MCP endpoint."""
+        try:
+            async with streamablehttp_client(
+                url=self.url,
+                headers=self.headers,
+                timeout=timedelta(seconds=10),
+                terminate_on_close=True
+            ) as (read_stream, write_stream, get_session_id):
+                logger.info("Connection established successfully")
+                return True
+        except Exception as e:
+            logger.error(f"Connection test failed: {e}")
+            return False
+