codekingpro/portable-devtools

1# Copyright 2023-present, the HuggingFace Inc. team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""Contains `WebhooksServer` and `webhook_endpoint` to create a webhook server easily."""15 16import atexit17import inspect18import os19from collections.abc import Callable20from functools import wraps21from typing import TYPE_CHECKING, Any, Optional22 23from .utils import experimental, is_fastapi_available, is_gradio_available24 25 26if TYPE_CHECKING:27    import gradio as gr28    from fastapi import Request29 30if is_fastapi_available():31    from fastapi import FastAPI, Request32    from fastapi.responses import JSONResponse33else:34    # Will fail at runtime if FastAPI is not available35    FastAPI = Request = JSONResponse = None  # type: ignore36 37 38_global_app: Optional["WebhooksServer"] = None39_is_local = os.environ.get("SPACE_ID") is None40 41 42@experimental43class WebhooksServer:44    """45    The [`WebhooksServer`] class lets you create an instance of a Gradio app that can receive Huggingface webhooks.46    These webhooks can be registered using the [`~WebhooksServer.add_webhook`] decorator. Webhook endpoints are added to47    the app as a POST endpoint to the FastAPI router. Once all the webhooks are registered, the `launch` method has to be48    called to start the app.49 50    It is recommended to accept [`WebhookPayload`] as the first argument of the webhook function. It is a Pydantic51    model that contains all the information about the webhook event. The data will be parsed automatically for you.52 53    Check out the [webhooks guide](../guides/webhooks_server) for a step-by-step tutorial on how to set up your54    WebhooksServer and deploy it on a Space.55 56    > [!WARNING]57    > `WebhooksServer` is experimental. Its API is subject to change in the future.58 59    > [!WARNING]60    > You must have `gradio` installed to use `WebhooksServer` (`pip install --upgrade gradio`).61 62    Args:63        ui (`gradio.Blocks`, optional):64            A Gradio UI instance to be used as the Space landing page. If `None`, a UI displaying instructions65            about the configured webhooks is created.66        webhook_secret (`str`, optional):67            A secret key to verify incoming webhook requests. You can set this value to any secret you want as long as68            you also configure it in your [webhooks settings panel](https://huggingface.co/settings/webhooks). You69            can also set this value as the `WEBHOOK_SECRET` environment variable. If no secret is provided, the70            webhook endpoints are opened without any security.71 72    Example:73 74        ```python75        import gradio as gr76        from huggingface_hub import WebhooksServer, WebhookPayload77 78        with gr.Blocks() as ui:79            ...80 81        app = WebhooksServer(ui=ui, webhook_secret="my_secret_key")82 83        @app.add_webhook("/say_hello")84        async def hello(payload: WebhookPayload):85            return {"message": "hello"}86 87        app.launch()88        ```89    """90 91    def __new__(cls, *args, **kwargs) -> "WebhooksServer":92        if not is_gradio_available():93            raise ImportError(94                "You must have `gradio` installed to use `WebhooksServer`. Please run `pip install --upgrade gradio`"95                " first."96            )97        if not is_fastapi_available():98            raise ImportError(99                "You must have `fastapi` installed to use `WebhooksServer`. Please run `pip install --upgrade fastapi`"100                " first."101            )102        return super().__new__(cls)103 104    def __init__(105        self,106        ui: Optional["gr.Blocks"] = None,107        webhook_secret: str | None = None,108    ) -> None:109        self._ui = ui110 111        self.webhook_secret = webhook_secret or os.getenv("WEBHOOK_SECRET")112        self.registered_webhooks: dict[str, Callable] = {}113        _warn_on_empty_secret(self.webhook_secret)114 115    def add_webhook(self, path: str | None = None) -> Callable:116        """117        Decorator to add a webhook to the [`WebhooksServer`] server.118 119        Args:120            path (`str`, optional):121                The URL path to register the webhook function. If not provided, the function name will be used as the122                path. In any case, all webhooks are registered under `/webhooks`.123 124        Raises:125            ValueError: If the provided path is already registered as a webhook.126 127        Example:128            ```python129            from huggingface_hub import WebhooksServer, WebhookPayload130 131            app = WebhooksServer()132 133            @app.add_webhook134            async def trigger_training(payload: WebhookPayload):135                if payload.repo.type == "dataset" and payload.event.action == "update":136                    # Trigger a training job if a dataset is updated137                    ...138 139            app.launch()140        ```141        """142        # Usage: directly as decorator. Example: `@app.add_webhook`143        if callable(path):144            # If path is a function, it means it was used as a decorator without arguments145            return self.add_webhook()(path)146 147        # Usage: provide a path. Example: `@app.add_webhook(...)`148        @wraps(FastAPI.post)149        def _inner_post(*args, **kwargs):150            func = args[0]151            abs_path = f"/webhooks/{(path or func.__name__).strip('/')}"152            if abs_path in self.registered_webhooks:153                raise ValueError(f"Webhook {abs_path} already exists.")154            self.registered_webhooks[abs_path] = func155 156        return _inner_post157 158    def launch(self, prevent_thread_lock: bool = False, **launch_kwargs: Any) -> None:159        """Launch the Gradio app and register webhooks to the underlying FastAPI server.160 161        Input parameters are forwarded to Gradio when launching the app.162        """163        ui = self._ui or self._get_default_ui()164 165        # Start Gradio App166        #   - as non-blocking so that webhooks can be added afterwards167        #   - as shared if launch locally (to debug webhooks)168        launch_kwargs.setdefault("share", _is_local)169        self.fastapi_app, _, _ = ui.launch(prevent_thread_lock=True, **launch_kwargs)170 171        # Register webhooks to FastAPI app172        for path, func in self.registered_webhooks.items():173            # Add secret check if required174            if self.webhook_secret is not None:175                func = _wrap_webhook_to_check_secret(func, webhook_secret=self.webhook_secret)176 177            # Add route to FastAPI app178            self.fastapi_app.post(path)(func)179 180        # Print instructions and block main thread181        space_host = os.environ.get("SPACE_HOST")182        url = "https://" + space_host if space_host is not None else (ui.share_url or ui.local_url)183        if url is None:184            raise ValueError("Cannot find the URL of the app. Please provide a valid `ui` or update `gradio` version.")185        url = url.strip("/")186        message = "\nWebhooks are correctly setup and ready to use:"187        message += "\n" + "\n".join(f"  - POST {url}{webhook}" for webhook in self.registered_webhooks)188        message += "\nGo to https://huggingface.co/settings/webhooks to setup your webhooks."189        print(message)190 191        if not prevent_thread_lock:192            ui.block_thread()193 194    def _get_default_ui(self) -> "gr.Blocks":195        """Default UI if not provided (lists webhooks and provides basic instructions)."""196        import gradio as gr197 198        with gr.Blocks() as ui:199            gr.Markdown("# This is an app to process 🤗 Webhooks")200            gr.Markdown(201                "Webhooks are a foundation for MLOps-related features. They allow you to listen for new changes on"202                " specific repos or to all repos belonging to particular set of users/organizations (not just your"203                " repos, but any repo). Check out this [guide](https://huggingface.co/docs/hub/webhooks) to get to"204                " know more about webhooks on the Huggingface Hub."205            )206            gr.Markdown(207                f"{len(self.registered_webhooks)} webhook(s) are registered:"208                + "\n\n"209                + "\n ".join(210                    f"- [{webhook_path}]({_get_webhook_doc_url(webhook.__name__, webhook_path)})"211                    for webhook_path, webhook in self.registered_webhooks.items()212                )213            )214            gr.Markdown(215                "Go to https://huggingface.co/settings/webhooks to setup your webhooks."216                + "\nYou app is running locally. Please look at the logs to check the full URL you need to set."217                if _is_local218                else (219                    "\nThis app is running on a Space. You can find the corresponding URL in the options menu"220                    " (top-right) > 'Embed the Space'. The URL looks like 'https://{username}-{repo_name}.hf.space'."221                )222            )223        return ui224 225 226@experimental227def webhook_endpoint(path: str | None = None) -> Callable:228    """Decorator to start a [`WebhooksServer`] and register the decorated function as a webhook endpoint.229 230    This is a helper to get started quickly. If you need more flexibility (custom landing page or webhook secret),231    you can use [`WebhooksServer`] directly. You can register multiple webhook endpoints (to the same server) by using232    this decorator multiple times.233 234    Check out the [webhooks guide](../guides/webhooks_server) for a step-by-step tutorial on how to set up your235    server and deploy it on a Space.236 237    > [!WARNING]238    > `webhook_endpoint` is experimental. Its API is subject to change in the future.239 240    > [!WARNING]241    > You must have `gradio` installed to use `webhook_endpoint` (`pip install --upgrade gradio`).242 243    Args:244        path (`str`, optional):245            The URL path to register the webhook function. If not provided, the function name will be used as the path.246            In any case, all webhooks are registered under `/webhooks`.247 248    Examples:249        The default usage is to register a function as a webhook endpoint. The function name will be used as the path.250        The server will be started automatically at exit (i.e. at the end of the script).251 252        ```python253        from huggingface_hub import webhook_endpoint, WebhookPayload254 255        @webhook_endpoint256        async def trigger_training(payload: WebhookPayload):257            if payload.repo.type == "dataset" and payload.event.action == "update":258                # Trigger a training job if a dataset is updated259                ...260 261        # Server is automatically started at the end of the script.262        ```263 264        Advanced usage: register a function as a webhook endpoint and start the server manually. This is useful if you265        are running it in a notebook.266 267        ```python268        from huggingface_hub import webhook_endpoint, WebhookPayload269 270        @webhook_endpoint271        async def trigger_training(payload: WebhookPayload):272            if payload.repo.type == "dataset" and payload.event.action == "update":273                # Trigger a training job if a dataset is updated274                ...275 276        # Start the server manually277        trigger_training.launch()278        ```279    """280    if callable(path):281        # If path is a function, it means it was used as a decorator without arguments282        return webhook_endpoint()(path)283 284    @wraps(WebhooksServer.add_webhook)285    def _inner(func: Callable) -> Callable:286        app = _get_global_app()287        app.add_webhook(path)(func)288        if len(app.registered_webhooks) == 1:289            # Register `app.launch` to run at exit (only once)290            atexit.register(app.launch)291 292        @wraps(app.launch)293        def _launch_now():294            # Run the app directly (without waiting atexit)295            atexit.unregister(app.launch)296            app.launch()297 298        func.launch = _launch_now  # type: ignore299        return func300 301    return _inner302 303 304def _get_global_app() -> WebhooksServer:305    global _global_app306    if _global_app is None:307        _global_app = WebhooksServer()308    return _global_app309 310 311def _warn_on_empty_secret(webhook_secret: str | None) -> None:312    if webhook_secret is None:313        print("Webhook secret is not defined. This means your webhook endpoints will be open to everyone.")314        print(315            "To add a secret, set `WEBHOOK_SECRET` as environment variable or pass it at initialization: "316            "\n\t`app = WebhooksServer(webhook_secret='my_secret', ...)`"317        )318        print(319            "For more details about webhook secrets, please refer to"320            " https://huggingface.co/docs/hub/webhooks#webhook-secret."321        )322    else:323        print("Webhook secret is correctly defined.")324 325 326def _get_webhook_doc_url(webhook_name: str, webhook_path: str) -> str:327    """Returns the anchor to a given webhook in the docs (experimental)"""328    return "/docs#/default/" + webhook_name + webhook_path.replace("/", "_") + "_post"329 330 331def _wrap_webhook_to_check_secret(func: Callable, webhook_secret: str) -> Callable:332    """Wraps a webhook function to check the webhook secret before calling the function.333 334    This is a hacky way to add the `request` parameter to the function signature. Since FastAPI based itself on route335    parameters to inject the values to the function, we need to hack the function signature to retrieve the `Request`336    object (and hence the headers). A far cleaner solution would be to use a middleware. However, since337    `fastapi==0.90.1`, a middleware cannot be added once the app has started. And since the FastAPI app is started by338    Gradio internals (and not by us), we cannot add a middleware.339 340    This method is called only when a secret has been defined by the user. If a request is sent without the341    "x-webhook-secret", the function will return a 401 error (unauthorized). If the header is sent but is incorrect,342    the function will return a 403 error (forbidden).343 344    Inspired by https://stackoverflow.com/a/33112180.345    """346    initial_sig = inspect.signature(func)347 348    @wraps(func)349    async def _protected_func(request: Request, **kwargs):350        request_secret = request.headers.get("x-webhook-secret")351        if request_secret is None:352            return JSONResponse({"error": "x-webhook-secret header not set."}, status_code=401)353        if request_secret != webhook_secret:354            return JSONResponse({"error": "Invalid webhook secret."}, status_code=403)355 356        # Inject `request` in kwargs if required357        if "request" in initial_sig.parameters:358            kwargs["request"] = request359 360        # Handle both sync and async routes361        if inspect.iscoroutinefunction(func):362            return await func(**kwargs)363        else:364            return func(**kwargs)365 366    # Update signature to include request367    if "request" not in initial_sig.parameters:368        _protected_func.__signature__ = initial_sig.replace(  # type: ignore369            parameters=(370                inspect.Parameter(name="request", kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=Request),371            )372            + tuple(initial_sig.parameters.values())373        )374 375    # Return protected route376    return _protected_func377