diff --git a/tools/mermaid/banner.png b/tools/mermaid/banner.png new file mode 100644 index 00000000..9853232b Binary files /dev/null and b/tools/mermaid/banner.png differ diff --git a/tools/mermaid/icon.png b/tools/mermaid/icon.png new file mode 100644 index 00000000..00448600 Binary files /dev/null and b/tools/mermaid/icon.png differ diff --git a/tools/mermaid/index.test.ts b/tools/mermaid/index.test.ts new file mode 100644 index 00000000..25943f32 --- /dev/null +++ b/tools/mermaid/index.test.ts @@ -0,0 +1,47 @@ +// tools/mermaid-generator/index.test.ts + +import { expect } from '@jest/globals'; +import { getToolTestClient } from '../../src/test/utils'; +import * as path from 'path'; + +describe('Mermaid Diagram Tool', () => { + // Adjust the folder name and file name if necessary to match + // where you put the mermaid generator code. + const toolPath = path.join(__dirname, 'tool.ts'); + const client = getToolTestClient(); + + it('generates a mermaid diagram from a simple description', async () => { + const testDescription = 'Create a flow that goes from A to B to C.'; + + // Execute tool + const response = await client.executeToolFromFile(toolPath, { description: testDescription }, {}, ['local:::__official_shinkai:::shinkai_llm_prompt_processor']); + // Expect the result to contain our output fields + expect(response).toHaveProperty('pngBase64'); + expect(response).toHaveProperty('finalMermaid'); + + // Validate some minimal checks + expect(typeof response.pngBase64).toBe('string'); + expect(response.pngBase64.length).toBeGreaterThan(0); + + // We can also do a minimal check of whether finalMermaid contains "graph" or something relevant + expect(response.finalMermaid).toMatch(/graph/i); + }, 20000); + + it('handles invalid descriptions gracefully (or tries multiple times)', async () => { + // Potentially we might feed in something bizarre to see if it times out or returns an error + // or a fallback. Adjust as needed. + const invalidDescription = '!!!@#$%^&*'; + + try { + const response = await client.executeToolFromFile(toolPath, { + description: invalidDescription + }, {}, ['local:::__official_shinkai:::shinkai_llm_prompt_processor']); + // We expect it might still succeed or might fail. Adjust as needed for your logic. + expect(response).toHaveProperty('pngBase64'); + expect(response).toHaveProperty('finalMermaid'); + } catch (err: any) { + // If your code throws an error after the maximum attempts, we check that logic here. + expect(err.message).toMatch(/Failed to produce a valid Mermaid diagram/i); + } + }, 30000); +}); diff --git a/tools/mermaid/metadata.json b/tools/mermaid/metadata.json new file mode 100644 index 00000000..ed88ca5f --- /dev/null +++ b/tools/mermaid/metadata.json @@ -0,0 +1,49 @@ +{ + "id": "mermaid", + "name": "mermaid", + "version": "1.0.0", + "description": "Generate diagrams and flowcharts using Mermaid syntax", + "author": "Shinkai", + "keywords": [ + "mermaid", + "diagram", + "flowchart", + "visualization", + "markdown" + ], + "configurations": { + "type": "object", + "properties": {}, + "required": [] + }, + "parameters": { + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "Mermaid diagram code" + }, + "format": { + "type": "string", + "enum": ["svg", "png"], + "default": "svg", + "description": "Output format for the diagram" + } + }, + "required": ["code"] + }, + "result": { + "type": "object", + "properties": { + "image": { + "type": "string", + "description": "Base64 encoded image data" + }, + "format": { + "type": "string", + "description": "Format of the generated image (svg or png)" + } + }, + "required": ["image", "format"] + } +} diff --git a/tools/mermaid/store.json b/tools/mermaid/store.json new file mode 100644 index 00000000..cd631dd4 --- /dev/null +++ b/tools/mermaid/store.json @@ -0,0 +1,3 @@ +{ + "categoryId": "0569c7f5-2942-407e-9aae-4e979718683b" +} diff --git a/tools/mermaid/tool.ts b/tools/mermaid/tool.ts new file mode 100644 index 00000000..4044e496 --- /dev/null +++ b/tools/mermaid/tool.ts @@ -0,0 +1,270 @@ +import axios from 'npm:axios@1.7.7'; +import { shinkaiLlmPromptProcessor } from './shinkai-local-tools.ts'; +import { encodeBase64 } from "https://deno.land/std@0.224.0/encoding/base64.ts"; +import { deflate } from "https://deno.land/x/compress@v0.4.5/zlib/deflate.ts"; +import { getHomePath } from './shinkai-local-support.ts'; +/** + * Configuration for the tool. + */ +type CONFIG = { + /** + * How many times to attempt LLM fixes if Kroki fails to parse the Mermaid diagram. + */ + maxRetries?: number; +}; + +/** + * Inputs for the tool: a single textual description from the user. + */ +type INPUTS = { + description: string; +}; + +/** + * Final output from the tool: + * - The base64-encoded PNG + * - The final (valid) Mermaid code that was successfully parsed. + */ +type OUTPUT = { + pngBase64: string; + finalMermaid: string; +}; + +/** + * This function: + * 1. Takes a textual description and asks an LLM to produce Mermaid code. + * 2. Sends the Mermaid code to Kroki (https://kroki.io/) to validate and render a PNG. + * 3. If Kroki fails to parse, it sends the error back to the LLM to refine the Mermaid code. + * 4. Repeats up to `maxRetries` times. If still invalid, throws an error. + */ +export async function run(config: CONFIG, inputs: INPUTS): Promise { + const { description } = inputs; + const maxRetries = config.maxRetries ?? 5; + + /** + * Helper: build the JSON payload Kroki expects to attempt rendering a Mermaid PNG. + */ + function buildKrokiPayload(mermaidSource: string) { + return { + diagram_source: mermaidSource, + diagram_type: 'mermaid', + output_format: 'png', + }; + } + + /** + * Attempt to render with Kroki. On success: return { ok: true, data: Buffer }. + * On failure: return { ok: false, error: string }. + */ + async function tryKrokiRender(mermaidCode: string) { + console.log('Attempting to render with Kroki:', { mermaidCode }); + + // Basic validation before sending to Kroki + if (!mermaidCode.trim().startsWith('graph')) { + console.log('Basic validation failed: Code does not start with "graph"'); + return { ok: false, error: 'Invalid Mermaid syntax: Must start with "graph"' }; + } + + try { + // First deflate the diagram + const encoder = new TextEncoder(); + const compressed = deflate(encoder.encode(mermaidCode.trim()), { level: 9 }); + // Then base64 encode it + const encodedDiagram = encodeBase64(compressed).replace(/\+/g, '-').replace(/\//g, '_'); + console.log('Encoded diagram:', { encodedDiagram }); + + console.log('Sending request to Kroki...'); + const resp = await axios.get(`https://kroki.io/mermaid/png/${encodedDiagram}`, { + responseType: 'arraybuffer', + timeout: 30000, + headers: { + 'Accept': 'image/png', + }, + validateStatus: (status) => status === 200, + }); + + console.log('Received successful response from Kroki'); + return { ok: true, data: new Uint8Array(resp.data) }; + } catch (err: any) { + console.log('Error from Kroki:', { + status: err.response?.status, + headers: err.response?.headers, + isAxiosError: err.isAxiosError, + message: err.message, + data: err.response?.data?.toString() + }); + + // Handle various error cases + if (err.response) { + const errorData = err.response.data; + let errorMessage = ''; + + try { + // Try to parse error as JSON if it's not binary data + if (err.response.headers['content-type']?.includes('application/json')) { + const jsonError = JSON.parse(errorData.toString()); + errorMessage = jsonError.error || jsonError.message || String(errorData); + } else { + errorMessage = errorData.toString(); + } + } catch (parseErr) { + console.log('Error parsing error response:', parseErr); + errorMessage = errorData.toString(); + } + + console.log('Formatted error message:', errorMessage); + return { + ok: false, + error: `Kroki error (HTTP ${err.response.status}): ${errorMessage}`, + }; + } + + // Network or other errors + return { ok: false, error: `Request failed: ${err.message}` }; + } + } + + /** + * Validate Mermaid syntax before sending to Kroki + */ + function validateMermaidSyntax(code: string): { isValid: boolean; error?: string } { + console.log('Validating Mermaid syntax for:', { code }); + + const trimmed = code.trim(); + + // Basic syntax checks + if (!trimmed.toLowerCase().startsWith('graph')) { + console.log('Validation failed: Does not start with "graph"'); + return { isValid: false, error: 'Diagram must start with "graph"' }; + } + + const lines = trimmed.split('\n').map(line => line.trim()).filter(line => line); + console.log('Processing lines:', { lines }); + + const firstLine = lines[0]; + + // Check graph direction + if (!firstLine.toLowerCase().match(/^graph\s+(td|lr)$/)) { + console.log('Validation failed: Invalid graph direction:', { firstLine }); + return { isValid: false, error: 'First line must be "graph TD" or "graph LR"' }; + } + + // Check for basic node definitions + const nodeLines = lines.slice(1); + for (const line of nodeLines) { + console.log('Checking node line:', { line }); + // More lenient regex that allows various amounts of whitespace + if (!line.match(/^[A-Za-z0-9]+(?:\[[^\]]+\])?\s*(?:-->|---|==>)\s*[A-Za-z0-9]+(?:\[[^\]]+\])?$/)) { + console.log('Validation failed: Invalid node definition:', { line }); + return { isValid: false, error: `Invalid node definition: ${line}` }; + } + } + + console.log('Validation successful'); + return { isValid: true }; + } + + /** + * LLM prompt to request a new or revised Mermaid code from the LLM. + */ + async function requestMermaid( + userDescription: string, + priorError?: string, + priorCode?: string + ): Promise { + let prompt = ''; + if (!priorError) { + // initial request + prompt = `Create a valid Mermaid.js diagram based on this description: "${userDescription}" + +Rules: +1. Start with either 'graph TD' (top-down) or 'graph LR' (left-right) +2. Use simple node names (A, B, C, etc.) with descriptive labels in brackets +3. Use standard arrows (-->) +4. Avoid special characters in labels +5. Return ONLY the Mermaid code, no explanations + +Example of valid format: +graph TD + A[Start] --> B[Process] + B --> C[End]`; + } else { + // revise with specific guidance based on prior error + prompt = `The following Mermaid code needs correction: +\`\`\` +${priorCode} +\`\`\` + +Error received: ${priorError} + +Please provide a corrected version following these rules: +1. Keep the diagram simple and minimal +2. Use only basic Mermaid syntax (graph TD/LR, basic nodes, arrows) +3. Ensure all nodes are properly defined before being referenced +4. Avoid special characters or complex styling +5. Return ONLY the corrected Mermaid code + +Example of valid format: +graph TD + A[Start] --> B[Process] + B --> C[End]`; + } + const resp = await shinkaiLlmPromptProcessor({ format: 'text', prompt }); + + // Clean up the response to extract just the Mermaid code + let code = resp.message.trim(); + // Remove any markdown code block markers + code = code.replace(/^```mermaid\n/m, '').replace(/^```\n/m, '').replace(/```$/m, ''); + return code.trim(); + } + + // Main logic: + console.log('Starting Mermaid diagram generation for description:', { description }); + let currentMermaid = await requestMermaid(description, undefined, undefined); + console.log('Initial Mermaid code generated:', { currentMermaid }); + + for (let attempt = 0; attempt < maxRetries; attempt++) { + console.log(`Attempt ${attempt + 1}/${maxRetries}`); + + // Validate syntax before sending to Kroki + const validation = validateMermaidSyntax(currentMermaid); + if (!validation.isValid) { + console.log('Validation failed:', validation.error); + // If invalid syntax, try to get a new diagram + currentMermaid = await requestMermaid( + description, + `Invalid Mermaid syntax: ${validation.error}`, + currentMermaid + ); + console.log('Generated new Mermaid code after validation failure:', { currentMermaid }); + continue; + } + + console.log('Validation passed, attempting to render'); + const renderResult = await tryKrokiRender(currentMermaid); + if (renderResult.ok && renderResult.data) { + console.log('Successfully rendered diagram'); + // Convert Uint8Array to base64 string + const pngBase64 = encodeBase64(renderResult.data); + + await Deno.writeFile(await getHomePath() + '/mermaid.png', renderResult.data); + + return { + pngBase64, + finalMermaid: currentMermaid, + }; + } else { + console.log('Render failed:', renderResult.error); + // Some error from Kroki. Let's refine + const errorMessage = renderResult.error || 'Unknown error'; + currentMermaid = await requestMermaid(description, errorMessage, currentMermaid); + console.log('Generated new Mermaid code after render failure:', { currentMermaid }); + } + } + + console.log('Exhausted all attempts, throwing error'); + // If we've exhausted attempts, throw an error + throw new Error( + `Failed to produce a valid Mermaid diagram after ${maxRetries} attempts. Last code:\n${currentMermaid}` + ); +}