huggingface-hub 0.36.0__py3-none-any.whl → 1.0.0__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]
 
@@ -81,9 +81,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`.")
@@ -129,27 +129,27 @@ 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
+                        - 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
+                        - 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
+                        - 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
@@ -247,16 +247,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.
@@ -278,8 +278,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
@@ -326,7 +326,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(
                         {
@@ -344,6 +344,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,24 +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]]
+    allowed_tools: NotRequired[list[str]]
 
 
 class HTTPServerConfig(TypedDict):
     type: Literal["http"]
     url: str
-    headers: Dict[str, str]
-    allowed_tools: NotRequired[List[str]]
+    headers: dict[str, str]
+    allowed_tools: NotRequired[list[str]]
 
 
 class SSEServerConfig(TypedDict):
     type: Literal["sse"]
     url: str
-    headers: Dict[str, str]
-    allowed_tools: NotRequired[List[str]]
+    headers: dict[str, str]
+    allowed_tools: NotRequired[list[str]]
 
 
 ServerConfig = Union[StdioServerConfig, HTTPServerConfig, SSEServerConfig]
@@ -41,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,7 +6,7 @@ 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
@@ -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,10 +84,10 @@ 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")

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,
@@ -6,7 +6,7 @@ from huggingface_hub.inference._providers.featherless_ai import (
 )
 from huggingface_hub.utils import logging
 
-from ._common import TaskProviderHelper, _fetch_inference_provider_mapping
+from ._common import AutoRouterConversationalTask, TaskProviderHelper, _fetch_inference_provider_mapping
 from .black_forest_labs import BlackForestLabsTextToImageTask
 from .cerebras import CerebrasConversationalTask
 from .clarifai import ClarifaiConversationalTask
@@ -73,7 +73,9 @@ PROVIDER_T = Literal[
 
 PROVIDER_OR_POLICY_T = Union[PROVIDER_T, Literal["auto"]]
 
-PROVIDERS: Dict[PROVIDER_T, Dict[str, TaskProviderHelper]] = {
+CONVERSATIONAL_AUTO_ROUTER = AutoRouterConversationalTask()
+
+PROVIDERS: dict[PROVIDER_T, dict[str, TaskProviderHelper]] = {
     "black-forest-labs": {
         "text-to-image": BlackForestLabsTextToImageTask(),
     },
@@ -206,13 +208,19 @@ def get_provider_helper(
 
     if provider is None:
         logger.info(
-            "Defaulting to 'auto' which will select the first provider available for the model, sorted by the user's order in https://hf.co/settings/inference-providers."
+            "No provider specified for task `conversational`. Defaulting to server-side auto routing."
+            if task == "conversational"
+            else "Defaulting to 'auto' which will select the first provider available for the model, sorted by the user's order in https://hf.co/settings/inference-providers."
         )
         provider = "auto"
 
     if provider == "auto":
         if model is None:
             raise ValueError("Specifying a model is required when provider is 'auto'")
+        if task == "conversational":
+            # Special case: we have a dedicated auto-router for conversational models. No need to fetch provider mapping.
+            return CONVERSATIONAL_AUTO_ROUTER
+
         provider_mapping = _fetch_inference_provider_mapping(model)
         provider = next(iter(provider_mapping)).provider
 

huggingface_hub/inference/_providers/_common.py

@@ -1,5 +1,5 @@
 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
@@ -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:
@@ -41,14 +41,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
@@ -75,11 +75,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.
@@ -126,7 +126,7 @@ class TaskProviderHelper:
 
     def get_response(
         self,
-        response: Union[bytes, Dict],
+        response: Union[bytes, dict],
         request_params: Optional[RequestParameters] = None,
     ) -> Any:
         """
@@ -185,8 +185,8 @@ class TaskProviderHelper:
         return provider_mapping
 
     def _normalize_headers(
-        self, headers: Dict[str, Any], payload: Optional[Dict[str, Any]], data: Optional[MimeBytes]
-    ) -> Dict[str, Any]:
+        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.
@@ -199,7 +199,7 @@ class TaskProviderHelper:
                 normalized_headers["content-type"] = "application/json"
         return normalized_headers
 
-    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict[str, Any]:
+    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.
@@ -234,8 +234,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.
@@ -246,9 +246,9 @@ class TaskProviderHelper:
     def _prepare_payload_as_bytes(
         self,
         inputs: Any,
-        parameters: Dict,
+        parameters: dict,
         provider_mapping_info: InferenceProviderMapping,
-        extra_payload: Optional[Dict],
+        extra_payload: Optional[dict],
     ) -> Optional[MimeBytes]:
         """Return the body to use for the request, as bytes.
 
@@ -272,13 +272,51 @@ 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})
 
 
+class AutoRouterConversationalTask(BaseConversationalTask):
+    """
+    Auto-router for conversational tasks.
+
+    We let the Hugging Face router select the best provider for the model, based on availability and user preferences.
+    This is a special case since the selection is done server-side (avoid 1 API call to fetch provider mapping).
+    """
+
+    def __init__(self):
+        super().__init__(provider="auto", base_url="https://router.huggingface.co")
+
+    def _prepare_base_url(self, api_key: str) -> str:
+        """Return the base URL to use for the request.
+
+        Usually not overwritten in subclasses."""
+        # Route to the proxy if the api_key is a HF TOKEN
+        if not api_key.startswith("hf_"):
+            raise ValueError("Cannot select auto-router when using non-Hugging Face API key.")
+        else:
+            return self.base_url  # No `/auto` suffix in the URL
+
+    def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
+        """
+        In auto-router, we don't need to fetch provider mapping info.
+        We just return a dummy mapping info with provider_id set to the HF model ID.
+        """
+        if model is None:
+            raise ValueError("Please provide an HF model ID.")
+
+        return InferenceProviderMapping(
+            provider="auto",
+            hf_model_id=model,
+            providerId=model,
+            status="live",
+            task="conversational",
+        )
+
+
 class BaseTextGenerationTask(TaskProviderHelper):
     """
     Base class for text-generation (completion) tasks.
@@ -292,13 +330,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.
     """
@@ -311,7 +349,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[str, Any]:
+    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":

huggingface_hub/inference/_providers/fal_ai.py

@@ -1,7 +1,7 @@
 import base64
 import time
 from abc import ABC
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 from urllib.parse import urlparse
 
 from huggingface_hub import constants
@@ -22,7 +22,7 @@ class FalAITask(TaskProviderHelper, ABC):
     def __init__(self, task: str):
         super().__init__(provider="fal-ai", base_url="https://fal.run", task=task)
 
-    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict[str, Any]:
+    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["authorization"] = f"Key {api_key}"
@@ -36,7 +36,7 @@ class FalAIQueueTask(TaskProviderHelper, ABC):
     def __init__(self, task: str):
         super().__init__(provider="fal-ai", base_url="https://queue.fal.run", task=task)
 
-    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict[str, Any]:
+    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["authorization"] = f"Key {api_key}"
@@ -50,7 +50,7 @@ class FalAIQueueTask(TaskProviderHelper, ABC):
 
     def get_response(
         self,
-        response: Union[bytes, Dict],
+        response: Union[bytes, dict],
         request_params: Optional[RequestParameters] = None,
     ) -> Any:
         response_dict = _as_dict(response)
@@ -91,8 +91,8 @@ class FalAIAutomaticSpeechRecognitionTask(FalAITask):
         super().__init__("automatic-speech-recognition")
 
     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]:
         if isinstance(inputs, str) and inputs.startswith(("http://", "https://")):
             # If input is a URL, pass it directly
             audio_url = inputs
@@ -108,7 +108,7 @@ class FalAIAutomaticSpeechRecognitionTask(FalAITask):
 
         return {"audio_url": audio_url, **filter_none(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:
         text = _as_dict(response)["text"]
         if not isinstance(text, str):
             raise ValueError(f"Unexpected output format from FalAI API. Expected string, got {type(text)}.")
@@ -120,9 +120,9 @@ class FalAITextToImageTask(FalAITask):
         super().__init__("text-to-image")
 
     def _prepare_payload_as_dict(
-        self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
-    ) -> Optional[Dict]:
-        payload: Dict[str, Any] = {
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        payload: dict[str, Any] = {
             "prompt": inputs,
             **filter_none(parameters),
         }
@@ -145,7 +145,7 @@ class FalAITextToImageTask(FalAITask):
 
         return payload
 
-    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:
         url = _as_dict(response)["images"][0]["url"]
         return get_session().get(url).content
 
@@ -155,11 +155,11 @@ class FalAITextToSpeechTask(FalAITask):
         super().__init__("text-to-speech")
 
     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 {"text": inputs, **filter_none(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:
         url = _as_dict(response)["audio"]["url"]
         return get_session().get(url).content
 
@@ -169,13 +169,13 @@ class FalAITextToVideoTask(FalAIQueueTask):
         super().__init__("text-to-video")
 
     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 {"prompt": inputs, **filter_none(parameters)}
 
     def get_response(
         self,
-        response: Union[bytes, Dict],
+        response: Union[bytes, dict],
         request_params: Optional[RequestParameters] = None,
     ) -> Any:
         output = super().get_response(response, request_params)
@@ -188,12 +188,12 @@ class FalAIImageToImageTask(FalAIQueueTask):
         super().__init__("image-to-image")
 
     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]:
         image_url = _as_url(inputs, default_mime_type="image/jpeg")
         if "target_size" in parameters:
             parameters["image_size"] = parameters.pop("target_size")
-        payload: Dict[str, Any] = {
+        payload: dict[str, Any] = {
             "image_url": image_url,
             **filter_none(parameters),
         }
@@ -209,7 +209,7 @@ class FalAIImageToImageTask(FalAIQueueTask):
 
     def get_response(
         self,
-        response: Union[bytes, Dict],
+        response: Union[bytes, dict],
         request_params: Optional[RequestParameters] = None,
     ) -> Any:
         output = super().get_response(response, request_params)
@@ -222,10 +222,10 @@ class FalAIImageToVideoTask(FalAIQueueTask):
         super().__init__("image-to-video")
 
     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]:
         image_url = _as_url(inputs, default_mime_type="image/jpeg")
-        payload: Dict[str, Any] = {
+        payload: dict[str, Any] = {
             "image_url": image_url,
             **filter_none(parameters),
         }
@@ -240,7 +240,7 @@ class FalAIImageToVideoTask(FalAIQueueTask):
 
     def get_response(
         self,
-        response: Union[bytes, Dict],
+        response: Union[bytes, dict],
         request_params: Optional[RequestParameters] = None,
     ) -> Any:
         output = super().get_response(response, request_params)