codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/agent/mermaid_fixer.py

@@ -0,0 +1,289 @@
+"""Mermaid syntax fixer — deterministic post-processing for LLM-generated diagrams.
+
+Ported from frontend preprocessMermaid() (StreamingMessage.tsx:125-276)
+with two additional categories (trailing semicolons, duplicate edges).
+Pure Python, no external dependencies.
+"""
+
+from __future__ import annotations
+
+import re
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def fix_mermaid_blocks(text: str) -> str:
+    """Find all ```mermaid blocks in *text* and fix each one.
+
+    Returns the original text unchanged if no mermaid blocks are present.
+    """
+    return _MERMAID_FENCE_RE.sub(_fix_block_match, text)
+
+
+def fix_mermaid_code(code: str) -> str:
+    """Fix a single mermaid code block (fence markers excluded)."""
+    code = _fix_comments(code)
+    code = _fix_keywords(code)
+    code = _fix_escaped_quotes(code)
+    code = _fix_arrows(code)
+    code = _fix_nested_brackets(code)
+    code = _fix_node_labels(code)
+    code = _fix_edge_labels(code)
+    code = _fix_subgraph_titles(code)
+    code = _fix_empty_labels(code)
+    code = _fix_direction(code)
+    code = _fix_trailing_semicolons(code)
+    code = _fix_duplicate_edges(code)
+    return code
+
+
+# ---------------------------------------------------------------------------
+# Regex for fenced mermaid blocks
+# ---------------------------------------------------------------------------
+
+_MERMAID_FENCE_RE = re.compile(r"```mermaid\s*\n([\s\S]*?)```")
+
+
+def _fix_block_match(m: re.Match[str]) -> str:
+    return "```mermaid\n" + fix_mermaid_code(m.group(1)) + "```"
+
+
+# ---------------------------------------------------------------------------
+# Category 1: Remove %% comments
+# ---------------------------------------------------------------------------
+
+
+def _fix_comments(code: str) -> str:
+    return re.sub(r"%%.*$", "", code, flags=re.MULTILINE)
+
+
+# ---------------------------------------------------------------------------
+# Category 2: Reserved keyword collision (`end` as node ID)
+# ---------------------------------------------------------------------------
+
+_END_WORD = re.compile(r"\bend\b")
+_END_SOLO_LINE = re.compile(r"^(\s*)end(\s*)$", re.MULTILINE)
+_DIRECTIVE_LINE = re.compile(
+    r"^(\s*(?:class|style|click|classDef|linkStyle)\b.*)$", re.MULTILINE
+)
+_PLACEHOLDER = "__END_PH__"
+
+
+def _fix_keywords(code: str) -> str:
+    # Protect: `end` alone on a line (subgraph closer)
+    code = _END_SOLO_LINE.sub(rf"\g<1>{_PLACEHOLDER}\g<2>", code)
+    # Protect: `end` inside class/style/click directives
+    code = _DIRECTIVE_LINE.sub(lambda m: _END_WORD.sub(_PLACEHOLDER, m.group(0)), code)
+    # Rename remaining bare `end` → `end_node`
+    code = _END_WORD.sub("end_node", code)
+    # Restore placeholders
+    code = code.replace(_PLACEHOLDER, "end")
+    return code
+
+
+# ---------------------------------------------------------------------------
+# Category 3: Escaped quotes
+# ---------------------------------------------------------------------------
+
+
+def _fix_escaped_quotes(code: str) -> str:
+    return code.replace('\\"', "&quot;")
+
+
+# ---------------------------------------------------------------------------
+# Category 4: Arrow syntax
+# ---------------------------------------------------------------------------
+
+_SINGLE_ARROW = re.compile(r"(\s)->([\s\w])")
+_DOTTED_ARROW = re.compile(r"(\s)-\.(\s)")
+
+
+def _fix_arrows(code: str) -> str:
+    code = _SINGLE_ARROW.sub(r"\1-->\2", code)
+    code = _DOTTED_ARROW.sub(r"\1-.->\2", code)
+    return code
+
+
+# ---------------------------------------------------------------------------
+# Category 5: Nested brackets & misquoted labels (depth counting)
+# ---------------------------------------------------------------------------
+
+
+def _fix_nested_brackets(code: str) -> str:
+    return re.sub(r"^(.*)$", _fix_line_brackets, code, flags=re.MULTILINE)
+
+
+def _fix_line_brackets(m: re.Match[str]) -> str:
+    line = m.group(0)
+    result: list[str] = []
+    i = 0
+    n = len(line)
+    while i < n:
+        if (
+            i > 0
+            and re.match(r"\w", line[i - 1])
+            and line[i] == "["
+            and (i + 1 < n)
+            and line[i + 1] not in ("[", "(", "/", "\\")
+        ):
+            depth = 1
+            j = i + 1
+            while j < n and depth > 0:
+                if line[j] == "[":
+                    depth += 1
+                elif line[j] == "]":
+                    depth -= 1
+                if depth > 0:
+                    j += 1
+
+            if depth == 0:
+                inner = line[i + 1 : j]
+                # Strip surrounding quotes
+                if inner.startswith('"') and inner.endswith('"'):
+                    inner = inner[1:-1]
+                elif inner.startswith('"'):
+                    inner = inner.replace('"', "")
+
+                if re.search(r'[\[\]"]', inner):
+                    escaped = (
+                        inner.replace("[", "&#91;")
+                        .replace("]", "&#93;")
+                        .replace('"', "&quot;")
+                    )
+                    result.append(f'["{escaped}"]')
+                else:
+                    result.append(line[i : j + 1])
+                i = j + 1
+            else:
+                result.append(line[i])
+                i += 1
+        else:
+            result.append(line[i])
+            i += 1
+    return "".join(result)
+
+
+# ---------------------------------------------------------------------------
+# Category 6: Node labels — quote and encode special chars
+# ---------------------------------------------------------------------------
+
+_SUBROUTINE_LABEL = re.compile(r"(\b\w+)\[\[([^\]\n]*)\]\]")
+_RECT_LABEL = re.compile(r"(\b\w+)\[(?!\[|\(|/|\\)([^\]\n]*)\]")
+
+
+def _sanitize_label(content: str) -> str:
+    return content.replace('"', "&quot;")
+
+
+def _fix_node_labels(code: str) -> str:
+    # [[content]] (subroutine) first
+    code = _SUBROUTINE_LABEL.sub(_quote_subroutine, code)
+    # [content] (rectangle)
+    code = _RECT_LABEL.sub(_quote_rect, code)
+    return code
+
+
+def _quote_subroutine(m: re.Match[str]) -> str:
+    node_id, content = m.group(1), m.group(2)
+    if content.startswith('"') and content.endswith('"'):
+        return m.group(0)
+    return f'{node_id}[["{_sanitize_label(content)}"]]'
+
+
+def _quote_rect(m: re.Match[str]) -> str:
+    node_id, content = m.group(1), m.group(2)
+    if content.startswith('"') and content.endswith('"'):
+        return m.group(0)
+    return f'{node_id}["{_sanitize_label(content)}"]'
+
+
+# ---------------------------------------------------------------------------
+# Category 7: Edge labels — brackets inside pipe-delimited labels
+# ---------------------------------------------------------------------------
+
+_EDGE_LABEL = re.compile(r"\|([^|]*)\|")
+
+
+def _fix_edge_labels(code: str) -> str:
+    return _EDGE_LABEL.sub(_escape_edge_label, code)
+
+
+def _escape_edge_label(m: re.Match[str]) -> str:
+    content = m.group(1)
+    if not re.search(r"[\[\]]", content):
+        return m.group(0)
+    return "|" + content.replace("[", "&#91;").replace("]", "&#93;") + "|"
+
+
+# ---------------------------------------------------------------------------
+# Category 8: Subgraph — quote unquoted titles with spaces
+# ---------------------------------------------------------------------------
+
+_SUBGRAPH_TITLE = re.compile(
+    r'^(\s*subgraph\s+)(?!")([^\n]*[ \t][^\n]*)$', re.MULTILINE
+)
+
+
+def _fix_subgraph_titles(code: str) -> str:
+    return _SUBGRAPH_TITLE.sub(_quote_subgraph_title, code)
+
+
+def _quote_subgraph_title(m: re.Match[str]) -> str:
+    prefix, title = m.group(1), m.group(2)
+    # Skip `subgraph id [title]` pattern
+    if re.match(r"^\w+\s*\[", title):
+        return m.group(0)
+    return f'{prefix}"{title.replace(chr(34), "&quot;")}"'
+
+
+# ---------------------------------------------------------------------------
+# Category 9: Empty labels — remove A[""] patterns
+# ---------------------------------------------------------------------------
+
+
+def _fix_empty_labels(code: str) -> str:
+    return re.sub(r'(\b\w+)\[""\]', r"\1", code)
+
+
+# ---------------------------------------------------------------------------
+# Category 10: Direction — fix invalid direction values
+# ---------------------------------------------------------------------------
+
+
+def _fix_direction(code: str) -> str:
+    return re.sub(
+        r"^(\s*direction\s+)(?!TB|BT|LR|RL|TD)(\w+)",
+        r"\g<1>TD",
+        code,
+        flags=re.MULTILINE,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Category 11 (NEW): Trailing semicolons
+# ---------------------------------------------------------------------------
+
+
+def _fix_trailing_semicolons(code: str) -> str:
+    return re.sub(r";\s*$", "", code, flags=re.MULTILINE)
+
+
+# ---------------------------------------------------------------------------
+# Category 12 (NEW): Duplicate edges
+# ---------------------------------------------------------------------------
+
+
+def _fix_duplicate_edges(code: str) -> str:
+    seen: set[str] = set()
+    lines: list[str] = []
+    for line in code.split("\n"):
+        stripped = line.strip()
+        # Only dedup edge lines (contain --> or --- or -.-> etc.)
+        if re.search(r"--[->]|---|\.-[>.]?->|==+>", stripped):
+            if stripped in seen:
+                continue
+            seen.add(stripped)
+        lines.append(line)
+    return "\n".join(lines)

jupyter_ext/agent/mermaid_middleware.py

@@ -0,0 +1,135 @@
+"""Mermaid Fix Middleware — deterministic Mermaid syntax repair on LLM output.
+
+Follows the same pattern as OutputGuardrailMiddleware (awrap_model_call)
+and AutoFixMiddleware (awrap_tool_call + request.override).
+"""
+
+from __future__ import annotations
+
+import datetime
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from langchain_core.messages import AIMessage
+
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    ModelCallResult,
+    ModelRequest,
+    ModelResponse,
+)
+from langgraph.prebuilt.tool_node import ToolCallRequest
+
+from .mermaid_fixer import fix_mermaid_blocks
+
+from ..log_paths import get_debug_log_path
+
+# ---------------------------------------------------------------------------
+# Debug logging (same pattern as guardrail/continuation middlewares)
+# ---------------------------------------------------------------------------
+
+_DEBUG_LOG = get_debug_log_path()
+
+_MERMAID_MARKER = "```mermaid"
+
+
+def _dbg(msg: str) -> None:
+    try:
+        ts = datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]
+        with open(_DEBUG_LOG, "a") as f:
+            f.write(f"[{ts}] MERMAID: {msg}\n")
+    except OSError:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# Helper: extract AIMessage from response (same as guardrail_middleware)
+# ---------------------------------------------------------------------------
+
+
+def _extract_ai_message(response: Any) -> AIMessage | None:
+    if isinstance(response, AIMessage):
+        return response
+    if hasattr(response, "message") and isinstance(response.message, AIMessage):
+        return response.message
+    if hasattr(response, "messages"):
+        for msg in reversed(response.messages):
+            if isinstance(msg, AIMessage):
+                return msg
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Middleware
+# ---------------------------------------------------------------------------
+
+
+class MermaidFixMiddleware(AgentMiddleware):
+    """Fix Mermaid syntax errors in LLM text responses and tool calls."""
+
+    # -- Hook 1: LLM text response + tool_calls with mermaid content ---------
+
+    def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> Awaitable[ModelCallResult]:
+        return self._impl_model(request, handler)
+
+    async def _impl_model(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        response = await handler(request)
+        ai_msg = _extract_ai_message(response)
+        if ai_msg is None:
+            return response
+
+        # Fix mermaid in text content
+        if isinstance(ai_msg.content, str) and _MERMAID_MARKER in ai_msg.content:
+            original = ai_msg.content
+            ai_msg.content = fix_mermaid_blocks(original)
+            if ai_msg.content != original:
+                _dbg("fixed mermaid in text content")
+
+        # Fix mermaid in jupyter_markdown tool_calls
+        tool_calls = getattr(ai_msg, "tool_calls", None)
+        if tool_calls:
+            for tc in tool_calls:
+                if tc.get("name") != "jupyter_markdown":
+                    continue
+                args = tc.get("args", {})
+                content = args.get("content", "")
+                if isinstance(content, str) and _MERMAID_MARKER in content:
+                    args["content"] = fix_mermaid_blocks(content)
+                    if args["content"] != content:
+                        _dbg("fixed mermaid in jupyter_markdown tool_call")
+
+        return response
+
+    # -- Hook 2: jupyter_markdown tool execution -----------------------------
+
+    async def awrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], Awaitable[Any]],
+    ) -> Any:
+        tool_name = request.tool_call.get("name", "")
+        if tool_name != "jupyter_markdown":
+            return await handler(request)
+
+        args = request.tool_call.get("args", {})
+        content = args.get("content", "")
+
+        if isinstance(content, str) and _MERMAID_MARKER in content:
+            fixed = fix_mermaid_blocks(content)
+            if fixed != content:
+                modified_call = {
+                    **request.tool_call,
+                    "args": {**args, "content": fixed},
+                }
+                request = request.override(tool_call=modified_call)
+                _dbg("fixed mermaid in jupyter_markdown tool execution")
+
+        return await handler(request)