"""Utility helpers for the visualize pipeline.""" from __future__ import annotations import json import re from typing import Any import defusedxml.ElementTree as ET _MERMAID_KEYWORDS = ( "graph", "flowchart", "sequenceDiagram", "classDiagram", "stateDiagram-v2", "stateDiagram", "erDiagram", "gantt", "mindmap", "pie", "journey", "gitGraph", "timeline", "quadrantChart", "requirementDiagram", "sankey-beta", "xychart-beta", "block-beta", "C4Context", ) def extract_json_object(text: str) -> dict[str, Any]: """Extract a JSON object from raw model output.""" raw = (text or "").strip() if not raw: return {} fenced = re.findall(r"```(?:json)?\s*([\s\S]*?)\s*```", raw) candidates = fenced + [raw] for candidate in candidates: try: parsed = json.loads(candidate) if isinstance(parsed, dict): return parsed except json.JSONDecodeError: parsed = _decode_first_json_object(candidate) if parsed is not None: return parsed start = raw.find("{") end = raw.rfind("}") if start != -1 and end != -1 and end > start: snippet = raw[start : end + 1] try: return json.loads(snippet) except json.JSONDecodeError: parsed = _decode_first_json_object(snippet) if parsed is not None: return parsed raise json.JSONDecodeError("No JSON object found", raw, 0) def _decode_first_json_object(text: str) -> dict[str, Any] | None: decoder = json.JSONDecoder() stripped = (text or "").lstrip() if not stripped: return None starts = [0] brace_index = stripped.find("{") if brace_index > 0: starts.append(brace_index) for start in starts: try: parsed, _end = decoder.raw_decode(stripped[start:]) except json.JSONDecodeError: continue if isinstance(parsed, dict): return parsed return None def extract_code_block(text: str, language: str = "") -> str: """Extract a fenced code block from LLM output. If *language* is given the block must start with that tag; otherwise any triple-backtick fence is accepted. """ if language: pattern = rf"```{re.escape(language)}\s*\n([\s\S]*?)\n```" else: pattern = r"```[A-Za-z]*\s*\n([\s\S]*?)\n```" match = re.search(pattern, text or "", re.IGNORECASE) if match: return match.group(1).strip() return (text or "").strip() def is_valid_html_document(html: str) -> bool: """Heuristic check that *html* looks like a renderable HTML fragment.""" if not html: return False lowered = html.lower() return " str: """Build a minimal, self-contained fallback HTML page. Used when the model fails to produce a renderable HTML document, so the user still gets *something* shown in the iframe instead of a blank panel. """ safe_title = (title or "Visualization").strip() or "Visualization" safe_summary = (summary or "").replace("\n", "
") or ( "The model did not return a renderable HTML document." ) safe_note = (note or "").replace("\n", "
") note_block = ( f'

Note:
{safe_note}

' if safe_note else "" ) return f""" {safe_title}

{safe_title}

{safe_summary}

{note_block}

""" def _strip_outer_fence(text: str) -> str: """Drop a single wrapping triple-backtick fence, if present.""" stripped = (text or "").strip() match = re.match(r"^```[A-Za-z]*\s*\n?([\s\S]*?)\n?```$", stripped) return match.group(1).strip() if match else stripped def validate_visualization(code: str, render_type: str) -> tuple[bool, str]: """Cheap, deterministic, local render-ability check. Returns ``(ok, error)``. When ``ok`` is False, ``error`` is a short, LLM-actionable message used to drive a single repair pass — none of these failures need an LLM call to *discover*. This replaces the generic LLM review for the text render types: only when local validation fails do we spend a model call (a targeted repair, not an open-ended review). """ text = (code or "").strip() if not text: return False, "Generated code is empty." if render_type == "svg": if " element." try: root = ET.fromstring(text) except ET.ParseError as exc: return False, f"SVG is not well-formed XML: {exc}" tag = root.tag.split("}")[-1].lower() if tag != "svg": return False, f"Root element must be