core[patch]: add retries and better messages to draw_mermaid_png (#30881)

This commit is contained in:
Vadym Barda 2025-04-17 14:25:37 -04:00 committed by GitHub
parent 75e50a3efd
commit d2cbfa379f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 73 additions and 10 deletions

View File

@ -630,6 +630,8 @@ class Graph:
draw_method: MermaidDrawMethod = MermaidDrawMethod.API, draw_method: MermaidDrawMethod = MermaidDrawMethod.API,
background_color: str = "white", background_color: str = "white",
padding: int = 10, padding: int = 10,
max_retries: int = 1,
retry_delay: float = 1.0,
frontmatter_config: Optional[dict[str, Any]] = None, frontmatter_config: Optional[dict[str, Any]] = None,
) -> bytes: ) -> bytes:
"""Draw the graph as a PNG image using Mermaid. """Draw the graph as a PNG image using Mermaid.
@ -645,6 +647,10 @@ class Graph:
Defaults to MermaidDrawMethod.API. Defaults to MermaidDrawMethod.API.
background_color: The color of the background. Defaults to "white". background_color: The color of the background. Defaults to "white".
padding: The padding around the graph. Defaults to 10. padding: The padding around the graph. Defaults to 10.
max_retries: The maximum number of retries (MermaidDrawMethod.API).
Defaults to 1.
retry_delay: The delay between retries (MermaidDrawMethod.API).
Defaults to 1.0.
frontmatter_config (dict[str, Any], optional): Mermaid frontmatter config. frontmatter_config (dict[str, Any], optional): Mermaid frontmatter config.
Can be used to customize theme and styles. Will be converted to YAML and Can be used to customize theme and styles. Will be converted to YAML and
added to the beginning of the mermaid graph. Defaults to None. added to the beginning of the mermaid graph. Defaults to None.
@ -680,6 +686,8 @@ class Graph:
draw_method=draw_method, draw_method=draw_method,
background_color=background_color, background_color=background_color,
padding=padding, padding=padding,
max_retries=max_retries,
retry_delay=retry_delay,
) )

View File

@ -2,7 +2,9 @@
import asyncio import asyncio
import base64 import base64
import random
import re import re
import time
from dataclasses import asdict from dataclasses import asdict
from pathlib import Path from pathlib import Path
from typing import Any, Literal, Optional from typing import Any, Literal, Optional
@ -254,6 +256,8 @@ def draw_mermaid_png(
draw_method: MermaidDrawMethod = MermaidDrawMethod.API, draw_method: MermaidDrawMethod = MermaidDrawMethod.API,
background_color: Optional[str] = "white", background_color: Optional[str] = "white",
padding: int = 10, padding: int = 10,
max_retries: int = 1,
retry_delay: float = 1.0,
) -> bytes: ) -> bytes:
"""Draws a Mermaid graph as PNG using provided syntax. """Draws a Mermaid graph as PNG using provided syntax.
@ -266,6 +270,10 @@ def draw_mermaid_png(
background_color (str, optional): Background color of the image. background_color (str, optional): Background color of the image.
Defaults to "white". Defaults to "white".
padding (int, optional): Padding around the image. Defaults to 10. padding (int, optional): Padding around the image. Defaults to 10.
max_retries (int, optional): Maximum number of retries (MermaidDrawMethod.API).
Defaults to 1.
retry_delay (float, optional): Delay between retries (MermaidDrawMethod.API).
Defaults to 1.0.
Returns: Returns:
bytes: PNG image bytes. bytes: PNG image bytes.
@ -283,7 +291,11 @@ def draw_mermaid_png(
) )
elif draw_method == MermaidDrawMethod.API: elif draw_method == MermaidDrawMethod.API:
img_bytes = _render_mermaid_using_api( img_bytes = _render_mermaid_using_api(
mermaid_syntax, output_file_path, background_color mermaid_syntax,
output_file_path=output_file_path,
background_color=background_color,
max_retries=max_retries,
retry_delay=retry_delay,
) )
else: else:
supported_methods = ", ".join([m.value for m in MermaidDrawMethod]) supported_methods = ", ".join([m.value for m in MermaidDrawMethod])
@ -371,9 +383,12 @@ async def _render_mermaid_using_pyppeteer(
def _render_mermaid_using_api( def _render_mermaid_using_api(
mermaid_syntax: str, mermaid_syntax: str,
*,
output_file_path: Optional[str] = None, output_file_path: Optional[str] = None,
background_color: Optional[str] = "white", background_color: Optional[str] = "white",
file_type: Optional[Literal["jpeg", "png", "webp"]] = "png", file_type: Optional[Literal["jpeg", "png", "webp"]] = "png",
max_retries: int = 1,
retry_delay: float = 1.0,
) -> bytes: ) -> bytes:
"""Renders Mermaid graph using the Mermaid.INK API.""" """Renders Mermaid graph using the Mermaid.INK API."""
try: try:
@ -400,15 +415,55 @@ def _render_mermaid_using_api(
f"https://mermaid.ink/img/{mermaid_syntax_encoded}" f"https://mermaid.ink/img/{mermaid_syntax_encoded}"
f"?type={file_type}&bgColor={background_color}" f"?type={file_type}&bgColor={background_color}"
) )
response = requests.get(image_url, timeout=10)
if response.status_code == requests.codes.ok:
img_bytes = response.content
if output_file_path is not None:
Path(output_file_path).write_bytes(response.content)
return img_bytes error_msg_suffix = (
msg = ( "To resolve this issue:\n"
f"Failed to render the graph using the Mermaid.INK API. " "1. Check your internet connection and try again\n"
f"Status code: {response.status_code}." "2. Try with higher retry settings: "
"`draw_mermaid_png(..., max_retries=5, retry_delay=2.0)`\n"
"3. Use the Pyppeteer rendering method which will render your graph locally "
"in a browser: `draw_mermaid_png(..., draw_method=MermaidDrawMethod.PYPPETEER)`"
) )
for attempt in range(max_retries + 1):
try:
response = requests.get(image_url, timeout=10)
if response.status_code == requests.codes.ok:
img_bytes = response.content
if output_file_path is not None:
Path(output_file_path).write_bytes(response.content)
return img_bytes
# If we get a server error (5xx), retry
if 500 <= response.status_code < 600 and attempt < max_retries:
# Exponential backoff with jitter
sleep_time = retry_delay * (2**attempt) * (0.5 + 0.5 * random.random()) # noqa: S311 not used for crypto
time.sleep(sleep_time)
continue
# For other status codes, fail immediately
msg = (
"Failed to reach https://mermaid.ink/ API while trying to render "
f"your graph. Status code: {response.status_code}.\n\n"
) + error_msg_suffix
raise ValueError(msg)
except (requests.RequestException, requests.Timeout) as e:
if attempt < max_retries:
# Exponential backoff with jitter
sleep_time = retry_delay * (2**attempt) * (0.5 + 0.5 * random.random()) # noqa: S311 not used for crypto
time.sleep(sleep_time)
else:
msg = (
"Failed to reach https://mermaid.ink/ API while trying to render "
f"your graph after {max_retries} retries. "
) + error_msg_suffix
raise ValueError(msg) from e
# This should not be reached, but just in case
msg = (
"Failed to reach https://mermaid.ink/ API while trying to render "
f"your graph after {max_retries} retries. "
) + error_msg_suffix
raise ValueError(msg) raise ValueError(msg)