updated code with working mcp

app.py

@@ -12,6 +12,7 @@ from py2puml.py2puml import py2puml
 from plantuml import PlantUML
 import pyan
 from pathlib import Path
+from utils import setup_testing_space, verify_testing_space, cleanup_testing_space
 
 if os.name == "nt":  # nt == Windows
     graphviz_bin = r"C:\\Program Files\\Graphviz\\bin"
@@ -19,55 +20,6 @@ if os.name == "nt":  # nt == Windows
         os.environ["PATH"] += os.pathsep + graphviz_bin
 
 
-def setup_testing_space():
-    """Create persistent testing_space directory and __init__.py at startup."""
-    testing_dir = os.path.join(os.getcwd(), "testing_space")
-    os.makedirs(testing_dir, exist_ok=True)
-
-    init_file = os.path.join(testing_dir, "__init__.py")
-    if not os.path.exists(init_file):
-        with open(init_file, "w", encoding="utf-8") as f:
-            f.write("# Testing space for py2puml analysis\n")
-        print("📁 Created testing_space directory and __init__.py")
-    else:
-        print("🔄 testing_space directory already exists")
-
-
-def cleanup_testing_space():
-    """Remove all .py files except __init__.py from testing_space."""
-    testing_dir = os.path.join(os.getcwd(), "testing_space")
-    if not os.path.exists(testing_dir):
-        print("⚠️ testing_space directory not found, creating it...")
-        setup_testing_space()
-        return
-
-    # Clean up any leftover .py files (keep only __init__.py)
-    files_removed = 0
-    for file in os.listdir(testing_dir):
-        if file.endswith(".py") and file != "__init__.py":
-            file_path = os.path.join(testing_dir, file)
-            try:
-                os.remove(file_path)
-                files_removed += 1
-            except Exception as e:
-                print(f"⚠️ Could not remove {file}: {e}")
-
-    if files_removed > 0:
-        print(f"🧹 Cleaned up {files_removed} leftover .py files from testing_space")
-
-
-def verify_testing_space():
-    """Verify testing_space contains only __init__.py."""
-    testing_dir = os.path.join(os.getcwd(), "testing_space")
-    if not os.path.exists(testing_dir):
-        return False
-
-    files = os.listdir(testing_dir)
-    expected_files = ["__init__.py"]
-
-    return files == expected_files
-
-
 def generate_call_graph_with_pyan3(
     python_code: str, filename: str = "analysis"
 ) -> Tuple[Optional[str], Optional[str], Dict[str, Any]]:
@@ -88,7 +40,7 @@ def generate_call_graph_with_pyan3(
     unique_filename = f"{filename}_{timestamp}"
 
     # Paths
-    testing_dir = os.path.join(os.getcwd(), "testing_space")
+    testing_dir = os.path.join(os.getcwd(), "inputs")
     code_file = os.path.join(testing_dir, f"{unique_filename}.py")
 
     try:
@@ -383,7 +335,7 @@ def generate_diagram(python_code: str, filename: str = "diagram") -> Optional[st
     unique_filename = f"{filename}_{timestamp}"
 
     # Paths
-    testing_dir = os.path.join(os.getcwd(), "testing_space")
+    testing_dir = os.path.join(os.getcwd(), "inputs")
     code_file = os.path.join(testing_dir, f"{unique_filename}.py")
 
     # Use PlantUML web service for rendering
@@ -394,7 +346,7 @@ def generate_diagram(python_code: str, filename: str = "diagram") -> Optional[st
         with open(code_file, "w", encoding="utf-8") as f:
             f.write(python_code)
 
-        print(f"📝 Created temporary file: testing_space/{unique_filename}.py")
+        print(f"📝 Created temporary file: inputs/{unique_filename}.py")
 
         # Generate PlantUML content using py2puml (no sys.path manipulation needed)
         print(f"📝 Generating PlantUML content...")
@@ -402,7 +354,7 @@ def generate_diagram(python_code: str, filename: str = "diagram") -> Optional[st
             os.path.join(
                 testing_dir, unique_filename
             ),  # path to the .py file (without extension)
-            f"testing_space.{unique_filename}",  # module name
+            f"inputs.{unique_filename}",  # module name
         )
         puml_content = "".join(puml_content_lines)
 
@@ -451,7 +403,7 @@ def generate_diagram(python_code: str, filename: str = "diagram") -> Optional[st
 
 
 def analyze_code_structure(python_code: str) -> str:
-    """Enhanced code analysis combining AST + pyan3 call graphs.
+    """Return a Markdown report with complexity metrics and recommendations.
 
     Args:
         python_code: The Python code to analyze
@@ -736,277 +688,450 @@ def get_sample_code(filename: str) -> str:
         return f.read()
 
 
-def generate_all_diagrams(python_code: str, filename: str = "diagram") -> Tuple[Optional[str], Optional[str], str]:
-    """Generate all diagrams and analysis at once.
-    
+def generate_all_diagrams(
+    python_code: str, filename: str = "diagram"
+) -> Tuple[Optional[str], Optional[str], str]:
+    """Generate class diagram, call-graph diagram and analysis in one call.
+
     Args:
         python_code: The Python code to analyze
         filename: Base filename for diagrams
-        
+
     Returns:
         Tuple of (uml_diagram_path, call_graph_path, analysis_text)
     """
     if not python_code.strip():
         return None, None, "No code provided for analysis."
-    
+
     print("🚀 Starting comprehensive diagram generation...")
-    
+
     # Step 1: Generate UML Class Diagram
     print("📊 Step 1/3: Generating UML class diagram...")
     uml_diagram_path = generate_diagram(python_code, filename)
-    
+
     # Step 2: Generate Call Graph
     print("🔗 Step 2/3: Generating call graph...")
     try:
         cleanup_testing_space()
-        dot_content, call_graph_path, structured_data = generate_call_graph_with_pyan3(python_code)
+        dot_content, call_graph_path, structured_data = generate_call_graph_with_pyan3(
+            python_code
+        )
     except Exception as e:
         print(f"⚠️ Call graph generation failed: {e}")
         call_graph_path = None
-    
+
     # Step 3: Generate Analysis
     print("📈 Step 3/3: Performing code analysis...")
     analysis_text = analyze_code_structure(python_code)
-    
+
     print("✅ All diagrams and analysis completed!")
-    
+
     return uml_diagram_path, call_graph_path, analysis_text
 
 
-# Create Gradio interface
+# def generate_class_diagram(python_code: str) -> str:
+#     """Create a UML class diagram from Python code and return the PNG path.
+
+#     Args:
+#         python_code (str): Python source to analyse.
+
+#     Returns:
+#         str: Local path to a PNG image with the UML diagram.
+#     """
+#     result = generate_diagram(python_code)
+#     return result if result is not None else ""
+
+
+# def generate_call_graph_diagram(python_code: str) -> str:
+#     """Create a function-call graph diagram (PNG).
+
+#     Args:
+#         python_code (str): Python source to analyse.
+
+#     Returns:
+#         str: Local path to the generated call-graph PNG.
+#     """
+#     _, png_path, _ = generate_call_graph_with_pyan3(python_code)
+#     return png_path or ""
+
+
+
+
+# # -----------------------------------------------------------------------------
+# # ❶  Four MCP‑exposed Interfaces ------------------------------------------------
+# # -----------------------------------------------------------------------------
+# iface_class = gr.Interface(
+#     fn=generate_class_diagram,
+#     inputs=gr.Textbox(lines=20, label="Python code"),
+#     outputs=gr.Image(label="UML diagram"),
+#     api_name="generate_class_diagram",
+#     description="Create a UML class diagram (PNG) from Python code.",
+# )
+
+# iface_call = gr.Interface(
+#     fn=generate_call_graph_diagram,
+#     inputs=gr.Textbox(lines=20, label="Python code"),
+#     outputs=gr.Image(label="Call‑graph"),
+#     api_name="generate_call_graph_diagram",
+#     description="Generate a function‑call graph (PNG) from Python code.",
+# )
+
+# iface_analysis = gr.Interface(
+#     fn=analyze_code_structure,
+#     inputs=gr.Textbox(lines=20, label="Python code"),
+#     outputs=gr.Textbox(lines=20, label="Analysis"),
+#     api_name="analyze_code_structure",
+#     description="Return a Markdown report with complexity metrics.",
+# )
+
+# iface_all = gr.Interface(
+#     fn=generate_all_diagrams,
+#     inputs=[
+#         gr.Textbox(lines=20, label="Python code"),
+#         gr.Textbox(value="diagram", label="Base filename"),
+#     ],
+#     outputs=[
+#         gr.Image(label="UML diagram"),
+#         gr.Image(label="Call‑graph"),
+#         gr.Textbox(lines=20, label="Analysis"),
+#     ],
+#     api_name="generate_all",
+#     description="Run class diagram, call graph and analysis in one call.",
+# )
+
+# # -----------------------------------------------------------------------------
+# # ❷  Web UI (Blocks) -----------------------------------------------------------
+# # -----------------------------------------------------------------------------
+# with gr.Blocks(
+#     title="Python Diagram Generator MCP Demo",
+#     theme=gr.themes.Soft(),
+#     css="""
+#     .gradio-container { max-width: 1400px !important; }
+#     .code-input       { font-family: 'Inter', monospace !important; }
+#     """,
+# ) as demo:
+
+#     # -- Header ----------------------------------------------------------------
+#     gr.Markdown(
+#         """
+#         # 🐍 Python UML Diagram Generator & MCP Server
+
+#         **Dual Functionality**
+#         1. 🖥️ *Web Interface* – generate UML & call‑graph diagrams + analysis  
+#         2. 🤖 *MCP Server*  – four tools ready for assistants such as Claude Desktop or Cursor
+#         """
+#     )
+
+#     # -- Main TAB --------------------------------------------------------------
+#     with gr.Tab("🎨 Diagram Generator"):
+#         with gr.Row():
+#             # ---------- Left column – inputs -----------------------------------
+#             with gr.Column(scale=1):
+#                 gr.Markdown("### Input")
+
+#                 example_files = list_example_files()
+#                 example_dropdown = gr.Dropdown(
+#                     label="Choose Example",
+#                     choices=example_files,
+#                     value=example_files[0] if example_files else None,
+#                 )
+
+#                 code_input = gr.Textbox(
+#                     label="Python Code",
+#                     placeholder="Paste your Python code here…",
+#                     lines=20,
+#                     max_lines=35,
+#                     value=get_sample_code(example_files[0]) if example_files else "",
+#                     elem_classes=["code-input"],
+#                 )
+
+#                 filename_input = gr.Textbox(
+#                     label="Diagram Name",
+#                     value="my_diagram",
+#                     placeholder="Enter a name for your diagram",
+#                 )
+
+#                 generate_diagrams_btn = gr.Button(
+#                     "🔄 Generate Diagrams", variant="primary", size="lg"
+#                 )
+
+#             # ---------- Right column – outputs ---------------------------------
+#             with gr.Column(scale=1):
+#                 gr.Markdown("### Generated UML Class Diagram")
+#                 uml_diagram_output = gr.Image(
+#                     label="UML Class Diagram", show_download_button=True, height=300
+#                 )
+
+#                 gr.Markdown("### Generated Call Graph Diagram")
+#                 call_graph_output = gr.Image(
+#                     label="Function Call Graph", show_download_button=True, height=300
+#                 )
+
+#         # ---------- Analysis row ----------------------------------------------
+#         gr.Markdown("### Code Analysis")
+#         analysis_output = gr.Textbox(
+#             label="Comprehensive Code Analysis",
+#             lines=15,
+#             max_lines=25,
+#             interactive=False,
+#             show_copy_button=True,
+#         )
+
+#     # -- About TAB -------------------------------------------------------------
+#     # with gr.Tab("ℹ️ About & Help"):
+#     #     gr.Markdown(open("ABOUT.md", "r", encoding="utf-8").read())  # keep long text external
+
+#     # -- Hidden TabbedInterface to register MCP tools -------------------------
+#     # (render=False avoids GUI duplication)
+#     # gr.TabbedInterface(
+#     #     [iface_class, iface_call, iface_analysis, iface_all],
+#     #     ["Class diagram", "Call graph", "Analysis", "All‑in‑one"],
+#     #     # render=False,
+#     # )
+
+#     # -------------------------------------------------------------------------
+#     # Event handlers -----------------------------------------------------------
+#     # -------------------------------------------------------------------------
+
+#     def _load_example(example_filename: str):
+#         """Load bundled example code into the textbox."""
+#         return get_sample_code(example_filename)
+
+#     example_dropdown.change(_load_example, example_dropdown, code_input)
+
+#     generate_diagrams_btn.click(
+#         fn=generate_all_diagrams,
+#         inputs=[code_input, filename_input],
+#         outputs=[uml_diagram_output, call_graph_output, analysis_output],
+#         show_progress=True,
+#     )
+
+#     # code_input.change(
+#     #     fn=analyze_code_structure,
+#     #     inputs=code_input,
+#     #     outputs=analysis_output,
+#     #     show_progress=False,
+#     # )
+
+# # -----------------------------------------------------------------------------
+# # ❸  Launch -------------------------------------------------------------------
+# # -----------------------------------------------------------------------------
+# if __name__ == "__main__":
+#     setup_testing_space()  # create a persistent working dir if needed
+
+#     demo.launch(
+#         mcp_server=True,   # Enable MCP endpoints (/gradio_api/mcp/*)
+#         show_api=False,    # Expose ONLY the 4 Interfaces as tools
+#         show_error=True,   # Display exceptions in the UI
+#         debug=True,        # Verbose server logs
+#     )
+
+
+
+# =============================================================================
+# ❶  Wrapper functions for diagram and analysis generation
+# These will be connected to the UI buttons and the MCP interfaces.
+# =============================================================================
+
+def generate_class_diagram_only(python_code: str) -> Optional[str]:
+    """Generates just the UML class diagram."""
+    if not python_code.strip():
+        gr.Warning("Input code is empty!")
+        return None
+    return generate_diagram(python_code)
+
+def generate_call_graph_only(python_code: str) -> Optional[str]:
+    """Generates just the call graph diagram."""
+    if not python_code.strip():
+        gr.Warning("Input code is empty!")
+        return None
+    _, png_path, _ = generate_call_graph_with_pyan3(python_code)
+    return png_path
+
+def analyze_code_only(python_code: str) -> str:
+    """Generates just the code analysis report."""
+    if not python_code.strip():
+        gr.Warning("Input code is empty!")
+        return "No code provided to analyze."
+    return analyze_code_structure(python_code)
+    
+def generate_all_outputs(python_code: str) -> Tuple[Optional[str], Optional[str], str]:
+    """Generates all three outputs: UML diagram, call graph, and analysis."""
+    if not python_code.strip():
+        gr.Warning("Input code is empty!")
+        return None, None, "No code provided to analyze."
+    
+    print("🚀 Starting comprehensive generation...")
+    uml_path = generate_diagram(python_code)
+    _, call_graph_path, _ = generate_call_graph_with_pyan3(python_code)
+    analysis_text = analyze_code_structure(python_code)
+    print("✅ All outputs generated!")
+    
+    return uml_path, call_graph_path, analysis_text
+
+# =============================================================================
+# ❷  Four MCP-exposed Interfaces
+# These are NOT rendered in the UI but are exposed as tools for agents.
+# =============================================================================
+
+iface_class = gr.Interface(
+    fn=generate_class_diagram_only,
+    inputs=gr.Textbox(lines=20, label="Python code"),
+    outputs=gr.Image(label="UML diagram"),
+    api_name="generate_class_diagram",
+    description="Create a UML class diagram (PNG) from Python code.",
+)
+
+iface_call = gr.Interface(
+    fn=generate_call_graph_only,
+    inputs=gr.Textbox(lines=20, label="Python code"),
+    outputs=gr.Image(label="Call‑graph"),
+    api_name="generate_call_graph_diagram",
+    description="Generate a function‑call graph (PNG) from Python code.",
+)
+
+iface_analysis = gr.Interface(
+    fn=analyze_code_only,
+    inputs=gr.Textbox(lines=20, label="Python code"),
+    outputs=gr.Markdown(label="Analysis"),
+    api_name="analyze_code_structure",
+    description="Return a Markdown report with complexity metrics.",
+)
+
+iface_all = gr.Interface(
+    fn=generate_all_outputs,
+    inputs=gr.Textbox(lines=20, label="Python code"),
+    outputs=[
+        gr.Image(label="UML diagram"),
+        gr.Image(label="Call‑graph"),
+        gr.Markdown(label="Analysis"),
+    ],
+    api_name="generate_all",
+    description="Run class diagram, call graph and analysis in one call.",
+)
+
+
+# =============================================================================
+# ❸  The Cleaned-up Web UI (using gr.Blocks)
+# =============================================================================
 with gr.Blocks(
-    title="Python UML Diagram Generator & MCP Server",
-    theme=gr.themes.Soft(),
-    css="""
-    .gradio-container {
-        max-width: 1400px !important;
-    }
-    .code-input {
-        font-family: 'Courier New', monospace !important;
-    }
-    """,
+    title="Python Code Visualizer & Analyzer",
+    theme=gr.themes.Soft(primary_hue="blue"),
+    css=""" .gradio-container { max-width: 1400px !important; } """,
 ) as demo:
-    # Header
+    # iface_class = gr.Interface(fn=generate_class_diagram_only, inputs=gr.Textbox(), outputs=gr.Image(), api_name="generate_class_diagram", description="Create a UML class diagram (PNG) from Python code.", visible =False)
+    # iface_call = gr.Interface(fn=generate_call_graph_only, inputs=gr.Textbox(), outputs=gr.Image(), api_name="generate_call_graph_diagram", description="Generate a function‑call graph (PNG) from Python code.", visible =False)
+    # iface_analysis = gr.Interface(fn=analyze_code_only, inputs=gr.Textbox(), outputs=gr.Markdown(), api_name="analyze_code_structure", description="Return a Markdown report with complexity metrics.", visible =False)
+    # iface_all = gr.Interface(fn=generate_all_outputs, inputs=gr.Textbox(), outputs=[gr.Image(), gr.Image(), gr.Markdown()], api_name="generate_all", description="Run class diagram, call graph and analysis in one call.", visible =False)
     gr.Markdown(
         """
-    # 🐍 Python UML Diagram Generator & MCP Server
-    
-    **Dual Functionality:**
-    - 🖥️ **Web Interface**: Generate UML class diagrams and call graphs from Python code
-    - 🤖 **MCP Server**: Provides tools for AI assistants (Claude Desktop, Cursor, etc.)
-    
-    Transform your Python code into comprehensive visual diagrams and analysis!
-    """
+        # 🐍 Python Code Visualizer & Analyzer
+        **Enter Python code, then choose an action to generate diagrams and analysis.**
+        This app also functions as an MCP Server, exposing four tools for AI assistants.
+        """
     )
 
-    with gr.Tab("🎨 Diagram Generator"):
-        with gr.Row():
-            with gr.Column(scale=1):
-                gr.Markdown("### Input")
-
-                example_files = list_example_files()
+    with gr.Row():
+        # ---------- Left column – inputs and actions -----------------------------------
+        with gr.Column(scale=2):
+            gr.Markdown("### 1. Input Code")
+            
+            example_files = list_example_files()
+            print(f"🔍 Found {len(example_files)} example files: {example_files}")
+            if example_files:
                 example_dropdown = gr.Dropdown(
-                    label="Choose Example",
+                    label="Load an Example",
                     choices=example_files,
-                    value=example_files[0] if example_files else None,
-                )
-
-                code_input = gr.Textbox(
-                    label="Python Code",
-                    placeholder="Paste your Python code here...",
-                    lines=20,
-                    max_lines=35,
-                    value=get_sample_code(example_files[0]) if example_files else "",
-                    elem_classes=["code-input"],
-                )
-
-                with gr.Row():
-                    filename_input = gr.Textbox(
-                        label="Diagram Name",
-                        value="my_diagram",
-                        placeholder="Enter a name for your diagram",
-                        scale=2,
-                    )
-
-                with gr.Row():
-                    generate_diagrams_btn = gr.Button(
-                        "🔄 Generate Diagrams", variant="primary", size="lg"
-                    )
-
-            with gr.Column(scale=1):
-                gr.Markdown("### Generated UML Class Diagram")
-
-                uml_diagram_output = gr.Image(
-                    label="UML Class Diagram",
-                    show_download_button=True,
-                    height=300,
-                )
-                
-                gr.Markdown("### Generated Call Graph Diagram")
-
-                call_graph_output = gr.Image(
-                    label="Function Call Graph",
-                    show_download_button=True,
-                    height=300,
+                    value=example_files[0],
                 )
+                # initial_code = get_sample_code(example_files[0])
+                # initial_code = "# Paste your Python code here\n\nclass MyClass:\n    pass"
+                initial_code = "Choose an example file from dropdown or paster your python code here "
+                # initial_code = get_sample_code("simple_class.py")
+            else:
+                initial_code = "# Paste your Python code here\n\nclass MyClass:\n    pass"
 
-        with gr.Row():
-            gr.Markdown("### Code Analysis")
-            
-        with gr.Row():
-            analysis_output = gr.Textbox(
-                label="Comprehensive Code Analysis",
+            code_input = gr.Textbox(
+                label="Python Code",
+                placeholder="Paste your Python code here…",
                 lines=15,
-                max_lines=25,
-                interactive=False,
-                show_copy_button=True,
+                max_lines=200,
+                value=initial_code,
+                elem_classes=["code-input"],
             )
 
-    with gr.Tab("ℹ️ About & Help"):
-        gr.Markdown(
-            """
-        ## About This Tool
-        
-        This Python UML Diagram Generator helps you visualize the structure of your Python code by creating comprehensive diagrams and analysis.
-        ### Inspiration:
-        The idea for this mcp server was inspired by a tweet made by karpathy [tweet](https://x.com/karpathy/status/1930305209747812559).
-        He makes the point that generated images are easy to discriminate by humans while going through a 300 line LLM generated code is time consuming.
-        This tool aims to provide a visual quick smell test for generated code so that user can quickly identify issues instead of going through the code line by line.
-        This is only a very rough and basic implementation of the idea.
-        Making compound AI systems instead of text-to-text chatbots is the necessary direction. 
-
-        ### ✨ Features:
-        - **UML Class Diagrams**: Automatically identifies classes, methods, attributes, and inheritance
-        - **Call Graph Diagrams**: Visualizes function dependencies and call relationships  
-        - **Code Analysis**: Provides detailed structure analysis with complexity metrics
-        - **MCP Integration**: Works with AI assistants via Model Context Protocol
-        
-        ### 📚 How to Use:
-        1. **Paste Code**: Enter your Python code in the text area
-        2. **Set Name**: Choose a name for your diagrams (optional)
-        3. **Generate**: Click "Generate Diagrams" to create all visualizations and analysis
-        4. **Download**: Save the generated diagram images
-        5. **Review**: Read the comprehensive code analysis
-        
-        ### 🔧 Technical Details:
-        - Built with **Gradio** for the web interface
-        - Uses **py2puml** for Python-to-PlantUML conversion
-        - **PlantUML** for UML diagram rendering
-        - **pyan3** and **Graphviz** for call graph generation
-        - **AST** (Abstract Syntax Tree) for code analysis
-        
-        ### 💡 Tips:
-        - Include type hints for better diagram quality
-        - Use meaningful class and method names
-        - Keep inheritance hierarchies clear
-        - Add docstrings for better understanding
-        - Works great with both class-based and function-based code
-        
-        ### 🐛 Troubleshooting:
-        - **No UML diagram generated**: Check if your code contains class definitions
-        - **No call graph generated**: Ensure your code has function definitions and calls
-        - **Syntax errors**: Ensure your Python code is valid
-        - **Import errors**: Stick to standard library imports for best results
-        
-        ## Model Context Protocol (MCP) Server
-        
-        This application automatically serves as an MCP server for AI assistants!
-        
-        ### 🌐 For Hugging Face Spaces (Public):
-        ```json
-        {
-          "mcpServers": {
-            "python-diagram-generator": {
-              "url": "https://your-username-space-name.hf.space/gradio_api/mcp/sse"
-            }
-          }
-        }
-        ```
-        
-        ### 🏠 For Local Development:
-        ```json
-        {
-          "mcpServers": {
-            "python-diagram-generator": {
-              "command": "npx",
-              "args": [
-                "mcp-remote",
-                "http://127.0.0.1:7860/gradio_api/mcp/sse",
-                "--transport",
-                "sse-only"
-              ]
-            }
-          }
-        }
-        ```
-        
-        ### 🔒 For Private Spaces:
-        ```json
-        {
-          "mcpServers": {
-            "python-diagram-generator": {
-              "url": "https://your-username-space-name.hf.space/gradio_api/mcp/sse",
-              "headers": {
-                "Authorization": "Bearer hf_your_token_here"
-              }
-            }
-          }
-        }
-        ```
-        
-        ### 📋 Setup Instructions:
-        1. Install Node.js and mcp-remote: `npm install -g mcp-remote`
-        2. Add the configuration above to your MCP client
-        3. Restart your MCP client (e.g., Claude Desktop)
-        4. Test with prompts like: "Generate a UML diagram for this Python code: [your code]"
-        
-        ---
-        
-        **Local MCP Endpoint**: `http://127.0.0.1:7860/gradio_api/mcp/sse`  
-        **MCP Schema**: View at `/gradio_api/mcp/schema`
-        
-        ### 🚀 Future Features:
-        - Logic flowcharts
-        - Data flow diagrams  
-        - State machine diagrams
-        - Multi-file analysis
-        - Enhanced UML features
-        """
-        )
-
+            gr.Markdown("### 2. Choose an Action")
+            with gr.Row():
+                class_btn = gr.Button("🖼️ Generate Class Diagram")
+                call_graph_btn = gr.Button("🔗 Generate Call Graph")
+                analyze_btn = gr.Button("📈 Analyze Code")
+                all_btn = gr.Button("✨ Generate All", variant="primary")
+
+        # ---------- Right column – outputs ---------------------------------
+        with gr.Column(scale=3):
+            gr.Markdown("### 3. Results")
+            with gr.Tabs():
+                with gr.TabItem("UML Class Diagram"):
+                    uml_output = gr.Image(label="UML Class Diagram", show_download_button=True, interactive=False)
+                with gr.TabItem("Function Call Graph"):
+                    call_graph_output = gr.Image(label="Function Call Graph", show_download_button=True, interactive=False)
+                with gr.TabItem("Code Analysis Report"):
+                    analysis_output = gr.Markdown(label="Comprehensive Code Analysis", elem_classes=["analysis-output"])
+
+    # -------------------------------------------------------------------------
     # Event handlers
-    def load_example(example_filename):
-        return get_sample_code(example_filename)
-
-    example_dropdown.change(
-        fn=load_example,
-        inputs=example_dropdown,
-        outputs=code_input,
+    # -------------------------------------------------------------------------
+
+    # Handler to load example code when dropdown changes
+    if example_files:
+        def _load_example(example_filename: str):
+            return get_sample_code(example_filename)
+        example_dropdown.change(fn=_load_example, inputs=example_dropdown, outputs=code_input, api_name = False)
+
+    # Handlers for the four action buttons
+    # class_btn.click(
+    #     fn=generate_class_diagram_only,
+    #     inputs=[code_input],
+    #     outputs=[uml_output],
+    #     api_name=False # Prevents this from creating a duplicate API endpoint
+    # )
+    class_btn.click(
+        fn=generate_class_diagram_only,
+        inputs=[code_input],
+        outputs=[uml_output],
+        # api_name=False # Prevents this from creating a duplicate API endpoint
+    )
+    
+    call_graph_btn.click(
+        fn=generate_call_graph_only,
+        inputs=[code_input],
+        outputs=[call_graph_output],
+        # api_name=False
     )
 
-    generate_diagrams_btn.click(
-        fn=generate_all_diagrams,
-        inputs=[code_input, filename_input],
-        outputs=[uml_diagram_output, call_graph_output, analysis_output],
-        show_progress=True,
+    analyze_btn.click(
+        fn=analyze_code_only,
+        inputs=[code_input],
+        outputs=[analysis_output],
+        # api_name=False
     )
 
-    code_input.change(
-        fn=analyze_code_structure,
-        inputs=code_input,
-        outputs=analysis_output,
-        show_progress=False,
+    all_btn.click(
+        fn=generate_all_outputs,
+        inputs=[code_input],
+        outputs=[uml_output, call_graph_output, analysis_output],
+        # api_name=False
     )
 
-# Launch configuration
+# =============================================================================
+# ❹  Launch the App and MCP Server
+# =============================================================================
 if __name__ == "__main__":
-    # Setup persistent testing space at startup
-    setup_testing_space()
+    setup_testing_space()  # Create a persistent working dir if needed
 
     demo.launch(
-        mcp_server=True,  # Enable MCP functionality
-        show_api=True,  # Show API documentation
-        show_error=True,  # Show errors in interface
-        share=True,  # Share the app publicly
-        # debug=True,  # Enable debug mode for development
-    )
+        mcp_server=True,    # Enable MCP endpoints (/gradio_api/mcp/*)
+        show_api=True,     # Expose ONLY the 4 Interfaces as tools
+        show_error=True,    # Display exceptions in the UI
+        debug=True,         # Verbose server logs
+    )

utils.py

@@ -0,0 +1,50 @@
+import os
+
+
+def setup_testing_space():
+    """Create persistent testing_space directory and __init__.py at startup."""
+    testing_dir = os.path.join(os.getcwd(), "inputs")
+    os.makedirs(testing_dir, exist_ok=True)
+
+    init_file = os.path.join(testing_dir, "__init__.py")
+    if not os.path.exists(init_file):
+        with open(init_file, "w", encoding="utf-8") as f:
+            f.write("# Testing space for py2puml analysis\n")
+        print("📁 Created testing_space directory and __init__.py")
+    else:
+        print("🔄 testing_space directory already exists")
+
+
+def cleanup_testing_space():
+    """Remove all .py files except __init__.py from testing_space."""
+    testing_dir = os.path.join(os.getcwd(), "inputs")
+    if not os.path.exists(testing_dir):
+        print("⚠️ testing_space directory not found, creating it...")
+        setup_testing_space()
+        return
+
+    # Clean up any leftover .py files (keep only __init__.py)
+    files_removed = 0
+    for file in os.listdir(testing_dir):
+        if file.endswith(".py") and file != "__init__.py":
+            file_path = os.path.join(testing_dir, file)
+            try:
+                os.remove(file_path)
+                files_removed += 1
+            except Exception as e:
+                print(f"⚠️ Could not remove {file}: {e}")
+
+    if files_removed > 0:
+        print(f"🧹 Cleaned up {files_removed} leftover .py files from testing_space")
+
+
+def verify_testing_space():
+    """Verify testing_space contains only __init__.py."""
+    testing_dir = os.path.join(os.getcwd(), "inputs")
+    if not os.path.exists(testing_dir):
+        return False
+
+    files = os.listdir(testing_dir)
+    expected_files = ["__init__.py"]
+
+    return files == expected_files

uv.lock

@@ -466,7 +466,7 @@ http = [
 
 [[package]]
 name = "gradio"
-version = "5.33.0"
+version = "5.33.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "aiofiles" },
@@ -499,9 +499,9 @@ dependencies = [
     { name = "urllib3", marker = "sys_platform == 'emscripten'" },
     { name = "uvicorn", marker = "sys_platform != 'emscripten'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/b0/97/908eb543fbce7c69250d6fbe87b6ccf4ce397d31bceb360b40316357c68c/gradio-5.33.0.tar.gz", hash = "sha256:0cba3a1596fda6cb0048dd7ddc2d57e6238a047c0df9dee5a4a0e5c2a74e8e50", size = 64888401, upload-time = "2025-06-04T21:47:57.431Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/5d/b9/1b59586d91dc5fa083c067789c06a6fbd288242351b506f5af26109188cb/gradio-5.33.1.tar.gz", hash = "sha256:f74c737aa92fc02b4d7dca7e50ee13ddce548aa16c9fcbe907ceabf93722f94d", size = 64908955, upload-time = "2025-06-10T00:51:35.594Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/4f/c3/c9b09b8d7efd63d83a9c8d9c53b02e1b77238e14305a7ee561e0a8990465/gradio-5.33.0-py3-none-any.whl", hash = "sha256:165e412e1510a22471901744722f99a52cb56465a7e9609f1e400cac9999e9d8", size = 54208887, upload-time = "2025-06-04T21:47:52.002Z" },
+    { url = "https://files.pythonhosted.org/packages/8a/de/209775586904b55395e046b0e201d74a0bf91aa06fd087fb0fdcbaa39134/gradio-5.33.1-py3-none-any.whl", hash = "sha256:c4329b04280d62041fbf0113e94fb5c4d20e0555ce1ac69174bf98225350159b", size = 54221598, upload-time = "2025-06-10T00:51:29.237Z" },
 ]
 
 [package.optional-dependencies]
@@ -512,7 +512,7 @@ mcp = [
 
 [[package]]
 name = "gradio-client"
-version = "1.10.2"
+version = "1.10.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "fsspec" },
@@ -522,9 +522,9 @@ dependencies = [
     { name = "typing-extensions" },
     { name = "websockets" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/d2/86/6684afe8691b024200fdc8983924f04b5f76bb401b9c700e5752a23595a0/gradio_client-1.10.2.tar.gz", hash = "sha256:bf71ba95714784fa77ca0cfb20189ad91c55e563c2dc71722d023a97f1815d7f", size = 321294, upload-time = "2025-05-30T13:59:55.756Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/0b/91/a31536da8fd18ed1c1d5b05929b7291a7bfe5963dfcd37a20a42fc2194f0/gradio_client-1.10.3.tar.gz", hash = "sha256:9e99b88e47f05dc3b68e40a3f3f83819f8d0ddcd43466ad385fe42e137825774", size = 321637, upload-time = "2025-06-10T00:51:46.408Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/9b/1b/b372308c263379ae3ebc440512432979458330113bdee26cef86c89bf48e/gradio_client-1.10.2-py3-none-any.whl", hash = "sha256:6de67b6224123d264c7887caa0586b2a9e2c369ec32ca38927cf8a841694edcd", size = 323311, upload-time = "2025-05-30T13:59:54.555Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/72/1e76abc821f8efaaeb2e3bd727a6c97bf87c6a9a0ffacfed0647e587824a/gradio_client-1.10.3-py3-none-any.whl", hash = "sha256:941e7f8d9a160f88487e9780a3db2736a40ea2b8b69d53ffdb306e47ef658b76", size = 323599, upload-time = "2025-06-10T00:51:45.204Z" },
 ]
 
 [[package]]
@@ -857,7 +857,7 @@ wheels = [
 
 [[package]]
 name = "mcp"
-version = "1.9.0"
+version = "1.9.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "anyio" },
@@ -870,9 +870,9 @@ dependencies = [
     { name = "starlette" },
     { name = "uvicorn", marker = "sys_platform != 'emscripten'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/bc/8d/0f4468582e9e97b0a24604b585c651dfd2144300ecffd1c06a680f5c8861/mcp-1.9.0.tar.gz", hash = "sha256:905d8d208baf7e3e71d70c82803b89112e321581bcd2530f9de0fe4103d28749", size = 281432, upload-time = "2025-05-15T18:51:06.615Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/f2/df/8fefc0c6c7a5c66914763e3ff3893f9a03435628f6625d5e3b0dc45d73db/mcp-1.9.3.tar.gz", hash = "sha256:587ba38448e81885e5d1b84055cfcc0ca56d35cd0c58f50941cab01109405388", size = 333045, upload-time = "2025-06-05T15:48:25.681Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/a5/d5/22e36c95c83c80eb47c83f231095419cf57cf5cca5416f1c960032074c78/mcp-1.9.0-py3-none-any.whl", hash = "sha256:9dfb89c8c56f742da10a5910a1f64b0d2ac2c3ed2bd572ddb1cfab7f35957178", size = 125082, upload-time = "2025-05-15T18:51:04.916Z" },
+    { url = "https://files.pythonhosted.org/packages/79/45/823ad05504bea55cb0feb7470387f151252127ad5c72f8882e8fe6cf5c0e/mcp-1.9.3-py3-none-any.whl", hash = "sha256:69b0136d1ac9927402ed4cf221d4b8ff875e7132b0b06edd446448766f34f9b9", size = 131063, upload-time = "2025-06-05T15:48:24.171Z" },
 ]
 
 [[package]]