graph drawMermaid functionality

[justinlevi]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

This is in reference to the previous PR: https://github.com/langchain-ai/langchainjs/pull/5978

I have the following graph


  createGraph() {
    const workflow = new StateGraph<AgentState>({
      channels: state,
    })
      .addNode(
        'REWRITE_AGENT_NODE',
        this.rewriteAgentNode.run.bind(this.rewriteAgentNode),
      )
      .addNode(
        'CRUD_AGENT_NODE',
        this.crudAgentNode.run.bind(this.crudAgentNode),
      )
      .addNode('CRUD_TOOLS_NODE', this.crudToolsNode)
      .addNode('ERROR_HANDLER', this.handleError.bind(this))

      .addConditionalEdges(
        'REWRITE_AGENT_NODE',
        this.shouldContinue.bind(this),
        {
          CONTINUE: 'CRUD_AGENT_NODE',
          ERROR_HANDLER: 'ERROR_HANDLER',
          END,
        },
      )
      .addConditionalEdges('CRUD_AGENT_NODE', this.shouldContinue.bind(this), {
        CONTINUE: 'CRUD_TOOLS_NODE',
        END,
      })
      .addEdge('CRUD_TOOLS_NODE', 'CRUD_AGENT_NODE')
      .addEdge('ERROR_HANDLER', END)
      .addEdge(START, 'REWRITE_AGENT_NODE');

    return workflow.compile();
  }

### Description

when I call drawMermaid on the graph, I'm seeing the following output

"%%{init: {'flowchart': {'curve': 'linear'}}}%%\ngraph TD;\n\tLangGraphInput[LangGraphInput]:::startclass;\n\tLangGraph([LangGraph]):::otherclass;\n\tLangGraphOutput[LangGraphOutput]:::endclass;\n\tLangGraphInput --> LangGraph;\n\tLangGraph --> LangGraphOutput;\n\tclassDef startclass fill:#ffdfba;\n\tclassDef endclass fill:#baffc9;\n\tclassDef otherclass fill:#fad7de;\n"

which formats down to:

"""
graph TD;
LangGraphInput[LangGraphInput]:::startclass;
LangGraph([LangGraph]):::otherclass;
LangGraphOutput[LangGraphOutput]:::endclass;
LangGraphInput --> LangGraph;
LangGraph --> LangGraphOutput;
classDef startclass fill:#ffdfba;
classDef endclass fill:#baffc9;
classDef otherclass fill:#fad7de;
"""

---


### System Info

mac, all the latest versions


### Question
I'm confused as I was expecting to see all of the nodes and edges of my graph. Is that not correct?

[justinlevi]

[dosubot[bot]]

The drawMermaid function is not displaying all the nodes and edges of your graph as expected because the graph structure might not be correctly represented or passed to the function. Based on the provided tests, the expected graph structure should have specific nodes and edges.

Here is the expected structure for the graph in the "Test graph sequence" test:

{
  "nodes": [
    {
      "id": 0,
      "type": "schema",
      "data": {
        "title": "PromptTemplateInput",
        "$schema": "http://json-schema.org/draft-07/schema#"
      }
    },
    {
      "id": 1,
      "type": "runnable",
      "data": {
        "id": ["langchain_core", "prompts", "prompt", "PromptTemplate"],
        "name": "PromptTemplate"
      }
    },
    {
      "id": 2,
      "type": "runnable",
      "data": {
        "id": ["langchain", "llms", "fake", "FakeLLM"],
        "name": "FakeLLM"
      }
    },
    {
      "id": 3,
      "type": "runnable",
      "data": {
        "id": [
          "langchain_core",
          "output_parsers",
          "list",
          "CommaSeparatedListOutputParser"
        ],
        "name": "CommaSeparatedListOutputParser"
      }
    },
    {
      "id": 4,
      "type": "schema",
      "data": {
        "title": "CommaSeparatedListOutputParserOutput",
        "$schema": "http://json-schema.org/draft-07/schema#"
      }
    }
  ],
  "edges": [
    { "source": 0, "target": 1 },
    { "source": 1, "target": 2 },
    { "source": 3, "target": 4 },
    { "source": 2, "target": 3 }
  ]
}

Ensure that the drawMermaid function receives this structure correctly. If the structure is correct and the function still does not display all nodes and edges, the issue might be within the drawMermaid function itself. You may need to debug the function to ensure it processes and renders the graph structure as expected.

Additionally, the Graph class in langchain-core/src/runnables/graph.ts provides methods to manipulate the graph, such as addNode, removeNode, addEdge, trimFirstNode, and trimLastNode. Ensure that these methods are correctly used to maintain the integrity of the graph structure. For example, trimFirstNode and trimLastNode methods remove nodes with only one outgoing or incoming edge, which might affect the display if not handled properly ^[1][2].

To continue talking to Dosu, mention @dosu.

[justinlevi]

I ended up hacking my around this by redefining the the drawGraph from here: https://github.com/langchain-ai/langchainjs/blob/824e7e05f6f47656c83e43f2faec917399dd3c00/langchain-core/src/runnables/graph_mermaid.ts

function _escapeNodeLabel(nodeLabel: string): string {
  // Escapes the node label for Mermaid syntax.
  return nodeLabel.replace(/[^a-zA-Z-_0-9]/g, '_');
}

// Adjusts Mermaid edge to map conditional nodes to pure nodes.
function _adjustMermaidEdge(edge: Edge, nodes: Record<string, string>) {
  const sourceNodeLabel = nodes[edge.source] ?? edge.source;
  const targetNodeLabel = nodes[edge.target] ?? edge.target;
  return [sourceNodeLabel, targetNodeLabel];
}

function _generateMermaidGraphStyles(
  nodeColors: Record<string, string>,
): string {
  let styles = '';
  for (const [className, color] of Object.entries(nodeColors)) {
    styles += `\tclassDef ${className}class fill:${color};\n`;
  }
  return styles;
}

export interface Edge {
  source: string;
  target: string;
  data?: string;
  conditional?: boolean;
}

/**
 * Draws a Mermaid graph using the provided graph data
 */
export function drawMermaid(
  nodes: Record<string, string>,
  edges: Edge[],
  config?: {
    firstNodeLabel?: string;
    lastNodeLabel?: string;
    curveStyle?: string;
    withStyles?: boolean;
    nodeColors?: Record<string, string>;
    wrapLabelNWords?: number;
  },
): string {
  const {
    firstNodeLabel,
    lastNodeLabel,
    nodeColors,
    withStyles = true,
    curveStyle = 'linear',
  } = config ?? {};
  // Initialize Mermaid graph configuration
  let mermaidGraph = withStyles
    ? `%%{init: {'flowchart': {'curve': '${curveStyle}'}}}%%\ngraph TD;\n`
    : 'graph TD;\n';
  if (withStyles) {
    // Node formatting templates
    const defaultClassLabel = 'default';
    const formatDict: Record<string, string> = {
      [defaultClassLabel]: '{0}([{1}]):::otherclass',
    };
    if (firstNodeLabel !== undefined) {
      formatDict[firstNodeLabel] = '{0}[{0}]:::startclass';
    }
    if (lastNodeLabel !== undefined) {
      formatDict[lastNodeLabel] = '{0}[{0}]:::endclass';
    }

    // Add nodes to the graph
    for (const node of Object.values(nodes)) {
      const nodeLabel = formatDict[node] ?? formatDict[defaultClassLabel];
      const escapedNodeLabel = _escapeNodeLabel(node);
      const nodeParts = node.split(':');
      const nodeSplit = nodeParts[nodeParts.length - 1];
      mermaidGraph += `\t${nodeLabel
        .replace(/\{0\}/g, escapedNodeLabel)
        .replace(/\{1\}/g, nodeSplit)};\n`;
    }
  }
  let subgraph = '';
  // Add edges to the graph
  for (const edge of edges) {
    const sourcePrefix = edge.source.includes(':')
      ? edge.source.split(':')[0]
      : undefined;
    const targetPrefix = edge.target.includes(':')
      ? edge.target.split(':')[0]
      : undefined;

    // Exit subgraph if source or target is not in the same subgraph
    if (
      subgraph !== '' &&
      (subgraph !== sourcePrefix || subgraph !== targetPrefix)
    ) {
      mermaidGraph += '\tend\n';
      subgraph = '';
    }
    // Enter subgraph if source and target are in the same subgraph
    if (
      subgraph === '' &&
      sourcePrefix !== undefined &&
      sourcePrefix === targetPrefix
    ) {
      mermaidGraph += `\tsubgraph ${sourcePrefix}\n`;
      subgraph = sourcePrefix;
    }

    const [source, target] = _adjustMermaidEdge(edge, nodes);
    let edgeLabel = '';

    // Format edge label based on conditionality
    if (edge.conditional) {
      edgeLabel = ` -. ${edge.data} .-> `; // Conditional edge
    } else {
      edgeLabel = ` --> `; // Regular edge
    }

    mermaidGraph += `\t${_escapeNodeLabel(
      source,
    )}${edgeLabel}${_escapeNodeLabel(target)};\n`;
  }
  if (subgraph !== '') {
    mermaidGraph += 'end\n';
  }

  // Add custom styles for nodes
  if (withStyles && nodeColors !== undefined) {
    mermaidGraph += _generateMermaidGraphStyles(nodeColors);
  }
  return mermaidGraph;
}

/**
 * Renders Mermaid graph using the Mermaid.INK API.
 */
export async function drawMermaidPng(
  mermaidSyntax: string,
  config?: {
    backgroundColor?: string;
  },
) {
  let { backgroundColor = 'white' } = config ?? {};
  // Use btoa for compatibility, assume ASCII
  const mermaidSyntaxEncoded = btoa(mermaidSyntax);
  // Check if the background color is a hexadecimal color code using regex
  if (backgroundColor !== undefined) {
    const hexColorPattern = /^#(?:[0-9a-fA-F]{3}){1,2}$/;
    if (!hexColorPattern.test(backgroundColor)) {
      backgroundColor = `!${backgroundColor}`;
    }
  }
  const imageUrl = `https://mermaid.ink/img/${mermaidSyntaxEncoded}?bgColor=${backgroundColor}`;
  const res = await fetch(imageUrl);
  if (!res.ok) {
    throw new Error(
      [
        `Failed to render the graph using the Mermaid.INK API.`,
        `Status code: ${res.status}`,
        `Status text: ${res.statusText}`,
      ].join('\n'),
    );
  }
  const content = await res.blob();
  return content;
}

I then build up my own nodes and edges from the graph manually.

 async drawGraph() {
    const nodes = {};
    const edges = [];

    // Access nodes directly from the workflow's nodes property
    Object.entries(this.workflow.nodes).forEach(([key, node]) => {
      nodes[key] = key; // Adjust according to your node structure
    });

    console.log('---BRANCHES---', nodes);
    console.log(JSON.stringify(this.workflow.branches, null, 2));

    // Access edges using the allEdges method
    this.workflow.allEdges.forEach(([source, target]) => {
      edges.push({
        source,
        target,
        conditional: false, // Default to false
      });
    });

    // Handle branching by checking the workflow's branches
    for (const [key, branch] of Object.entries(this.workflow.branches)) {
      for (const [condition, targets] of Object.entries(branch)) {
        // targets is expected to be an object with ends
        if (targets.ends) {
          for (const [targetCondition, target] of Object.entries(
            targets.ends,
          )) {
            edges.push({
              source: key,
              target,
              conditional: true, // Mark as conditional for branching
              data: targetCondition, // Store the condition for edge labeling
            });
          }
        }
      }
    }

    console.log('---NODES---', nodes);
    console.log('---EDGES---', edges);

    // Generate and log the Mermaid graph
    const mermaidGraph = drawMermaid(nodes, edges, {
      firstNodeLabel: '__start__',
      lastNodeLabel: END,
      withStyles: true,
      nodeColors: {
        RETRIEVE: 'lightblue',
        GRADE_DOCUMENTS: 'lightgreen',
        GENERATE: 'lightcoral',
        WEBSEARCH: 'lightyellow',
      },
    });
    console.log(mermaidGraph);


  }

This outputs a nicely formatted Mermaid graph based on the actual graph I have defined

%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
  RETRIEVE([RETRIEVE]):::otherclass;
  GRADE_DOCUMENTS([GRADE_DOCUMENTS]):::otherclass;
  GENERATE([GENERATE]):::otherclass;
  WEBSEARCH([WEBSEARCH]):::otherclass;
  RETRIEVE --> GRADE_DOCUMENTS;
  WEBSEARCH --> GENERATE;
  __start__ -. WEBSEARCH .-> WEBSEARCH;
  __start__ -. RETRIEVE .-> RETRIEVE;
  GRADE_DOCUMENTS -. WEBSEARCH .-> WEBSEARCH;
  GRADE_DOCUMENTS -. GENERATE .-> GENERATE;
  GENERATE -. not supported .-> GENERATE;
  GENERATE -. useful .-> __end__;
  GENERATE -. not useful .-> WEBSEARCH;
  classDef RETRIEVEclass fill:lightblue;
  classDef GRADE_DOCUMENTSclass fill:lightgreen;
  classDef GENERATEclass fill:lightcoral;
  classDef WEBSEARCHclass fill:lightyellow;
  ```
  
This graph is originally from this repo: https://github.com/emarco177/langgaph-course/tree/main

[justinlevi]

[wi-ski]

@justinlevi - You're awesome.

[wi-ski]

is this something that should get PR-ed at least? like did it get fixed?