huggingface-hub 0.34.4__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

huggingface_hub/inference/_mcp/mcp_client.py

@@ -3,9 +3,9 @@ import logging
 from contextlib import AsyncExitStack
 from datetime import timedelta
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, AsyncIterable, Dict, List, Literal, Optional, Union, overload
+from typing import TYPE_CHECKING, Any, AsyncIterable, Literal, Optional, TypedDict, Union, overload
 
-from typing_extensions import NotRequired, TypeAlias, TypedDict, Unpack
+from typing_extensions import NotRequired, TypeAlias, Unpack
 
 from ...utils._runtime import get_hf_hub_version
 from .._generated._async_client import AsyncInferenceClient
@@ -32,14 +32,14 @@ ServerType: TypeAlias = Literal["stdio", "sse", "http"]
 
 class StdioServerParameters_T(TypedDict):
     command: str
-    args: NotRequired[List[str]]
-    env: NotRequired[Dict[str, str]]
+    args: NotRequired[list[str]]
+    env: NotRequired[dict[str, str]]
     cwd: NotRequired[Union[str, Path, None]]
 
 
 class SSEServerParameters_T(TypedDict):
     url: str
-    headers: NotRequired[Dict[str, Any]]
+    headers: NotRequired[dict[str, Any]]
     timeout: NotRequired[float]
     sse_read_timeout: NotRequired[float]
 
@@ -84,9 +84,9 @@ class MCPClient:
         api_key: Optional[str] = None,
     ):
         # Initialize MCP sessions as a dictionary of ClientSession objects
-        self.sessions: Dict[ToolName, "ClientSession"] = {}
+        self.sessions: dict[ToolName, "ClientSession"] = {}
         self.exit_stack = AsyncExitStack()
-        self.available_tools: List[ChatCompletionInputTool] = []
+        self.available_tools: list[ChatCompletionInputTool] = []
         # To be able to send the model in the payload if `base_url` is provided
         if model is None and base_url is None:
             raise ValueError("At least one of `model` or `base_url` should be set in `MCPClient`.")
@@ -132,28 +132,34 @@ class MCPClient:
                 - "stdio": Standard input/output server (local)
                 - "sse": Server-sent events (SSE) server
                 - "http": StreamableHTTP server
-            **params (`Dict[str, Any]`):
+            **params (`dict[str, Any]`):
                 Server parameters that can be either:
                     - For stdio servers:
                         - command (str): The command to run the MCP server
-                        - args (List[str], optional): Arguments for the command
-                        - env (Dict[str, str], optional): Environment variables for the command
+                        - args (list[str], optional): Arguments for the command
+                        - env (dict[str, str], optional): Environment variables for the command
                         - cwd (Union[str, Path, None], optional): Working directory for the command
+                        - allowed_tools (list[str], optional): List of tool names to allow from this server
                     - For SSE servers:
                         - url (str): The URL of the SSE server
-                        - headers (Dict[str, Any], optional): Headers for the SSE connection
+                        - headers (dict[str, Any], optional): Headers for the SSE connection
                         - timeout (float, optional): Connection timeout
                         - sse_read_timeout (float, optional): SSE read timeout
+                        - allowed_tools (list[str], optional): List of tool names to allow from this server
                     - For StreamableHTTP servers:
                         - url (str): The URL of the StreamableHTTP server
-                        - headers (Dict[str, Any], optional): Headers for the StreamableHTTP connection
+                        - headers (dict[str, Any], optional): Headers for the StreamableHTTP connection
                         - timeout (timedelta, optional): Connection timeout
                         - sse_read_timeout (timedelta, optional): SSE read timeout
                         - terminate_on_close (bool, optional): Whether to terminate on close
+                        - allowed_tools (list[str], optional): List of tool names to allow from this server
         """
         from mcp import ClientSession, StdioServerParameters
         from mcp import types as mcp_types
 
+        # Extract allowed_tools configuration if provided
+        allowed_tools = params.pop("allowed_tools", [])
+
         # Determine server type and create appropriate parameters
         if type == "stdio":
             # Handle stdio server
@@ -211,7 +217,15 @@ class MCPClient:
         response = await session.list_tools()
         logger.debug("Connected to server with tools:", [tool.name for tool in response.tools])
 
-        for tool in response.tools:
+        # Filter tools based on allowed_tools configuration
+        filtered_tools = [tool for tool in response.tools if tool.name in allowed_tools]
+
+        if allowed_tools:
+            logger.debug(
+                f"Tool filtering applied. Using {len(filtered_tools)} of {len(response.tools)} available tools: {[tool.name for tool in filtered_tools]}"
+            )
+
+        for tool in filtered_tools:
             if tool.name in self.sessions:
                 logger.warning(f"Tool '{tool.name}' already defined by another server. Skipping.")
                 continue
@@ -235,16 +249,16 @@ class MCPClient:
 
     async def process_single_turn_with_tools(
         self,
-        messages: List[Union[Dict, ChatCompletionInputMessage]],
-        exit_loop_tools: Optional[List[ChatCompletionInputTool]] = None,
+        messages: list[Union[dict, ChatCompletionInputMessage]],
+        exit_loop_tools: Optional[list[ChatCompletionInputTool]] = None,
         exit_if_first_chunk_no_tool: bool = False,
     ) -> AsyncIterable[Union[ChatCompletionStreamOutput, ChatCompletionInputMessage]]:
         """Process a query using `self.model` and available tools, yielding chunks and tool outputs.
 
         Args:
-            messages (`List[Dict]`):
+            messages (`list[dict]`):
                 List of message objects representing the conversation history
-            exit_loop_tools (`List[ChatCompletionInputTool]`, *optional*):
+            exit_loop_tools (`list[ChatCompletionInputTool]`, *optional*):
                 List of tools that should exit the generator when called
             exit_if_first_chunk_no_tool (`bool`, *optional*):
                 Exit if no tool is present in the first chunks. Default to False.
@@ -266,8 +280,8 @@ class MCPClient:
             stream=True,
         )
 
-        message: Dict[str, Any] = {"role": "unknown", "content": ""}
-        final_tool_calls: Dict[int, ChatCompletionStreamOutputDeltaToolCall] = {}
+        message: dict[str, Any] = {"role": "unknown", "content": ""}
+        final_tool_calls: dict[int, ChatCompletionStreamOutputDeltaToolCall] = {}
         num_of_chunks = 0
 
         # Read from stream
@@ -286,16 +300,19 @@ class MCPClient:
             # Process tool calls
             if delta.tool_calls:
                 for tool_call in delta.tool_calls:
-                    # Aggregate chunks into tool calls
-                    if tool_call.index not in final_tool_calls:
-                        if (
-                            tool_call.function.arguments is None or tool_call.function.arguments == "{}"
-                        ):  # Corner case (depends on provider)
-                            tool_call.function.arguments = ""
-                        final_tool_calls[tool_call.index] = tool_call
-
-                    elif tool_call.function.arguments:
-                        final_tool_calls[tool_call.index].function.arguments += tool_call.function.arguments
+                    idx = tool_call.index
+                    # first chunk for this tool call
+                    if idx not in final_tool_calls:
+                        final_tool_calls[idx] = tool_call
+                        if final_tool_calls[idx].function.arguments is None:
+                            final_tool_calls[idx].function.arguments = ""
+                        continue
+                    # safety before concatenating text to .function.arguments
+                    if final_tool_calls[idx].function.arguments is None:
+                        final_tool_calls[idx].function.arguments = ""
+
+                    if tool_call.function.arguments:
+                        final_tool_calls[idx].function.arguments += tool_call.function.arguments
 
             # Optionally exit early if no tools in first chunks
             if exit_if_first_chunk_no_tool and num_of_chunks <= 2 and len(final_tool_calls) == 0:
@@ -311,7 +328,7 @@ class MCPClient:
                 message["role"] = "assistant"
             # Convert final_tool_calls to the format expected by OpenAI
             if final_tool_calls:
-                tool_calls_list: List[Dict[str, Any]] = []
+                tool_calls_list: list[dict[str, Any]] = []
                 for tc in final_tool_calls.values():
                     tool_calls_list.append(
                         {
@@ -329,6 +346,17 @@ class MCPClient:
         # Process tool calls one by one
         for tool_call in final_tool_calls.values():
             function_name = tool_call.function.name
+            if function_name is None:
+                message = ChatCompletionInputMessage.parse_obj_as_instance(
+                    {
+                        "role": "tool",
+                        "tool_call_id": tool_call.id,
+                        "content": "Invalid tool call with no function name.",
+                    }
+                )
+                messages.append(message)
+                yield message
+                continue  # move to next tool call
             try:
                 function_args = json.loads(tool_call.function.arguments or "{}")
             except json.JSONDecodeError as err:

huggingface_hub/inference/_mcp/types.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Literal, TypedDict, Union
+from typing import Literal, TypedDict, Union
 
 from typing_extensions import NotRequired
 
@@ -13,21 +13,24 @@ class InputConfig(TypedDict, total=False):
 class StdioServerConfig(TypedDict):
     type: Literal["stdio"]
     command: str
-    args: List[str]
-    env: Dict[str, str]
+    args: list[str]
+    env: dict[str, str]
     cwd: str
+    allowed_tools: NotRequired[list[str]]
 
 
 class HTTPServerConfig(TypedDict):
     type: Literal["http"]
     url: str
-    headers: Dict[str, str]
+    headers: dict[str, str]
+    allowed_tools: NotRequired[list[str]]
 
 
 class SSEServerConfig(TypedDict):
     type: Literal["sse"]
     url: str
-    headers: Dict[str, str]
+    headers: dict[str, str]
+    allowed_tools: NotRequired[list[str]]
 
 
 ServerConfig = Union[StdioServerConfig, HTTPServerConfig, SSEServerConfig]
@@ -38,5 +41,5 @@ class AgentConfig(TypedDict):
     model: str
     provider: str
     apiKey: NotRequired[str]
-    inputs: List[InputConfig]
-    servers: List[ServerConfig]
+    inputs: list[InputConfig]
+    servers: list[ServerConfig]

huggingface_hub/inference/_mcp/utils.py

@@ -6,12 +6,12 @@ Formatting utilities taken from the JS SDK: https://github.com/huggingface/huggi
 
 import json
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Tuple
+from typing import TYPE_CHECKING, Optional
 
 from huggingface_hub import snapshot_download
 from huggingface_hub.errors import EntryNotFoundError
 
-from .constants import DEFAULT_AGENT, DEFAULT_REPO_ID, FILENAME_CONFIG, FILENAME_PROMPT
+from .constants import DEFAULT_AGENT, DEFAULT_REPO_ID, FILENAME_CONFIG, PROMPT_FILENAMES
 from .types import AgentConfig
 
 
@@ -36,7 +36,7 @@ def format_result(result: "mcp_types.CallToolResult") -> str:
     if len(content) == 0:
         return "[No content]"
 
-    formatted_parts: List[str] = []
+    formatted_parts: list[str] = []
 
     for item in content:
         if item.type == "text":
@@ -84,17 +84,21 @@ def _get_base64_size(base64_str: str) -> int:
     return (len(base64_str) * 3) // 4 - padding
 
 
-def _load_agent_config(agent_path: Optional[str]) -> Tuple[AgentConfig, Optional[str]]:
+def _load_agent_config(agent_path: Optional[str]) -> tuple[AgentConfig, Optional[str]]:
     """Load server config and prompt."""
 
-    def _read_dir(directory: Path) -> Tuple[AgentConfig, Optional[str]]:
+    def _read_dir(directory: Path) -> tuple[AgentConfig, Optional[str]]:
         cfg_file = directory / FILENAME_CONFIG
         if not cfg_file.exists():
             raise FileNotFoundError(f" Config file not found in {directory}! Please make sure it exists locally")
 
         config: AgentConfig = json.loads(cfg_file.read_text(encoding="utf-8"))
-        prompt_file = directory / FILENAME_PROMPT
-        prompt: Optional[str] = prompt_file.read_text(encoding="utf-8") if prompt_file.exists() else None
+        prompt: Optional[str] = None
+        for filename in PROMPT_FILENAMES:
+            prompt_file = directory / filename
+            if prompt_file.exists():
+                prompt = prompt_file.read_text(encoding="utf-8")
+                break
         return config, prompt
 
     if agent_path is None:

huggingface_hub/inference/_providers/__init__.py

@@ -1,4 +1,4 @@
-from typing import Dict, Literal, Optional, Union
+from typing import Literal, Optional, Union
 
 from huggingface_hub.inference._providers.featherless_ai import (
     FeatherlessConversationalTask,
@@ -65,7 +65,7 @@ PROVIDER_T = Literal[
 
 PROVIDER_OR_POLICY_T = Union[PROVIDER_T, Literal["auto"]]
 
-PROVIDERS: Dict[PROVIDER_T, Dict[str, TaskProviderHelper]] = {
+PROVIDERS: dict[PROVIDER_T, dict[str, TaskProviderHelper]] = {
     "black-forest-labs": {
         "text-to-image": BlackForestLabsTextToImageTask(),
     },

huggingface_hub/inference/_providers/_common.py

@@ -1,9 +1,9 @@
 from functools import lru_cache
-from typing import Any, Dict, List, Optional, Union, overload
+from typing import Any, Optional, Union, overload
 
 from huggingface_hub import constants
 from huggingface_hub.hf_api import InferenceProviderMapping
-from huggingface_hub.inference._common import RequestParameters
+from huggingface_hub.inference._common import MimeBytes, RequestParameters
 from huggingface_hub.inference._generated.types.chat_completion import ChatCompletionInputMessage
 from huggingface_hub.utils import build_hf_headers, get_token, logging
 
@@ -14,7 +14,7 @@ logger = logging.get_logger(__name__)
 # Dev purposes only.
 # If you want to try to run inference for a new model locally before it's registered on huggingface.co
 # for a given Inference Provider, you can add it to the following dictionary.
-HARDCODED_MODEL_INFERENCE_MAPPING: Dict[str, Dict[str, InferenceProviderMapping]] = {
+HARDCODED_MODEL_INFERENCE_MAPPING: dict[str, dict[str, InferenceProviderMapping]] = {
     # "HF model ID" => InferenceProviderMapping object initialized with "Model ID on Inference Provider's side"
     #
     # Example:
@@ -38,14 +38,14 @@ HARDCODED_MODEL_INFERENCE_MAPPING: Dict[str, Dict[str, InferenceProviderMapping]
 
 
 @overload
-def filter_none(obj: Dict[str, Any]) -> Dict[str, Any]: ...
+def filter_none(obj: dict[str, Any]) -> dict[str, Any]: ...
 @overload
-def filter_none(obj: List[Any]) -> List[Any]: ...
+def filter_none(obj: list[Any]) -> list[Any]: ...
 
 
-def filter_none(obj: Union[Dict[str, Any], List[Any]]) -> Union[Dict[str, Any], List[Any]]:
+def filter_none(obj: Union[dict[str, Any], list[Any]]) -> Union[dict[str, Any], list[Any]]:
     if isinstance(obj, dict):
-        cleaned: Dict[str, Any] = {}
+        cleaned: dict[str, Any] = {}
         for k, v in obj.items():
             if v is None:
                 continue
@@ -72,11 +72,11 @@ class TaskProviderHelper:
         self,
         *,
         inputs: Any,
-        parameters: Dict[str, Any],
-        headers: Dict,
+        parameters: dict[str, Any],
+        headers: dict,
         model: Optional[str],
         api_key: Optional[str],
-        extra_payload: Optional[Dict[str, Any]] = None,
+        extra_payload: Optional[dict[str, Any]] = None,
     ) -> RequestParameters:
         """
         Prepare the request to be sent to the provider.
@@ -108,13 +108,22 @@ class TaskProviderHelper:
             raise ValueError("Both payload and data cannot be set in the same request.")
         if payload is None and data is None:
             raise ValueError("Either payload or data must be set in the request.")
+
+        # normalize headers to lowercase and add content-type if not present
+        normalized_headers = self._normalize_headers(headers, payload, data)
+
         return RequestParameters(
-            url=url, task=self.task, model=provider_mapping_info.provider_id, json=payload, data=data, headers=headers
+            url=url,
+            task=self.task,
+            model=provider_mapping_info.provider_id,
+            json=payload,
+            data=data,
+            headers=normalized_headers,
         )
 
     def get_response(
         self,
-        response: Union[bytes, Dict],
+        response: Union[bytes, dict],
         request_params: Optional[RequestParameters] = None,
     ) -> Any:
         """
@@ -172,7 +181,22 @@ class TaskProviderHelper:
             )
         return provider_mapping
 
-    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
+    def _normalize_headers(
+        self, headers: dict[str, Any], payload: Optional[dict[str, Any]], data: Optional[MimeBytes]
+    ) -> dict[str, Any]:
+        """Normalize the headers to use for the request.
+
+        Override this method in subclasses for customized headers.
+        """
+        normalized_headers = {key.lower(): value for key, value in headers.items() if value is not None}
+        if normalized_headers.get("content-type") is None:
+            if data is not None and data.mime_type is not None:
+                normalized_headers["content-type"] = data.mime_type
+            elif payload is not None:
+                normalized_headers["content-type"] = "application/json"
+        return normalized_headers
+
+    def _prepare_headers(self, headers: dict, api_key: str) -> dict[str, Any]:
         """Return the headers to use for the request.
 
         Override this method in subclasses for customized headers.
@@ -207,8 +231,8 @@ class TaskProviderHelper:
         return ""
 
     def _prepare_payload_as_dict(
-        self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
-    ) -> Optional[Dict]:
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
         """Return the payload to use for the request, as a dict.
 
         Override this method in subclasses for customized payloads.
@@ -219,10 +243,10 @@ class TaskProviderHelper:
     def _prepare_payload_as_bytes(
         self,
         inputs: Any,
-        parameters: Dict,
+        parameters: dict,
         provider_mapping_info: InferenceProviderMapping,
-        extra_payload: Optional[Dict],
-    ) -> Optional[bytes]:
+        extra_payload: Optional[dict],
+    ) -> Optional[MimeBytes]:
         """Return the body to use for the request, as bytes.
 
         Override this method in subclasses for customized body data.
@@ -245,10 +269,10 @@ class BaseConversationalTask(TaskProviderHelper):
 
     def _prepare_payload_as_dict(
         self,
-        inputs: List[Union[Dict, ChatCompletionInputMessage]],
-        parameters: Dict,
+        inputs: list[Union[dict, ChatCompletionInputMessage]],
+        parameters: dict,
         provider_mapping_info: InferenceProviderMapping,
-    ) -> Optional[Dict]:
+    ) -> Optional[dict]:
         return filter_none({"messages": inputs, **parameters, "model": provider_mapping_info.provider_id})
 
 
@@ -265,13 +289,13 @@ class BaseTextGenerationTask(TaskProviderHelper):
         return "/v1/completions"
 
     def _prepare_payload_as_dict(
-        self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
-    ) -> Optional[Dict]:
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
         return filter_none({"prompt": inputs, **parameters, "model": provider_mapping_info.provider_id})
 
 
 @lru_cache(maxsize=None)
-def _fetch_inference_provider_mapping(model: str) -> List["InferenceProviderMapping"]:
+def _fetch_inference_provider_mapping(model: str) -> list["InferenceProviderMapping"]:
     """
     Fetch provider mappings for a model from the Hub.
     """
@@ -284,7 +308,7 @@ def _fetch_inference_provider_mapping(model: str) -> List["InferenceProviderMapp
     return provider_mapping
 
 
-def recursive_merge(dict1: Dict, dict2: Dict) -> Dict:
+def recursive_merge(dict1: dict, dict2: dict) -> dict:
     return {
         **dict1,
         **{

huggingface_hub/inference/_providers/black_forest_labs.py

@@ -1,5 +1,5 @@
 import time
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 from huggingface_hub.hf_api import InferenceProviderMapping
 from huggingface_hub.inference._common import RequestParameters, _as_dict
@@ -18,7 +18,7 @@ class BlackForestLabsTextToImageTask(TaskProviderHelper):
     def __init__(self):
         super().__init__(provider="black-forest-labs", base_url="https://api.us1.bfl.ai", task="text-to-image")
 
-    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
+    def _prepare_headers(self, headers: dict, api_key: str) -> dict[str, Any]:
         headers = super()._prepare_headers(headers, api_key)
         if not api_key.startswith("hf_"):
             _ = headers.pop("authorization")
@@ -29,8 +29,8 @@ class BlackForestLabsTextToImageTask(TaskProviderHelper):
         return f"/v1/{mapped_model}"
 
     def _prepare_payload_as_dict(
-        self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
-    ) -> Optional[Dict]:
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
         parameters = filter_none(parameters)
         if "num_inference_steps" in parameters:
             parameters["steps"] = parameters.pop("num_inference_steps")
@@ -39,7 +39,7 @@ class BlackForestLabsTextToImageTask(TaskProviderHelper):
 
         return {"prompt": inputs, **parameters}
 
-    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
+    def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
         """
         Polling mechanism for Black Forest Labs since the API is asynchronous.
         """
@@ -50,7 +50,7 @@ class BlackForestLabsTextToImageTask(TaskProviderHelper):
 
             response = session.get(url, headers={"Content-Type": "application/json"})  # type: ignore
             response.raise_for_status()  # type: ignore
-            response_json: Dict = response.json()  # type: ignore
+            response_json: dict = response.json()  # type: ignore
             status = response_json.get("status")
             logger.info(
                 f"Polling generation result from {url}. Current status: {status}. "

huggingface_hub/inference/_providers/cohere.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from huggingface_hub.hf_api import InferenceProviderMapping
 
@@ -17,8 +17,8 @@ class CohereConversationalTask(BaseConversationalTask):
         return "/compatibility/v1/chat/completions"
 
     def _prepare_payload_as_dict(
-        self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
-    ) -> Optional[Dict]:
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
         payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
         response_format = parameters.get("response_format")
         if isinstance(response_format, dict) and response_format.get("type") == "json_schema":