Revamp stuff

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 get_hf_tools, call_hf_tool, SimpleHFMCPClient
+from utils.huggingface_mcp_llamaindex import get_hf_tools, call_hf_tool, diagnose_connection
 from prompts.devstral_coding_prompt import devstral_code_gen_sys_prompt, devstral_code_gen_user_prompt
 from dotenv import load_dotenv
 import os
@@ -49,7 +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")
+MCP_TOKEN = os.getenv("MCP_TOKEN")
 def get_file_hash(file_path):
     """Generate a hash of the file for caching purposes"""
     try:
@@ -245,17 +245,19 @@ async def generate_plan(history, file_cache):
             if ai_msg:
                 conversation_history += f"Assistant: {ai_msg}\n"
     
-    simple_client = SimpleHFMCPClient(HF_TOKEN)
-    connection_ok = await simple_client.test_connection()
+    print("Running connection diagnostics...")
+    diagnostics = await diagnose_connection(MCP_TOKEN)
+    print(f"Diagnostics: {json.dumps(diagnostics, indent=2)}")
     
-    if not connection_ok:
-        print("Basic connection test failed")
+    if not diagnostics["connection_test"]:
+        print("Basic connection failed - check token and network")
         return
         
-    print("Basic connection test passed")
-
+    if not diagnostics["tools_test"]:
+        print("Tools retrieval failed - check server status")
+        return
     # try:
-    hf_query_gen_tool_details = await get_hf_tools(hf_token=HF_TOKEN)
+    hf_query_gen_tool_details = await get_hf_tools(hf_token=MCP_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))
@@ -271,7 +273,7 @@ async def generate_plan(history, file_cache):
     
     # Call tool to get tool calls
     try:
-        tool_calls = await asyncio.gather(*[call_hf_tool(HF_TOKEN, step['tool'], step['args']) for step in parsed_plan])
+        tool_calls = await asyncio.gather(*[call_hf_tool(MCP_TOKEN, step['tool'], step['args']) for step in parsed_plan])
     except Exception as e:
         tool_calls = []
     print(tool_calls)

utils/huggingface_mcp_llamaindex.py

@@ -21,18 +21,22 @@ logger = logging.getLogger(__name__)
 class HuggingFaceMCPClient:
     """Client for interacting with Hugging Face MCP endpoint."""
     
-    def __init__(self, hf_token: str, timeout: int = 30):
+    def __init__(self, hf_token: str, timeout: int = 60):
         """
         Initialize the Hugging Face MCP client.
         
         Args:
             hf_token: Hugging Face API token
-            timeout: Timeout in seconds for HTTP requests
+            timeout: Timeout in seconds for HTTP requests (increased for Spaces)
         """
         self.hf_token = hf_token
         self.url = "https://huggingface.co/mcp"
-        self.headers = {"Authorization": f"Bearer {hf_token}"}
+        self.headers = {
+            "Authorization": f"Bearer {hf_token}",
+            "User-Agent": "hf-mcp-client/1.0.0"  # Add user agent
+        }
         self.timeout = timedelta(seconds=timeout)
+        self.sse_read_timeout = timedelta(seconds=timeout * 2)  # Longer SSE timeout
         self.request_id_counter = 0
         
     def _get_next_request_id(self) -> int:
@@ -43,18 +47,35 @@ class HuggingFaceMCPClient:
     async def _send_request_and_get_response(
         self, 
         method: str, 
-        params: Optional[Dict[str, Any]] = None
+        params: Optional[Dict[str, Any]] = None,
+        max_retries: int = 3
     ) -> Any:
         """
-        Send a JSON-RPC request and wait for the response.
+        Send a JSON-RPC request and wait for the response with retry logic.
         
         Args:
             method: The JSON-RPC method name
             params: Optional parameters for the method
+            max_retries: Maximum number of retry attempts
             
         Returns:
             The response result or raises an exception
         """
+        last_exception = None
+        
+        for attempt in range(max_retries):
+            try:
+                return await self._attempt_request(method, params)
+            except Exception as e:
+                last_exception = e
+                logger.warning(f"Attempt {attempt + 1} failed: {e}")
+                if attempt < max_retries - 1:
+                    await asyncio.sleep(2 ** attempt)  # Exponential backoff
+                
+        raise last_exception
+    
+    async def _attempt_request(self, method: str, params: Optional[Dict[str, Any]]) -> Any:
+        """Single attempt to send request."""
         request_id = self._get_next_request_id()
         
         # Create JSON-RPC request
@@ -72,11 +93,12 @@ class HuggingFaceMCPClient:
             url=self.url,
             headers=self.headers,
             timeout=self.timeout,
+            sse_read_timeout=self.sse_read_timeout,
             terminate_on_close=True
         ) as (read_stream, write_stream, get_session_id):
             
             try:
-                # Send initialization request first
+                # Send initialization request first with proper client info
                 init_request = JSONRPCRequest(
                     jsonrpc="2.0",
                     id=self._get_next_request_id(),
@@ -84,7 +106,9 @@ class HuggingFaceMCPClient:
                     params={
                         "protocolVersion": "2024-11-05",
                         "capabilities": {
-                            "tools": {}
+                            "tools": {},
+                            "resources": {},
+                            "prompts": {}
                         },
                         "clientInfo": {
                             "name": "hf-mcp-client",
@@ -96,19 +120,25 @@ class HuggingFaceMCPClient:
                 init_message = JSONRPCMessage(init_request)
                 init_session_message = SessionMessage(init_message)
                 
+                logger.info("Sending initialization request...")
                 await write_stream.send(init_session_message)
                 
-                # Wait for initialization response
+                # Wait for initialization response with better error handling
                 init_response_received = False
                 timeout_counter = 0
-                max_iterations = 100  # Prevent infinite loops
+                max_iterations = 150  # Increased for Spaces environment
                 
                 while not init_response_received and timeout_counter < max_iterations:
                     try:
-                        response = await read_stream.receive()
+                        # Use asyncio.wait_for to add timeout per receive operation
+                        response = await asyncio.wait_for(
+                            read_stream.receive(), 
+                            timeout=30.0  # 30 second timeout per receive
+                        )
                         timeout_counter += 1
                         
                         if isinstance(response, Exception):
+                            logger.error(f"Received exception during init: {response}")
                             raise response
                         
                         if isinstance(response, SessionMessage):
@@ -116,16 +146,30 @@ class HuggingFaceMCPClient:
                             if isinstance(msg, JSONRPCResponse) and msg.id == init_request.id:
                                 logger.info("MCP client initialized successfully")
                                 init_response_received = True
+                                # Log the session ID if available
+                                session_id = get_session_id()
+                                if session_id:
+                                    logger.info(f"Session ID: {session_id}")
                             elif isinstance(msg, JSONRPCError) and msg.id == init_request.id:
-                                raise Exception(f"Initialization failed: {msg.error}")
+                                error_msg = f"Initialization failed: {msg.error}"
+                                logger.error(error_msg)
+                                raise Exception(error_msg)
+                            else:
+                                logger.debug(f"Received other message during init: {type(msg)}")
+                                
+                    except asyncio.TimeoutError:
+                        logger.warning(f"Timeout waiting for response (attempt {timeout_counter})")
+                        if timeout_counter > 10:  # After 10 timeouts, give up
+                            raise Exception("Initialization timeout: no response from server")
                     except Exception as e:
-                        if "ClosedResourceError" in str(type(e)):
+                        if "ClosedResourceError" in str(type(e)) or "StreamClosed" in str(e):
                             logger.error("Stream closed during initialization")
                             raise Exception("Connection closed during initialization")
+                        logger.error(f"Error during initialization: {e}")
                         raise
                 
                 if not init_response_received:
-                    raise Exception("Initialization timeout")
+                    raise Exception("Initialization timeout: maximum iterations reached")
                 
                 # Send initialized notification
                 initialized_notification = JSONRPCNotification(
@@ -136,12 +180,14 @@ class HuggingFaceMCPClient:
                 init_notif_message = JSONRPCMessage(initialized_notification)
                 init_notif_session_message = SessionMessage(init_notif_message)
                 
+                logger.info("Sending initialized notification...")
                 await write_stream.send(init_notif_session_message)
                 
-                # Small delay to let the notification process
-                await asyncio.sleep(0.1)
+                # Longer delay for Spaces environment
+                await asyncio.sleep(1.0)
                 
                 # Now send our actual request
+                logger.info(f"Sending actual request: {method}")
                 await write_stream.send(session_message)
                 
                 # Wait for the response to our request
@@ -150,26 +196,41 @@ class HuggingFaceMCPClient:
                 
                 while not response_received and timeout_counter < max_iterations:
                     try:
-                        response = await read_stream.receive()
+                        response = await asyncio.wait_for(
+                            read_stream.receive(), 
+                            timeout=30.0
+                        )
                         timeout_counter += 1
                         
                         if isinstance(response, Exception):
+                            logger.error(f"Received exception during request: {response}")
                             raise response
                         
                         if isinstance(response, SessionMessage):
                             msg = response.message.root
                             if isinstance(msg, JSONRPCResponse) and msg.id == request_id:
+                                logger.info(f"Request '{method}' completed successfully")
                                 return msg.result
                             elif isinstance(msg, JSONRPCError) and msg.id == request_id:
-                                raise Exception(f"Request failed: {msg.error}")
+                                error_msg = f"Request failed: {msg.error}"
+                                logger.error(error_msg)
+                                raise Exception(error_msg)
+                            else:
+                                logger.debug(f"Received other message during request: {type(msg)}")
+                                
+                    except asyncio.TimeoutError:
+                        logger.warning(f"Timeout waiting for request response (attempt {timeout_counter})")
+                        if timeout_counter > 10:
+                            raise Exception("Request timeout: no response from server")
                     except Exception as e:
-                        if "ClosedResourceError" in str(type(e)):
+                        if "ClosedResourceError" in str(type(e)) or "StreamClosed" in str(e):
                             logger.error("Stream closed during request processing")
                             raise Exception("Connection closed during request processing")
+                        logger.error(f"Error during request processing: {e}")
                         raise
                 
                 if not response_received:
-                    raise Exception("Request timeout")
+                    raise Exception("Request timeout: maximum iterations reached")
                     
             except Exception as e:
                 logger.error(f"Error during MCP communication: {e}")
@@ -178,8 +239,8 @@ class HuggingFaceMCPClient:
                 # Ensure streams are properly closed
                 try:
                     await write_stream.aclose()
-                except:
-                    pass
+                except Exception as close_error:
+                    logger.debug(f"Error closing write stream: {close_error}")
     
     async def get_all_tools(self) -> List[Dict[str, Any]]:
         """
@@ -243,7 +304,7 @@ async def get_hf_tools(hf_token: str) -> List[Dict[str, Any]]:
     Returns:
         List of tool definitions
     """
-    client = HuggingFaceMCPClient(hf_token)
+    client = HuggingFaceMCPClient(hf_token, timeout=90)  # Longer timeout for Spaces
     return await client.get_all_tools()
 
 
@@ -259,31 +320,65 @@ async def call_hf_tool(hf_token: str, tool_name: str, args: Dict[str, Any]) -> A
     Returns:
         The tool's response
     """
-    client = HuggingFaceMCPClient(hf_token)
+    client = HuggingFaceMCPClient(hf_token, timeout=90)
     return await client.call_tool(tool_name, args)
 
 
-# Alternative simpler implementation for debugging
-class SimpleHFMCPClient:
-    """Simplified version for debugging connection issues."""
+# Diagnostic function for Spaces environment
+async def diagnose_connection(hf_token: str) -> Dict[str, Any]:
+    """
+    Diagnose connection issues with detailed logging.
     
-    def __init__(self, hf_token: str):
-        self.hf_token = hf_token
-        self.url = "https://huggingface.co/mcp"
-        self.headers = {"Authorization": f"Bearer {hf_token}"}
+    Args:
+        hf_token: Hugging Face API 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
+    Returns:
+        Diagnostic information
+    """
+    diagnostics = {
+        "environment": "huggingface_spaces" if os.getenv("SPACE_ID") else "local",
+        "space_id": os.getenv("SPACE_ID"),
+        "token_length": len(hf_token) if hf_token else 0,
+        "has_token": bool(hf_token),
+        "connection_test": False,
+        "initialization_test": False,
+        "tools_test": False,
+        "error": None
+    }
+    
+    try:
+        # Test basic connection
+        from mcp.client.streamable_http import streamablehttp_client
+        
+        headers = {
+            "Authorization": f"Bearer {hf_token}",
+            "User-Agent": "hf-mcp-diagnostic/1.0.0"
+        }
+        
+        async with streamablehttp_client(
+            url="https://huggingface.co/mcp",
+            headers=headers,
+            timeout=timedelta(seconds=30),
+            terminate_on_close=True
+        ) as (read_stream, write_stream, get_session_id):
+            diagnostics["connection_test"] = True
+            logger.info("Basic connection test passed")
+            
+            # Test initialization
+            client = HuggingFaceMCPClient(hf_token, timeout=60)
+            try:
+                tools = await client.get_all_tools()
+                diagnostics["initialization_test"] = True
+                diagnostics["tools_test"] = True
+                diagnostics["tool_count"] = len(tools)
+                logger.info(f"Full diagnostic passed - found {len(tools)} tools")
+            except Exception as init_error:
+                diagnostics["error"] = str(init_error)
+                logger.error(f"Initialization failed: {init_error}")
+                
+    except Exception as conn_error:
+        diagnostics["error"] = str(conn_error)
+        logger.error(f"Connection test failed: {conn_error}")
+    
+    return diagnostics