Update app.py

app.py

@@ -2,68 +2,55 @@ import os
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-# ✅ Set up Hugging Face client with token from environment
+# ✅ Set up Hugging Face client
 client = InferenceClient(
-        provider="auto",
-        api_key=os.environ["HF_TOKEN"],
-    )
+    model="deepseek-ai/DeepSeek-R1-0528",
+    token=os.environ["HF_TOKEN"]
+)
 
 # 🧠 Function to generate software architecture
 def generate_software_spec(name, description, architecture, components, deployment, platform, extra):
-    print("generate_software_spec")
     prompt = f"""
-        You are a software architect assistant. Based on the following input, generate:
-        1. A **Sequence Diagram** in Mermaid syntax
-        2. A **Business Process Flow** in Mermaid syntax
-        3. Code Snippets for each component based on the selected tech stack.
-        
-        ---
-        
-        App Name: {name}
-        Description: {description}
-        Architecture: {architecture}
-        Deployment: {deployment}
-        Platform: {platform}
-        Components & Tech Stack:
-        {components}
-        Extra requirements: {extra or "None"}
-        
-        Return the result in this format:
-        
-        ### Sequence Diagram (Mermaid)
-        ```mermaid
-        <sequence_diagram_here>
-        ```
-        
-        ### Business Process Flow (Mermaid)
-        ```mermaid
-        <flowchart_here>
-        ```
-        
-        ### Code Snippets
-        #### Component: <Component 1>
-        ```<language>
-        <code_here>
-        ```
-        
-        #### Component: <Component 2>
-        ```<language>
-        <code_here>
-        ```
-        """
+You are a software architect assistant. Based on the following input, generate:
+1. A **Sequence Diagram** in Mermaid syntax
+2. A **Business Process Flow** in Mermaid syntax
+3. Code Snippets for each component based on the selected tech stack.
 
-    print(prompt)
-    # Generate response from Hugging Face Inference Client
-    response = client.chat_completion(
-        messages=[
-            {"role": "system", "content": "You are a helpful software architecture assistant."},
-            {"role": "user", "content": prompt}
-        ],
-        model="deepseek-ai/DeepSeek-R1-0528",
-        max_tokens=2048,
-    )
-    print(response)
-    return response.choices[0].message.content if response and response.choices else "Error: No response received."
+App Name: {name}
+Description: {description}
+Architecture: {architecture}
+Deployment: {deployment}
+Platform: {platform}
+Components & Tech Stack:
+{components}
+Extra requirements: {extra or "None"}
+
+Return the result in this format:
+
+### Sequence Diagram (Mermaid)
+```mermaid
+<sequence_diagram_here>
+```
+
+### Business Process Flow (Mermaid)
+```mermaid
+<flowchart_here>
+```
+
+### Code Snippets
+#### Component: <Component 1>
+```<language>
+<code_here>
+```
+
+#### Component: <Component 2>
+```<language>
+<code_here>
+```
+    """
+    # 🔄 Get text response
+    response = client.text_generation(prompt=prompt, max_tokens=2048)
+    return response if response else "❌ Error: No response received."
 
 # 🎨 Gradio UI
 with gr.Blocks() as demo:
@@ -72,11 +59,11 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             name = gr.Textbox(label="App Name", placeholder="MyApp")
-            description = gr.Textbox(label="Short Description", lines=2, placeholder="A system that manages online tutoring...")
-            architecture = gr.Radio(["Monolithic", "Backend-Frontend"], label="Architecture Style")
-            deployment = gr.Radio(["Serverless", "VM"], label="Deployment Style")
-            platform = gr.Radio(["Web App", "Mobile App"], label="Target Platform")
-            components = gr.Textbox(label="Components & Tech Stack (e.g. Backend: Python + Flask, Frontend: React)", lines=4)
+            description = gr.Textbox(label="Short Description", lines=2, placeholder="A system for ...")
+            architecture = gr.Radio(["Monolithic", "Backend-Frontend"], label="Architecture Style", value="Backend-Frontend")
+            deployment = gr.Radio(["Serverless", "VM"], label="Deployment Style", value="Serverless")
+            platform = gr.Radio(["Web App", "Mobile App"], label="Target Platform", value="Web App")
+            components = gr.Textbox(label="Components & Tech Stack", lines=4, placeholder="Backend: Python + Flask, Frontend: React")
             extra = gr.Textbox(label="Extra Requirements (optional)", lines=2)
             submit = gr.Button("Generate Design")
             status = gr.Markdown(visible=False)
@@ -86,15 +73,23 @@ with gr.Blocks() as demo:
 
     def wrapper(name, description, architecture, components, deployment, platform, extra):
         result = generate_software_spec(name, description, architecture, components, deployment, platform, extra)
-        html_content = """
-        <div>{}</div>
-        <script type='module'>
-            import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs';
-            mermaid.initialize({{startOnLoad: true}});
-        </script>
-        """.format(result)
-
-        return html_content, gr.update(visible=False)
+
+        # Extract Mermaid blocks and wrap them properly
+        def render_mermaid_blocks(text):
+            import re
+            blocks = re.findall(r"```mermaid\n(.*?)```", text, re.DOTALL)
+            html_blocks = "".join([f"<div class='mermaid'>{block}</div>" for block in blocks])
+            return html_blocks or text
+
+        mermaid_html = render_mermaid_blocks(result)
+        final_html = f"""
+<link rel=\"stylesheet\" href=\"https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.css\">
+<script src=\"https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js\"></script>
+<script>mermaid.initialize({{startOnLoad:true}});</script>
+{mermaid_html}
+<pre style='white-space: pre-wrap;'>{result}</pre>
+        """
+        return final_html, gr.update(visible=False)
 
     def show_spinner():
         return gr.update(visible=True, value="<center>⏳ Generating... Please wait.</center>")