Webhook tasks using FlyteAgents

flytekit/extras/webhook/__init__.py

@@ -0,0 +1,4 @@
+from .agent import WebhookAgent
+from .task import WebhookTask
+
+__all__ = ["WebhookTask", "WebhookAgent"]

flytekit/extras/webhook/agent.py

@@ -0,0 +1,81 @@
+import asyncio
+import http
+from typing import Optional
+
+import aiohttp
+from flyteidl.core.execution_pb2 import TaskExecution
+
+from flytekit.extend.backend.base_agent import AgentRegistry, Resource, SyncAgentBase
+from flytekit.interaction.string_literals import literal_map_string_repr
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+from flytekit.utils.dict_formatter import format_dict
+
+from .constants import BODY_KEY, HEADERS_KEY, METHOD_KEY, SHOW_BODY_KEY, SHOW_URL_KEY, TASK_TYPE, URL_KEY
+
+
+class WebhookAgent(SyncAgentBase):
+    name = "Webhook Agent"
+
+    def __init__(self):
+        super().__init__(task_type_name=TASK_TYPE)
+        self._session = None
+        self._lock = asyncio.Lock()
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        async with self._lock:
+            if self._session is None:
+                self._session = aiohttp.ClientSession()
+            return self._session
+
+    async def do(
+        self, task_template: TaskTemplate, output_prefix: str, inputs: Optional[LiteralMap] = None, **kwargs
+    ) -> Resource:
+        try:
+            custom_dict = task_template.custom
+            input_dict = {
+                "inputs": literal_map_string_repr(inputs),
+            }
+
+            final_dict = format_dict("test", custom_dict, input_dict)
+            url = final_dict.get(URL_KEY)
+            body = final_dict.get(BODY_KEY)
+            headers = final_dict.get(HEADERS_KEY)
+            method = final_dict.get(METHOD_KEY)
+            method = http.HTTPMethod(method)
+            show_body = final_dict.get(SHOW_BODY_KEY, False)
+            show_url = final_dict.get(SHOW_URL_KEY, False)
+
+            session = await self._get_session()
+
+            text = None
+            if method == http.HTTPMethod.GET:
+                response = await session.get(url, headers=headers)
+                text = await response.text()
+            else:
+                response = await session.post(url, json=body, headers=headers)
+                text = await response.text()
+            if response.status != 200:
+                return Resource(
+                    phase=TaskExecution.FAILED,
+                    message=f"Webhook failed with status code {response.status}, response: {text}",
+                )
+            final_response = {
+                "status_code": response.status,
+                "body": text,
+            }
+            if show_body:
+                final_response["input_body"] = body
+            if show_url:
+                final_response["url"] = url
+
+            return Resource(
+                phase=TaskExecution.SUCCEEDED,
+                outputs={"info": final_response},
+                message="Webhook was successfully invoked!",
+            )
+        except Exception as e:
+            return Resource(phase=TaskExecution.FAILED, message=str(e))
+
+
+AgentRegistry.register(WebhookAgent())

flytekit/extras/webhook/constants.py

@@ -0,0 +1,8 @@
+TASK_TYPE = "webhook"
+
+URL_KEY = "url"
+METHOD_KEY = "method"
+HEADERS_KEY = "headers"
+BODY_KEY = "body"
+SHOW_BODY_KEY = "show_body"
+SHOW_URL_KEY = "show_url"

flytekit/extras/webhook/task.py

@@ -0,0 +1,110 @@
+import http
+from typing import Any, Dict, Optional, Type
+
+from flytekit import Documentation
+from flytekit.configuration import SerializationSettings
+from flytekit.core.base_task import PythonTask
+from flytekit.extend.backend.base_agent import SyncAgentExecutorMixin
+
+from ...core.interface import Interface
+from .constants import BODY_KEY, HEADERS_KEY, METHOD_KEY, SHOW_BODY_KEY, SHOW_URL_KEY, TASK_TYPE, URL_KEY
+
+
+class WebhookTask(SyncAgentExecutorMixin, PythonTask):
+    """
+    This is the simplest form of a BigQuery Task, that can be used even for tasks that do not produce any output.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        method: http.HTTPMethod = http.HTTPMethod.POST,
+        headers: Optional[Dict[str, str]] = None,
+        body: Optional[Dict[str, Any]] = None,
+        dynamic_inputs: Optional[Dict[str, Type]] = None,
+        show_body: bool = False,
+        show_url: bool = False,
+        description: Optional[str] = None,
+        # secret_requests: Optional[List[Secret]] = None,  TODO Secret support is coming soon
+    ):
+        """
+        This task is used to invoke a webhook. The webhook can be invoked with a POST or GET method.
+
+        All the parameters can be formatted using python format strings. The following parameters are available for
+        formatting:
+        - dynamic_inputs: These are the dynamic inputs to the task. The keys are the names of the inputs and the values
+            are the values of the inputs. All inputs are available under the prefix `inputs.`.
+            For example, if the inputs are {"input1": 10, "input2": "hello"}, then you can
+            use {inputs.input1} and {inputs.input2} in the URL and the body. Define the dynamic_inputs argument in the
+            constructor to use these inputs. The dynamic inputs should not be actual values, but the types of the inputs.
+
+        TODO Coming soon secrets support
+        - secrets: These are the secrets that are requested by the task. The keys are the names of the secrets and the
+            values are the values of the secrets. All secrets are available under the prefix `secrets.`.
+            For example, if the secret requested are Secret(name="secret1") and Secret(name="secret), then you can use
+            {secrets.secret1} and {secrets.secret2} in the URL and the body. Define the secret_requests argument in the
+            constructor to use these secrets. The secrets should not be actual values, but the types of the secrets.
+
+        :param name: Name of this task, should be unique in the project
+        :param url: The endpoint or URL to invoke for this webhook. This can be a static string or a python format string,
+            where the format arguments are the dynamic_inputs to the task, secrets etc. Refer to the description for more
+            details of available formatting parameters.
+        :param method: The HTTP method to use for the request. Default is POST.
+        :param headers: The headers to send with the request. This can be a static dictionary or a python format string,
+            where the format arguments are the dynamic_inputs to the task, secrets etc. Refer to the description for more
+            details of available formatting parameters.
+        :param body: The body to send with the request. This can be a static dictionary or a python format string,
+            where the format arguments are the dynamic_inputs to the task, secrets etc. Refer to the description for more
+            details of available formatting parameters.
+        :param dynamic_inputs: The dynamic inputs to the task. The keys are the names of the inputs and the values
+            are the types of the inputs. These inputs are available under the prefix `inputs.` to be used in the URL,
+            headers and body and other formatted fields.
+        :param secret_requests: The secrets that are requested by the task. (TODO not yet supported)
+        :param show_body: If True, the body of the request will be logged in the UI as the output of the task.
+        :param show_url: If True, the URL of the request will be logged in the UI as the output of the task.
+        :param description: Description of the task
+        """
+        if method not in {http.HTTPMethod.GET, http.HTTPMethod.POST}:
+            raise ValueError(f"Method should be either GET or POST. Got {method}")
+        if method == http.HTTPMethod.GET:
+            if body:
+                raise ValueError("GET method cannot have a body")
+            if show_body:
+                raise ValueError("GET method cannot show body")
+        outputs = {
+            "status_code": int,
+        }
+        if show_body:
+            outputs["body"] = dict
+        if show_url:
+            outputs["url"] = bool
+
+        interface = Interface(
+            inputs=dynamic_inputs or {},
+            outputs={"info": dict},
+        )
+        super().__init__(
+            name=name,
+            interface=interface,
+            task_type=TASK_TYPE,
+            # secret_requests=secret_requests,
+            docs=Documentation(short_description=description) if description else None,
+        )
+        self._url = url
+        self._method = method
+        self._headers = headers
+        self._body = body
+        self._show_body = show_body
+        self._show_url = show_url
+
+    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
+        config = {
+            URL_KEY: self._url,
+            METHOD_KEY: self._method.value,
+            HEADERS_KEY: self._headers or {},
+            BODY_KEY: self._body or {},
+            SHOW_BODY_KEY: self._show_body,
+            SHOW_URL_KEY: self._show_url,
+        }
+        return config

flytekit/utils/dict_formatter.py

@@ -0,0 +1,86 @@
+import re
+from typing import Any, Dict, Optional
+
+
+def get_nested_value(d: Dict[str, Any], keys: list[str]) -> Any:
+    """
+    Retrieve the nested value from a dictionary based on a list of keys.
+    """
+    for key in keys:
+        if key not in d:
+            raise ValueError(f"Could not find the key {key} in {d}.")
+        d = d[key]
+    return d
+
+
+def replace_placeholder(
+    service: str,
+    original_dict: str,
+    placeholder: str,
+    replacement: str,
+) -> str:
+    """
+    Replace a placeholder in the original string and handle the specific logic for the sagemaker service and idempotence token.
+    """
+    temp_dict = original_dict.replace(f"{{{placeholder}}}", replacement)
+    if service == "sagemaker" and placeholder in [
+        "inputs.idempotence_token",
+        "idempotence_token",
+    ]:
+        if len(temp_dict) > 63:
+            truncated_token = replacement[: 63 - len(original_dict.replace(f"{{{placeholder}}}", ""))]
+            return original_dict.replace(f"{{{placeholder}}}", truncated_token)
+        else:
+            return temp_dict
+    return temp_dict
+
+
+def format_dict(
+    service: str,
+    original_dict: Any,
+    update_dict: Dict[str, Any],
+    idempotence_token: Optional[str] = None,
+) -> Any:
+    """
+    Recursively update a dictionary with format strings with values from another dictionary where the keys match
+    the format string. This goes a little beyond regular python string formatting and uses `.` to denote nested keys.
+
+    For example, if original_dict is {"EndpointConfigName": "{endpoint_config_name}"},
+    and update_dict is {"endpoint_config_name": "my-endpoint-config"},
+    then the result will be {"EndpointConfigName": "my-endpoint-config"}.
+
+    For nested keys if the original_dict is {"EndpointConfigName": "{inputs.endpoint_config_name}"},
+    and update_dict is {"inputs": {"endpoint_config_name": "my-endpoint-config"}},
+    then the result will be {"EndpointConfigName": "my-endpoint-config"}.
+
+    :param service: The AWS service to use
+    :param original_dict: The dictionary to update (in place)
+    :param update_dict: The dictionary to use for updating
+    :param idempotence_token: Hash of config -- this is to ensure the execution ID is deterministic
+    :return: The updated dictionary
+    """
+    if original_dict is None:
+        return None
+
+    if isinstance(original_dict, str) and "{" in original_dict and "}" in original_dict:
+        matches = re.findall(r"\{([^}]+)\}", original_dict)
+        for match in matches:
+            if "." in match:
+                keys = match.split(".")
+                nested_value = get_nested_value(update_dict, keys)
+                if f"{{{match}}}" == original_dict:
+                    return nested_value
+                else:
+                    original_dict = replace_placeholder(service, original_dict, match, nested_value)
+            elif match == "idempotence_token" and idempotence_token:
+                original_dict = replace_placeholder(service, original_dict, match, idempotence_token)
+        return original_dict
+
+    if isinstance(original_dict, list):
+        return [format_dict(service, item, update_dict, idempotence_token) for item in original_dict]
+
+    if isinstance(original_dict, dict):
+        for key, value in original_dict.items():
+            original_dict[key] = format_dict(service, value, update_dict, idempotence_token)
+
+    return original_dict