mcp tool of the new derive graph (flow chart) from sympy code.

axiomatic_mcp/servers/equations/server.py

@@ -10,6 +10,7 @@
 from ...providers.toolset_provider import get_mcp_tools
 from ...shared.api_client import AxiomaticAPIClient
 from ...shared.documents.pdf_to_markdown import pdf_to_markdown
+from ...shared.utils.code_utils import get_python_code
 from ...shared.utils.prompt_utils import get_feedback_prompt
 
 
@@ -45,7 +46,7 @@ async def _get_document_content(document: Path | str) -> str:
     name="AxEquationExplorer Server",
     instructions="""This server provides tools to compose and analyze equations.
     """
-    + get_feedback_prompt("find_functional_form, check_equation"),
+    + get_feedback_prompt("find_functional_form, check_equation, generate_derivation_graph"),
     version="0.0.1",
     middleware=get_mcp_middleware(),
     tools=get_mcp_tools(),
@@ -61,7 +62,10 @@ async def _get_document_content(document: Path | str) -> str:
     tags=["equations", "compose", "derive", "find", "function-finder"],
 )
 async def find_expression(
-    document: Annotated[Path | str, "Either a file path to a PDF document or the document content as a string"],
+    document: Annotated[
+        Path | str,
+        "Either a file path to a PDF document or the document content as a string",
+    ],
     task: Annotated[str, "The task to be done for expression composition"],
 ) -> ToolResult:
     """If you have scientific text with equations, but you don't see the equation you're
@@ -110,8 +114,14 @@ async def find_expression(
     tags=["equations", "check", "error-correction", "validate"],
 )
 async def check_equation(
-    document: Annotated[Path | str, "Either a file path to a PDF document or the document content as a string"],
-    task: Annotated[str, "The task to be done for equation checking (e.g., 'check if E=mc² is correct')"],
+    document: Annotated[
+        Path | str,
+        "Either a file path to a PDF document or the document content as a string",
+    ],
+    task: Annotated[
+        str,
+        "The task to be done for equation checking (e.g., 'check if E=mc² is correct')",
+    ],
 ) -> ToolResult:
     """Use this tool to validate equations or check for errors in mathematical expressions.
     For example: 'Check if the equation F = ma is dimensionally consistent' or
@@ -131,3 +141,67 @@ async def check_equation(
 
     except Exception as e:
         raise ToolError(f"Failed to check equations in document: {e!s}") from e
+
+
+@mcp.tool(
+    name="generate_derivation_graph",
+    description=(
+        "Generates a Mermaid flowchart representing the mathematical derivation steps in SymPy code. "
+        "Analyzes the provided SymPy code and creates a visual flowchart showing the derivation flow, "
+        "intermediate calculations, and dependencies between steps. "
+        "This is particularly useful for: understanding complex mathematical derivations, "
+        "visualizing the flow from input variables to final results, identifying dependencies between "
+        "intermediate calculations, and creating documentation of derivation steps. "
+        "Example usage: 'Create a flowchart showing the derivation steps in this Euler-Lagrange code', "
+        "'Visualize how the Maxwell equations are derived in this SymPy script', or "
+        "'Generate a graph showing the dependencies in this quantum mechanics derivation'. "
+        "The tool returns a Mermaid flowchart text that can be rendered in Markdown or visualization tools."
+    ),
+    tags=["equations", "derivation", "graph", "flowchart", "visualization", "sympy"],
+)
+async def generate_derivation_graph(
+    sympy_code: Annotated[
+        Path | str,
+        "Either a file path to a Python file containing SymPy code or the SymPy code as a string",
+    ],
+) -> ToolResult:
+    try:
+        code_content, input_file_path = get_python_code(sympy_code)
+
+        # Send as multipart form-data with sympy_code as a text field
+        form_data = {"sympy_code": (None, code_content)}
+
+        response = AxiomaticAPIClient().post(
+            "/equations/compose/derivation-graph",
+            files=form_data,
+        )
+
+        mermaid_text = response.get("mermaid_text", "")
+        metadata = response.get("metadata")
+
+        if not mermaid_text:
+            raise ToolError("No mermaid_text returned from service")
+
+        # Save output file next to input file if it was a file, otherwise use current directory
+        output_path = input_file_path.parent / f"{input_file_path.stem}_derivation.mmd" if input_file_path else Path.cwd() / "derivation_graph.mmd"
+
+        with Path.open(output_path, "w", encoding="utf-8") as f:
+            f.write(mermaid_text)
+
+        if mermaid_text.strip().startswith("```"):
+            mermaid_display = f"\nMermaid Flowchart:\n{mermaid_text}"
+        else:
+            mermaid_display = f"\nMermaid Flowchart:\n```mermaid\n{mermaid_text}\n```"
+
+        result_content = [
+            TextContent(type="text", text=f"Mermaid flowchart saved to: {output_path}"),
+            TextContent(type="text", text=mermaid_display),
+        ]
+
+        if metadata:
+            result_content.append(TextContent(type="text", text=f"\nMetadata: {metadata}"))
+
+        return ToolResult(content=result_content)
+
+    except Exception as e:
+        raise ToolError(f"Failed to generate derivation graph: {e!s}") from e

axiomatic_mcp/shared/utils/code_utils.py

@@ -0,0 +1,44 @@
+"""Utilities for handling code input processing."""
+
+from pathlib import Path
+
+
+def get_python_code(code_input: Path | str) -> tuple[str, Path | None]:
+    """Extract Python code from either a file path or direct string.
+
+    Args:
+        code_input: Either a Path object to a .py file or a string (file path or direct code)
+
+    Returns:
+        tuple[str, Path | None]:
+            - code_content: The Python code to process
+            - file_path: Path to the source file if input was a file, None if direct code
+                        (Used to determine where to save output files)
+
+    Raises:
+        ValueError: If file not found or unsupported file type
+    """
+    if isinstance(code_input, Path):
+        if not code_input.exists():
+            raise ValueError(f"File not found: {code_input}")
+
+        if code_input.suffix.lower() != ".py":
+            raise ValueError(f"Unsupported file type: {code_input.suffix}. Only .py files are supported")
+
+        with Path.open(code_input, encoding="utf-8") as f:
+            return f.read(), code_input
+
+    if len(code_input) < 500 and "\n" not in code_input:
+        potential_path = Path(code_input)
+
+        if potential_path.suffix:
+            if not potential_path.exists():
+                raise ValueError(f"File not found: {potential_path}")
+
+            if potential_path.suffix.lower() != ".py":
+                raise ValueError(f"Unsupported file type: {potential_path.suffix}. Only .py files are supported")
+
+            with Path.open(potential_path, encoding="utf-8") as f:
+                return f.read(), potential_path
+
+    return code_input, None