Coverage for haystack/core/pipeline/draw.py: 79%
199 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import base64
6import colorsys
7import json
8import random
9import zlib
10from typing import Any
12import httpx
13import networkx
15from haystack import logging
16from haystack.core.errors import PipelineDrawingError
17from haystack.core.pipeline.descriptions import find_pipeline_inputs, find_pipeline_outputs
18from haystack.core.type_utils import _type_name
20logger = logging.getLogger(__name__)
23def generate_color_variations(n: int, base_color: str | None = "#3498DB", variation_range: float = 0.4) -> list[str]:
24 """
25 Generate n different variations of a base color.
27 :param n: Number of variations to generate
28 :param base_color: Hex color code, default is a shade of blue (#3498DB)
29 :param variation_range: Range for varying brightness and saturation (0-1)
31 :returns:
32 list: List of hex color codes representing variations of the base color
33 """
34 # convert hex to RGB
35 base_color = base_color.lstrip("#") # type:ignore
36 r = int(base_color[0:2], 16) / 255.0
37 g = int(base_color[2:4], 16) / 255.0
38 b = int(base_color[4:6], 16) / 255.0
40 # convert RGB to HSV (Hue, Saturation, Value)
41 h, s, v = colorsys.rgb_to_hsv(r, g, b)
43 variations = []
44 for _ in range(n):
45 # vary saturation and brightness within the specified range
46 new_s = max(0, min(1, s + random.uniform(-variation_range, variation_range)))
47 new_v = max(0, min(1, v + random.uniform(-variation_range, variation_range)))
49 # keep hue the same for color consistency
50 new_h = h
52 # Convert back to RGB and then to hex
53 new_r, new_g, new_b = colorsys.hsv_to_rgb(new_h, new_s, new_v)
54 hex_color = f"#{int(new_r * 255):02x}{int(new_g * 255):02x}{int(new_b * 255):02x}"
56 variations.append(hex_color)
58 return variations
61def _prepare_for_drawing(graph: networkx.MultiDiGraph) -> networkx.MultiDiGraph:
62 """
63 Add some extra nodes to show the inputs and outputs of the pipeline.
65 Also adds labels to edges.
66 """
67 # Label the edges
68 for inp, outp, key, data in graph.edges(keys=True, data=True):
69 data["label"] = (
70 f"{data['from_socket'].name} -> {data['to_socket'].name}{' (opt.)' if not data['mandatory'] else ''}"
71 )
72 graph.add_edge(inp, outp, key=key, **data)
74 # Add inputs fake node
75 graph.add_node("input")
76 for node, in_sockets in find_pipeline_inputs(graph).items():
77 for in_socket in in_sockets:
78 if not in_socket.senders and in_socket.is_mandatory:
79 # If this socket has no sender it could be a socket that receives input
80 # directly when running the Pipeline. We can't know that for sure, in doubt
81 # we draw it as receiving input directly.
82 graph.add_edge("input", node, label=in_socket.name, conn_type=_type_name(in_socket.type))
84 # Add outputs fake node
85 graph.add_node("output")
86 for node, out_sockets in find_pipeline_outputs(graph).items():
87 for out_socket in out_sockets:
88 graph.add_edge(node, "output", label=out_socket.name, conn_type=_type_name(out_socket.type))
90 return graph
93ARROWTAIL_MANDATORY = "--"
94ARROWTAIL_OPTIONAL = "-."
95ARROWHEAD_MANDATORY = "-->"
96ARROWHEAD_OPTIONAL = ".->"
97MERMAID_STYLED_TEMPLATE = """
98%%{{ init: {params} }}%%
100graph TD;
102{connections}
104classDef component text-align:center;
105{style_definitions}
106"""
109def _validate_mermaid_params(params: dict[str, Any]) -> None:
110 """
111 Validates and sets default values for Mermaid parameters.
113 :param params:
114 Dictionary of customization parameters to modify the output. Refer to Mermaid documentation for more details.
115 Supported keys:
116 - format: Output format ('img', 'svg', or 'pdf'). Default: 'img'.
117 - type: Image type for /img endpoint ('jpeg', 'png', 'webp'). Default: 'png'.
118 - theme: Mermaid theme ('default', 'neutral', 'dark', 'forest'). Default: 'neutral'.
119 - bgColor: Background color in hexadecimal (e.g., 'FFFFFF') or named format (e.g., '!white').
120 - width: Width of the output image (integer).
121 - height: Height of the output image (integer).
122 - scale: Scaling factor (1–3). Only applicable if 'width' or 'height' is specified.
123 - fit: Whether to fit the diagram size to the page (PDF only, boolean).
124 - paper: Paper size for PDFs (e.g., 'a4', 'a3'). Ignored if 'fit' is true.
125 - landscape: Landscape orientation for PDFs (boolean). Ignored if 'fit' is true.
127 :raises ValueError:
128 If any parameter is invalid or does not match the expected format.
129 """
130 valid_img_types = {"jpeg", "png", "webp"}
131 valid_themes = {"default", "neutral", "dark", "forest"}
132 valid_formats = {"img", "svg", "pdf"}
134 params.setdefault("format", "img")
135 params.setdefault("type", "png")
136 params.setdefault("theme", "neutral")
138 if params["format"] not in valid_formats:
139 raise ValueError(f"Invalid image format: {params['format']}. Valid options are: {valid_formats}.")
141 if params["format"] == "img" and params["type"] not in valid_img_types:
142 raise ValueError(f"Invalid image type: {params['type']}. Valid options are: {valid_img_types}.")
144 if params["theme"] not in valid_themes:
145 raise ValueError(f"Invalid theme: {params['theme']}. Valid options are: {valid_themes}.")
147 if "width" in params and not isinstance(params["width"], int):
148 raise ValueError("Width must be an integer.")
149 if "height" in params and not isinstance(params["height"], int):
150 raise ValueError("Height must be an integer.")
152 if "scale" in params and not 1 <= params["scale"] <= 3:
153 raise ValueError("Scale must be a number between 1 and 3.")
154 if "scale" in params and not ("width" in params or "height" in params):
155 raise ValueError("Scale is only allowed when width or height is set.")
157 if "bgColor" in params and not isinstance(params["bgColor"], str):
158 raise ValueError("Background color must be a string.")
160 # PDF specific parameters
161 if params["format"] == "pdf":
162 if "fit" in params and not isinstance(params["fit"], bool):
163 raise ValueError("Fit must be a boolean.")
164 if "paper" in params and not isinstance(params["paper"], str):
165 raise ValueError("Paper size must be a string (e.g., 'a4', 'a3').")
166 if "landscape" in params and not isinstance(params["landscape"], bool):
167 raise ValueError("Landscape must be a boolean.")
168 if "fit" in params and ("paper" in params or "landscape" in params):
169 logger.warning("`fit` overrides `paper` and `landscape` for PDFs. Ignoring `paper` and `landscape`.")
172# Magic-byte signatures used to verify a Mermaid server response matches the requested output format.
173_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
174_JPEG_SIGNATURE = b"\xff\xd8\xff"
175_PDF_SIGNATURE = b"%PDF-"
176_RIFF_SIGNATURE = b"RIFF"
177_WEBP_SIGNATURE = b"WEBP"
178_SVG_PREFIXES = (b"<?xml", b"<svg")
181def _validate_image_response(resp: httpx.Response, params: dict[str, Any]) -> None:
182 """
183 Validate that the Mermaid server response actually contains the expected image/SVG/PDF data.
185 `Pipeline.draw()` writes the raw response body to disk, so a misconfigured or malicious
186 `server_url` could otherwise cause arbitrary content (e.g. an HTML error page or a crafted
187 payload) to be written verbatim to the output path. As defense-in-depth we check both the
188 `Content-Type` header (which the server controls and could spoof) and the response body's
189 magic-byte signature (which is harder to forge while still producing a usable payload).
191 :param resp:
192 The HTTP response returned by the Mermaid server.
193 :param params:
194 Validated Mermaid parameters; used to determine the expected output format.
195 :raises PipelineDrawingError:
196 If the response is empty or does not match the expected format.
197 """
198 content = resp.content
199 if not content:
200 raise PipelineDrawingError("The Mermaid server returned an empty response; no image will be saved.")
202 output_format = params.get("format", "img")
203 img_type = params.get("type", "png")
205 # (human-readable label, expected Content-Type prefix, body signature check)
206 content_type_prefixes: tuple[str, ...]
207 if output_format == "svg":
208 expected_label = "SVG"
209 content_type_prefixes = ("image/svg+xml", "text/xml", "application/xml")
210 stripped = content.lstrip()[:512].lower()
211 body_ok = stripped.startswith(_SVG_PREFIXES) or _SVG_PREFIXES[1] in stripped
212 elif output_format == "pdf":
213 expected_label = "PDF"
214 content_type_prefixes = ("application/pdf",)
215 body_ok = content.startswith(_PDF_SIGNATURE)
216 elif img_type == "jpeg":
217 expected_label = "JPEG image"
218 content_type_prefixes = ("image/jpeg",)
219 body_ok = content.startswith(_JPEG_SIGNATURE)
220 elif img_type == "webp":
221 expected_label = "WebP image"
222 content_type_prefixes = ("image/webp",)
223 body_ok = content[0:4] == _RIFF_SIGNATURE and content[8:12] == _WEBP_SIGNATURE
224 else: # png (default)
225 expected_label = "PNG image"
226 content_type_prefixes = ("image/png",)
227 body_ok = content.startswith(_PNG_SIGNATURE)
229 # The Content-Type header is server-controlled, so a mismatch is only a warning: the
230 # authoritative check is the body signature below.
231 content_type = resp.headers.get("content-type", "").split(";")[0].strip().lower()
232 if content_type and not content_type.startswith(content_type_prefixes):
233 logger.warning(
234 "The Mermaid server returned an unexpected Content-Type '{content_type}' (expected {expected}).",
235 content_type=content_type,
236 expected=expected_label,
237 )
239 if not body_ok:
240 raise PipelineDrawingError(
241 f"The Mermaid server response does not look like a valid {expected_label}. "
242 f"This can happen if 'server_url' points to a server that is not a Mermaid renderer. "
243 f"To avoid writing untrusted content to disk, no file will be saved."
244 )
247def _to_mermaid_image(
248 graph: networkx.MultiDiGraph,
249 server_url: str = "https://mermaid.ink",
250 params: dict | None = None,
251 timeout: int = 30,
252 super_component_mapping: dict[str, str] | None = None,
253) -> bytes:
254 """
255 Renders a pipeline using a Mermaid server.
257 :param graph:
258 The graph to render as a Mermaid pipeline.
259 :param server_url:
260 Base URL of the Mermaid server (default: 'https://mermaid.ink').
261 :param params:
262 Dictionary of customization parameters. See `validate_mermaid_params` for valid keys.
263 :param timeout:
264 Timeout in seconds for the request to the Mermaid server.
265 :param super_component_mapping:
266 Mapping of component names to the name of the SuperComponent that contains them, used to
267 group those components visually in the rendered diagram.
268 :returns:
269 The image, SVG, or PDF data returned by the Mermaid server as bytes.
270 :raises ValueError:
271 If any parameter is invalid or does not match the expected format.
272 :raises PipelineDrawingError:
273 If there is an issue connecting to the Mermaid server or the server returns an error.
274 """
276 if params is None:
277 params = {}
279 _validate_mermaid_params(params)
281 theme = params.get("theme")
282 init_params = json.dumps({"theme": theme})
284 # Copy the graph to avoid modifying the original
285 graph_styled = _to_mermaid_text(graph.copy(), init_params, super_component_mapping)
286 json_string = json.dumps({"code": graph_styled})
288 # Compress the JSON string with zlib (RFC 1950)
289 compressor = zlib.compressobj(level=9, wbits=15)
290 compressed_data = compressor.compress(json_string.encode("utf-8")) + compressor.flush()
291 compressed_url_safe_base64 = base64.urlsafe_b64encode(compressed_data).decode("utf-8").strip()
293 # Determine the correct endpoint
294 endpoint_format = params.get("format", "img") # Default to /img endpoint
295 if endpoint_format not in {"img", "svg", "pdf"}:
296 raise ValueError(f"Invalid format: {endpoint_format}. Valid options are 'img', 'svg', or 'pdf'.")
298 # Construct the URL without query parameters
299 url = f"{server_url}/{endpoint_format}/pako:{compressed_url_safe_base64}"
301 # Add query parameters adhering to mermaid.ink documentation
302 query_params = []
303 for key, value in params.items():
304 if key not in {"theme", "format"}: # Exclude theme (handled in init_params) and format (endpoint-specific)
305 if value is True:
306 query_params.append(f"{key}")
307 else:
308 query_params.append(f"{key}={value}")
310 if query_params:
311 url += "?" + "&".join(query_params)
313 logger.debug("Rendering graph at {url}", url=url)
314 try:
315 resp = httpx.get(url, timeout=timeout)
316 if resp.status_code >= 400:
317 logger.warning(
318 "Failed to draw the pipeline: {server_url} returned status {status_code}",
319 server_url=server_url,
320 status_code=resp.status_code,
321 )
322 logger.info("Exact URL requested: {url}", url=url)
323 logger.warning("No pipeline diagram will be saved.")
324 resp.raise_for_status()
326 except Exception as exc:
327 logger.warning(
328 "Failed to draw the pipeline: could not connect to {server_url} ({error})", server_url=server_url, error=exc
329 )
330 logger.info("Exact URL requested: {url}", url=url)
331 logger.warning("No pipeline diagram will be saved.")
332 raise PipelineDrawingError(f"There was an issue with {server_url}, see the stacktrace for details.") from exc
334 # Validate the response before it gets written to disk by the caller, so that a misconfigured
335 # or malicious server cannot cause arbitrary content to be saved to the output path.
336 _validate_image_response(resp, params)
338 return resp.content
341def _to_mermaid_text(
342 graph: networkx.MultiDiGraph, init_params: str | dict, super_component_mapping: dict[str, str] | None = None
343) -> str:
344 """
345 Converts a Networkx graph into Mermaid syntax.
347 The output of this function can be used in the documentation with `mermaid` codeblocks and will be
348 automatically rendered.
350 :param graph: The graph to convert to Mermaid syntax
351 :param init_params: Initialization parameters for Mermaid
352 :param super_component_mapping: Mapping of component names to super component names
353 """
354 # Copy the graph to avoid modifying the original
355 graph = _prepare_for_drawing(graph.copy())
356 sockets = {
357 comp: "".join(
358 [
359 f"<li>{name} ({_type_name(socket.type)})</li>"
360 for name, socket in data.get("input_sockets", {}).items()
361 if (not socket.is_mandatory and not socket.senders) or socket.is_variadic
362 ]
363 )
364 for comp, data in graph.nodes(data=True)
365 }
366 optional_inputs = {
367 comp: f"<br><br>Optional inputs:<ul style='text-align:left;'>{sockets}</ul>" if sockets else ""
368 for comp, sockets in sockets.items()
369 }
371 # Create node definitions
372 states = {}
373 super_component_components = super_component_mapping.keys() if super_component_mapping else {}
375 # color variations for super components
376 super_component_colors = {}
377 if super_component_components:
378 unique_super_components = set(super_component_mapping.values()) # type:ignore
379 color_variations = generate_color_variations(n=len(unique_super_components))
380 super_component_colors = dict(zip(unique_super_components, color_variations, strict=True))
382 # Generate style definitions for each super component
383 style_definitions = []
384 for super_comp, color in super_component_colors.items():
385 style_definitions.append(f"classDef {super_comp} fill:{color},color:white;")
387 for comp, data in graph.nodes(data=True):
388 if comp in ["input", "output"]:
389 continue
391 # styling based on whether the component is a SuperComponent
392 if comp in super_component_components:
393 super_component_name = super_component_mapping[comp] # type:ignore
394 style = super_component_name
395 else:
396 style = "component"
397 node_def = f'{comp}["<b>{comp}</b><br><small><i>{type(data["instance"]).__name__}{optional_inputs[comp]}</i></small>"]:::{style}' # noqa: E501
398 states[comp] = node_def
400 connections_list = []
401 for from_comp, to_comp, conn_data in graph.edges(data=True):
402 if from_comp != "input" and to_comp != "output":
403 arrowtail = ARROWTAIL_MANDATORY if conn_data["mandatory"] else ARROWTAIL_OPTIONAL
404 arrowhead = ARROWHEAD_MANDATORY if conn_data["mandatory"] else ARROWHEAD_OPTIONAL
405 label = f'"{conn_data["label"]}<br><small><i>{conn_data["conn_type"]}</i></small>"'
406 conn_string = f"{states[from_comp]} {arrowtail} {label} {arrowhead} {states[to_comp]}"
407 connections_list.append(conn_string)
409 input_connections = [
410 f'i{{*}}--"{conn_data["label"]}<br><small><i>{conn_data["conn_type"]}</i></small>"--> {states[to_comp]}'
411 for _, to_comp, conn_data in graph.out_edges("input", data=True)
412 ]
413 output_connections = [
414 f'{states[from_comp]}--"{conn_data["label"]}<br><small><i>{conn_data["conn_type"]}</i></small>"--> o{{*}}'
415 for from_comp, _, conn_data in graph.in_edges("output", data=True)
416 ]
417 connections = "\n".join(connections_list + input_connections + output_connections)
419 # Create legend
420 legend_nodes = []
421 if super_component_colors:
422 legend_nodes.append("subgraph Legend")
423 for super_comp in super_component_colors:
424 legend_id = f"legend_{super_comp}"
425 legend_nodes.append(f'{legend_id}["{super_comp}"]:::{super_comp}')
426 legend_nodes.append("end")
427 connections += "\n" + "\n".join(legend_nodes)
429 # Add style definitions to the template
430 graph_styled = MERMAID_STYLED_TEMPLATE.format(
431 params=init_params, connections=connections, style_definitions="\n".join(style_definitions)
432 )
433 logger.debug("Mermaid diagram:\n{diagram}", diagram=graph_styled)
435 return graph_styled