/haystack/core/pipeline/draw.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import base64
import colorsys
import json
import random
import zlib
from typing import Any

import httpx
import networkx

from haystack import logging
from haystack.core.errors import PipelineDrawingError
from haystack.core.pipeline.descriptions import find_pipeline_inputs, find_pipeline_outputs
from haystack.core.type_utils import _type_name

logger = logging.getLogger(__name__)


def generate_color_variations(n: int, base_color: str | None = "#3498DB", variation_range: float = 0.4) -> list[str]:
    """
    Generate n different variations of a base color.

    :param n: Number of variations to generate
    :param base_color: Hex color code, default is a shade of blue (#3498DB)
    :param variation_range: Range for varying brightness and saturation (0-1)

    :returns:
        list: List of hex color codes representing variations of the base color
    """
    # convert hex to RGB
    base_color = base_color.lstrip("#")  # type:ignore
    r = int(base_color[0:2], 16) / 255.0
    g = int(base_color[2:4], 16) / 255.0
    b = int(base_color[4:6], 16) / 255.0

    # convert RGB to HSV (Hue, Saturation, Value)
    h, s, v = colorsys.rgb_to_hsv(r, g, b)

    variations = []
    for _ in range(n):
        # vary saturation and brightness within the specified range
        new_s = max(0, min(1, s + random.uniform(-variation_range, variation_range)))
        new_v = max(0, min(1, v + random.uniform(-variation_range, variation_range)))

        # keep hue the same for color consistency
        new_h = h

        # Convert back to RGB and then to hex
        new_r, new_g, new_b = colorsys.hsv_to_rgb(new_h, new_s, new_v)
        hex_color = f"#{int(new_r * 255):02x}{int(new_g * 255):02x}{int(new_b * 255):02x}"

        variations.append(hex_color)

    return variations


def _prepare_for_drawing(graph: networkx.MultiDiGraph) -> networkx.MultiDiGraph:
    """
    Add some extra nodes to show the inputs and outputs of the pipeline.

    Also adds labels to edges.
    """
    # Label the edges
    for inp, outp, key, data in graph.edges(keys=True, data=True):
        data["label"] = (
            f"{data['from_socket'].name} -> {data['to_socket'].name}{' (opt.)' if not data['mandatory'] else ''}"
        )
        graph.add_edge(inp, outp, key=key, **data)

    # Add inputs fake node
    graph.add_node("input")
    for node, in_sockets in find_pipeline_inputs(graph).items():
        for in_socket in in_sockets:
            if not in_socket.senders and in_socket.is_mandatory:
                # If this socket has no sender it could be a socket that receives input
                # directly when running the Pipeline. We can't know that for sure, in doubt
                # we draw it as receiving input directly.
                graph.add_edge("input", node, label=in_socket.name, conn_type=_type_name(in_socket.type))

    # Add outputs fake node
    graph.add_node("output")
    for node, out_sockets in find_pipeline_outputs(graph).items():
        for out_socket in out_sockets:
            graph.add_edge(node, "output", label=out_socket.name, conn_type=_type_name(out_socket.type))

    return graph


ARROWTAIL_MANDATORY = "--"
ARROWTAIL_OPTIONAL = "-."
ARROWHEAD_MANDATORY = "-->"
ARROWHEAD_OPTIONAL = ".->"
MERMAID_STYLED_TEMPLATE = """
%%{{ init: {params} }}%%

graph TD;

{connections}

classDef component text-align:center;
{style_definitions}
"""


def _validate_mermaid_params(params: dict[str, Any]) -> None:
    """
    Validates and sets default values for Mermaid parameters.

    :param params:
        Dictionary of customization parameters to modify the output. Refer to Mermaid documentation for more details.
        Supported keys:
            - format: Output format ('img', 'svg', or 'pdf'). Default: 'img'.
            - type: Image type for /img endpoint ('jpeg', 'png', 'webp'). Default: 'png'.
            - theme: Mermaid theme ('default', 'neutral', 'dark', 'forest'). Default: 'neutral'.
            - bgColor: Background color in hexadecimal (e.g., 'FFFFFF') or named format (e.g., '!white').
            - width: Width of the output image (integer).
            - height: Height of the output image (integer).
            - scale: Scaling factor (1–3). Only applicable if 'width' or 'height' is specified.
            - fit: Whether to fit the diagram size to the page (PDF only, boolean).
            - paper: Paper size for PDFs (e.g., 'a4', 'a3'). Ignored if 'fit' is true.
            - landscape: Landscape orientation for PDFs (boolean). Ignored if 'fit' is true.

    :raises ValueError:
        If any parameter is invalid or does not match the expected format.
    """
    valid_img_types = {"jpeg", "png", "webp"}
    valid_themes = {"default", "neutral", "dark", "forest"}
    valid_formats = {"img", "svg", "pdf"}

    params.setdefault("format", "img")
    params.setdefault("type", "png")
    params.setdefault("theme", "neutral")

    if params["format"] not in valid_formats:
        raise ValueError(f"Invalid image format: {params['format']}. Valid options are: {valid_formats}.")

    if params["format"] == "img" and params["type"] not in valid_img_types:
        raise ValueError(f"Invalid image type: {params['type']}. Valid options are: {valid_img_types}.")

    if params["theme"] not in valid_themes:
        raise ValueError(f"Invalid theme: {params['theme']}. Valid options are: {valid_themes}.")

    if "width" in params and not isinstance(params["width"], int):
        raise ValueError("Width must be an integer.")
    if "height" in params and not isinstance(params["height"], int):
        raise ValueError("Height must be an integer.")

    if "scale" in params and not 1 <= params["scale"] <= 3:
        raise ValueError("Scale must be a number between 1 and 3.")
    if "scale" in params and not ("width" in params or "height" in params):
        raise ValueError("Scale is only allowed when width or height is set.")

    if "bgColor" in params and not isinstance(params["bgColor"], str):
        raise ValueError("Background color must be a string.")

    # PDF specific parameters
    if params["format"] == "pdf":
        if "fit" in params and not isinstance(params["fit"], bool):
            raise ValueError("Fit must be a boolean.")
        if "paper" in params and not isinstance(params["paper"], str):
            raise ValueError("Paper size must be a string (e.g., 'a4', 'a3').")
        if "landscape" in params and not isinstance(params["landscape"], bool):
            raise ValueError("Landscape must be a boolean.")
        if "fit" in params and ("paper" in params or "landscape" in params):
            logger.warning("`fit` overrides `paper` and `landscape` for PDFs. Ignoring `paper` and `landscape`.")


# Magic-byte signatures used to verify a Mermaid server response matches the requested output format.
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
_JPEG_SIGNATURE = b"\xff\xd8\xff"
_PDF_SIGNATURE = b"%PDF-"
_RIFF_SIGNATURE = b"RIFF"
_WEBP_SIGNATURE = b"WEBP"
_SVG_PREFIXES = (b"<?xml", b"<svg")


def _validate_image_response(resp: httpx.Response, params: dict[str, Any]) -> None:
    """
    Validate that the Mermaid server response actually contains the expected image/SVG/PDF data.

    `Pipeline.draw()` writes the raw response body to disk, so a misconfigured or malicious
    `server_url` could otherwise cause arbitrary content (e.g. an HTML error page or a crafted
    payload) to be written verbatim to the output path. As defense-in-depth we check both the
    `Content-Type` header (which the server controls and could spoof) and the response body's
    magic-byte signature (which is harder to forge while still producing a usable payload).

    :param resp:
        The HTTP response returned by the Mermaid server.
    :param params:
        Validated Mermaid parameters; used to determine the expected output format.
    :raises PipelineDrawingError:
        If the response is empty or does not match the expected format.
    """
    content = resp.content
    if not content:
        raise PipelineDrawingError("The Mermaid server returned an empty response; no image will be saved.")

    output_format = params.get("format", "img")
    img_type = params.get("type", "png")

    # (human-readable label, expected Content-Type prefix, body signature check)
    content_type_prefixes: tuple[str, ...]
    if output_format == "svg":
        expected_label = "SVG"
        content_type_prefixes = ("image/svg+xml", "text/xml", "application/xml")
        stripped = content.lstrip()[:512].lower()
        body_ok = stripped.startswith(_SVG_PREFIXES) or _SVG_PREFIXES[1] in stripped
    elif output_format == "pdf":
        expected_label = "PDF"
        content_type_prefixes = ("application/pdf",)
        body_ok = content.startswith(_PDF_SIGNATURE)
    elif img_type == "jpeg":
        expected_label = "JPEG image"
        content_type_prefixes = ("image/jpeg",)
        body_ok = content.startswith(_JPEG_SIGNATURE)
    elif img_type == "webp":
        expected_label = "WebP image"
        content_type_prefixes = ("image/webp",)
        body_ok = content[0:4] == _RIFF_SIGNATURE and content[8:12] == _WEBP_SIGNATURE
    else:  # png (default)
        expected_label = "PNG image"
        content_type_prefixes = ("image/png",)
        body_ok = content.startswith(_PNG_SIGNATURE)

    # The Content-Type header is server-controlled, so a mismatch is only a warning: the
    # authoritative check is the body signature below.
    content_type = resp.headers.get("content-type", "").split(";")[0].strip().lower()
    if content_type and not content_type.startswith(content_type_prefixes):
        logger.warning(
            "The Mermaid server returned an unexpected Content-Type '{content_type}' (expected {expected}).",
            content_type=content_type,
            expected=expected_label,
        )

    if not body_ok:
        raise PipelineDrawingError(
            f"The Mermaid server response does not look like a valid {expected_label}. "
            f"This can happen if 'server_url' points to a server that is not a Mermaid renderer. "
            f"To avoid writing untrusted content to disk, no file will be saved."
        )


def _to_mermaid_image(
    graph: networkx.MultiDiGraph,
    server_url: str = "https://mermaid.ink",
    params: dict | None = None,
    timeout: int = 30,
    super_component_mapping: dict[str, str] | None = None,
) -> bytes:
    """
    Renders a pipeline using a Mermaid server.

    :param graph:
        The graph to render as a Mermaid pipeline.
    :param server_url:
        Base URL of the Mermaid server (default: 'https://mermaid.ink').
    :param params:
        Dictionary of customization parameters. See `validate_mermaid_params` for valid keys.
    :param timeout:
        Timeout in seconds for the request to the Mermaid server.
    :returns:
        The image, SVG, or PDF data returned by the Mermaid server as bytes.
    :raises ValueError:
        If any parameter is invalid or does not match the expected format.
    :raises PipelineDrawingError:
        If there is an issue connecting to the Mermaid server or the server returns an error.
    """

    if params is None:
        params = {}

    _validate_mermaid_params(params)

    theme = params.get("theme")
    init_params = json.dumps({"theme": theme})

    # Copy the graph to avoid modifying the original
    graph_styled = _to_mermaid_text(graph.copy(), init_params, super_component_mapping)
    json_string = json.dumps({"code": graph_styled})

    # Compress the JSON string with zlib (RFC 1950)
    compressor = zlib.compressobj(level=9, wbits=15)
    compressed_data = compressor.compress(json_string.encode("utf-8")) + compressor.flush()
    compressed_url_safe_base64 = base64.urlsafe_b64encode(compressed_data).decode("utf-8").strip()

    # Determine the correct endpoint
    endpoint_format = params.get("format", "img")  # Default to /img endpoint
    if endpoint_format not in {"img", "svg", "pdf"}:
        raise ValueError(f"Invalid format: {endpoint_format}. Valid options are 'img', 'svg', or 'pdf'.")

    # Construct the URL without query parameters
    url = f"{server_url}/{endpoint_format}/pako:{compressed_url_safe_base64}"

    # Add query parameters adhering to mermaid.ink documentation
    query_params = []
    for key, value in params.items():
        if key not in {"theme", "format"}:  # Exclude theme (handled in init_params) and format (endpoint-specific)
            if value is True:
                query_params.append(f"{key}")
            else:
                query_params.append(f"{key}={value}")

    if query_params:
        url += "?" + "&".join(query_params)

    logger.debug("Rendering graph at {url}", url=url)
    try:
        resp = httpx.get(url, timeout=timeout)
        if resp.status_code >= 400:
            logger.warning(
                "Failed to draw the pipeline: {server_url} returned status {status_code}",
                server_url=server_url,
                status_code=resp.status_code,
            )
            logger.info("Exact URL requested: {url}", url=url)
            logger.warning("No pipeline diagram will be saved.")
            resp.raise_for_status()

    except Exception as exc:
        logger.warning(
            "Failed to draw the pipeline: could not connect to {server_url} ({error})", server_url=server_url, error=exc
        )
        logger.info("Exact URL requested: {url}", url=url)
        logger.warning("No pipeline diagram will be saved.")
        raise PipelineDrawingError(f"There was an issue with {server_url}, see the stacktrace for details.") from exc

    # Validate the response before it gets written to disk by the caller, so that a misconfigured
    # or malicious server cannot cause arbitrary content to be saved to the output path.
    _validate_image_response(resp, params)

    return resp.content


def _to_mermaid_text(
    graph: networkx.MultiDiGraph, init_params: str | dict, super_component_mapping: dict[str, str] | None = None
) -> str:
    """
    Converts a Networkx graph into Mermaid syntax.

    The output of this function can be used in the documentation with `mermaid` codeblocks and will be
    automatically rendered.

    :param graph: The graph to convert to Mermaid syntax
    :param init_params: Initialization parameters for Mermaid
    :param super_component_mapping: Mapping of component names to super component names
    """
    # Copy the graph to avoid modifying the original
    graph = _prepare_for_drawing(graph.copy())
    sockets = {
        comp: "".join(
            [
                f"<li>{name} ({_type_name(socket.type)})</li>"
                for name, socket in data.get("input_sockets", {}).items()
                if (not socket.is_mandatory and not socket.senders) or socket.is_variadic
            ]
        )
        for comp, data in graph.nodes(data=True)
    }
    optional_inputs = {
        comp: f"<br><br>Optional inputs:<ul style='text-align:left;'>{sockets}</ul>" if sockets else ""
        for comp, sockets in sockets.items()
    }

    # Create node definitions
    states = {}
    super_component_components = super_component_mapping.keys() if super_component_mapping else {}

    # color variations for super components
    super_component_colors = {}
    if super_component_components:
        unique_super_components = set(super_component_mapping.values())  # type:ignore
        color_variations = generate_color_variations(n=len(unique_super_components))
        super_component_colors = dict(zip(unique_super_components, color_variations, strict=True))

    # Generate style definitions for each super component
    style_definitions = []
    for super_comp, color in super_component_colors.items():
        style_definitions.append(f"classDef {super_comp} fill:{color},color:white;")

    for comp, data in graph.nodes(data=True):
        if comp in ["input", "output"]:
            continue

        # styling based on whether the component is a SuperComponent
        if comp in super_component_components:
            super_component_name = super_component_mapping[comp]  # type:ignore
            style = super_component_name
        else:
            style = "component"
        node_def = f'{comp}["<b>{comp}</b><br><small><i>{type(data["instance"]).__name__}{optional_inputs[comp]}</i></small>"]:::{style}'  # noqa: E501
        states[comp] = node_def

    connections_list = []
    for from_comp, to_comp, conn_data in graph.edges(data=True):
        if from_comp != "input" and to_comp != "output":
            arrowtail = ARROWTAIL_MANDATORY if conn_data["mandatory"] else ARROWTAIL_OPTIONAL
            arrowhead = ARROWHEAD_MANDATORY if conn_data["mandatory"] else ARROWHEAD_OPTIONAL
            label = f'"{conn_data["label"]}<br><small><i>{conn_data["conn_type"]}</i></small>"'
            conn_string = f"{states[from_comp]} {arrowtail} {label} {arrowhead} {states[to_comp]}"
            connections_list.append(conn_string)

    input_connections = [
        f'i{{&ast;}}--"{conn_data["label"]}<br><small><i>{conn_data["conn_type"]}</i></small>"--> {states[to_comp]}'
        for _, to_comp, conn_data in graph.out_edges("input", data=True)
    ]
    output_connections = [
        f'{states[from_comp]}--"{conn_data["label"]}<br><small><i>{conn_data["conn_type"]}</i></small>"--> o{{&ast;}}'
        for from_comp, _, conn_data in graph.in_edges("output", data=True)
    ]
    connections = "\n".join(connections_list + input_connections + output_connections)

    # Create legend
    legend_nodes = []
    if super_component_colors:
        legend_nodes.append("subgraph Legend")
        for super_comp in super_component_colors:
            legend_id = f"legend_{super_comp}"
            legend_nodes.append(f'{legend_id}["{super_comp}"]:::{super_comp}')
        legend_nodes.append("end")
        connections += "\n" + "\n".join(legend_nodes)

    # Add style definitions to the template
    graph_styled = MERMAID_STYLED_TEMPLATE.format(
        params=init_params, connections=connections, style_definitions="\n".join(style_definitions)
    )
    logger.debug("Mermaid diagram:\n{diagram}", diagram=graph_styled)

    return graph_styled