Update app.py

app.py

@@ -47,151 +47,92 @@ def get_model_and_tokenizer() -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
             raise e
     return models["7B"], tokenizers["7B"]
 
-# --- Default Prompt Templates for Multiple Presets ---
+# --- Refactored Default Prompt Templates ---
 default_prompts = {
     "coding": {
-        "brainstorm": """**Coding Brainstorm (Round 1)**
-As a Senior Code Analyst, analyze the following problem and list key challenges and potential approaches.
-
-**User Request:**
-{user_prompt}
-
-**Guidelines:**
-1. Identify coding challenges.
-2. Suggest potential methods and approaches.
-3. Highlight any critical edge cases.
-""",
-        "round2": """**Advanced Reasoning & Code Generation (Round 2)**
-Based on your initial analysis:
-
-**Initial Analysis:**
-{brainstorm_response}
-
-**User Request:**
-{user_prompt}
-
-**Task:**
-1. Generate production-ready code with advanced reasoning.
-2. Include a pun-filled birthday message with a coding twist within your output.
-3. Comment the code clearly.
-""",
-        "synthesis": """**Synthesis & Final Refinement (Round 3)**
-Review the detailed code and reasoning below, and synthesize a final, refined response that:
-1. Combines the brainstorming insights and advanced code generation.
-2. Summarizes the solution succinctly.
-3. Provides any additional improvements.
- 
-**Detailed Code & Reasoning:**
-{round2_response}
-""",
-        "rationale": """**Pun Generation and Rationale (Round 4)**
-Based on the final refined response below, generate a clear, stand-alone pun-filled birthday message with a coding twist, then explain in detail why that pun was chosen.
-
-Final Refined Response:
-{final_response}
-
-Your answer should:
-1. Clearly output the pun as a separate line.
-2. Explain the pun’s connection to birthdays and coding concepts (e.g., binary, syntax).
-3. Describe any creative insights behind the choice.
-"""
+        "brainstorm": (
+            "**Round 1: Brainstorm & Analysis**\n"
+            "Please analyze the following coding challenge or question. Consider the overall problem, "
+            "potential edge cases, and any assumptions you might need to make. Explain your reasoning as you think aloud.\n\n"
+            "**User Request:**\n{user_prompt}\n"
+        ),
+        "round2": (
+            "**Round 2: Detailed Reasoning & Strategy**\n"
+            "Based on your initial analysis, please break down the problem into logical steps. "
+            "Outline a plan or strategy that could be used to solve the challenge, highlighting key algorithms, structures, or design considerations.\n\n"
+            "**Initial Analysis:**\n{brainstorm_response}\n\n"
+            "**User Request:**\n{user_prompt}\n"
+        ),
+        "synthesis": (
+            "**Round 3: Synthesis & Implementation**\n"
+            "Taking into account the steps outlined previously, synthesize a coherent solution. "
+            "Provide a detailed explanation of how the code addresses the problem while encouraging best practices and clear logic.\n\n"
+            "**Detailed Strategy:**\n{round2_response}\n"
+        ),
+        "rationale": (
+            "**Round 4: Reflection & Final Output**\n"
+            "Review your solution and provide a final, well-rounded response that summarizes your reasoning and the implementation strategy. "
+            "Explain any key decisions made during the process and how they contribute to an effective solution.\n\n"
+            "**Final Draft:**\n{final_response}\n"
+        )
     },
     "math": {
-        "brainstorm": """**Math Problem Brainstorm (Round 1)**
-As an expert mathematician, analyze the following problem and outline key concepts and strategies.
-
-**Problem:**
-{user_prompt}
-
-**Guidelines:**
-1. Identify the mathematical concepts involved.
-2. List potential strategies or methods.
-3. Note any assumptions or conditions.
-""",
-        "round2": """**Solution Strategy Development (Round 2)**
-Based on the initial analysis:
-
-**Initial Analysis:**
-{brainstorm_response}
-
-**Problem:**
-{user_prompt}
-
-**Task:**
-1. Develop a detailed strategy to solve the problem.
-2. Include potential methods and intermediate steps.
-""",
-        "synthesis": """**Solution Synthesis (Round 3)**
-Review the strategy and previous analysis below, and produce a refined, step-by-step solution that:
-1. Clearly explains the solution path.
-2. Highlights key steps and justifications.
-3. Summarizes the final answer.
- 
-**Detailed Strategy:**
-{round2_response}
-""",
-        "rationale": """**Solution Rationale (Round 4)**
-Based on the final refined solution below, provide a detailed explanation of the key steps and mathematical insights.
-
-Final Refined Solution:
-{final_response}
-
-Your response should:
-1. Clearly explain why each step was taken.
-2. Detail any assumptions and mathematical principles used.
-3. Summarize the creative reasoning behind the solution.
-"""
+        "brainstorm": (
+            "**Round 1: Problem Analysis & Exploration**\n"
+            "Carefully analyze the mathematical problem provided. Describe the underlying concepts and any assumptions you are making. "
+            "Detail your initial reasoning and potential methods to tackle the problem.\n\n"
+            "**Problem:**\n{user_prompt}\n"
+        ),
+        "round2": (
+            "**Round 2: Detailed Reasoning & Methodology**\n"
+            "Based on your initial exploration, break down the problem into sequential steps or methodologies. "
+            "Explain the reasoning behind each step and how they connect to solve the problem.\n\n"
+            "**Initial Analysis:**\n{brainstorm_response}\n\n"
+            "**Problem:**\n{user_prompt}\n"
+        ),
+        "synthesis": (
+            "**Round 3: Synthesis & Step-by-Step Solution**\n"
+            "Integrate your previous reasoning into a structured solution. Clearly explain each step of your calculation or proof, "
+            "ensuring that your logical progression is easy to follow.\n\n"
+            "**Detailed Methodology:**\n{round2_response}\n"
+        ),
+        "rationale": (
+            "**Round 4: Reflection & Final Explanation**\n"
+            "Present your final solution along with a detailed explanation of the reasoning behind each step. "
+            "Discuss any assumptions and insights that helped you arrive at the final answer.\n\n"
+            "**Final Solution:**\n{final_response}\n"
+        )
     },
     "writing": {
-        "brainstorm": """**Creative Brainstorm (Round 1)**
-As a seasoned writer, brainstorm creative ideas for the following writing prompt.
-
-**Writing Prompt:**
-{user_prompt}
-
-**Guidelines:**
-1. List key themes and creative directions.
-2. Suggest multiple approaches to the narrative.
-3. Highlight any unique stylistic ideas.
-""",
-        "round2": """**Outline Generation (Round 2)**
-Based on the brainstorming below:
-
-**Brainstormed Ideas:**
-{brainstorm_response}
-
-**Writing Prompt:**
-{user_prompt}
-
-**Task:**
-1. Generate a detailed outline for a creative piece.
-2. Organize the ideas into a coherent structure.
-3. Provide bullet points or sections for the narrative.
-""",
-        "synthesis": """**Draft Writing (Round 3)**
-Review the outline below and produce a refined draft of the creative piece that:
-1. Synthesizes the brainstorming insights and the outline.
-2. Provides a coherent and engaging narrative.
-3. Includes stylistic and thematic elements.
- 
-**Outline:**
-{round2_response}
-""",
-        "rationale": """**Final Editing and Rationale (Round 4)**
-Based on the final draft below, refine the piece further and provide a detailed explanation of your creative choices.
-
-Final Draft:
-{final_response}
-
-Your answer should:
-1. Present the final refined text.
-2. Explain the narrative choices, stylistic decisions, and thematic connections.
-3. Detail any creative insights that influenced the final version.
-"""
+        "brainstorm": (
+            "**Round 1: Creative Exploration & Conceptualization**\n"
+            "Read the following writing prompt and explore its themes, tone, and potential narrative directions. "
+            "Outline your initial thoughts and reasoning behind various creative choices.\n\n"
+            "**Writing Prompt:**\n{user_prompt}\n"
+        ),
+        "round2": (
+            "**Round 2: Detailed Outline & Narrative Structure**\n"
+            "Based on your brainstorming, create a detailed outline that organizes the narrative or essay. "
+            "Explain the reasoning behind your structure, the flow of ideas, and how you plan to incorporate creative elements.\n\n"
+            "**Initial Brainstorming:**\n{brainstorm_response}\n\n"
+            "**Writing Prompt:**\n{user_prompt}\n"
+        ),
+        "synthesis": (
+            "**Round 3: Draft Synthesis & Refinement**\n"
+            "Integrate your outline and creative ideas into a coherent draft. Provide a well-rounded narrative that is both engaging and logically structured. "
+            "Explain your thought process as you refine the narrative.\n\n"
+            "**Outline & Strategy:**\n{round2_response}\n"
+        ),
+        "rationale": (
+            "**Round 4: Reflection & Final Editing**\n"
+            "Review your draft and provide a final version that reflects thoughtful editing and creative reasoning. "
+            "Explain the choices made in refining the text, from structure to stylistic decisions.\n\n"
+            "**Final Draft:**\n{final_response}\n"
+        )
     }
 }
 
-# --- Domain Detection ---
+# --- Domain Detection (Retained for fallback) ---
 def detect_domain(user_prompt: str) -> str:
     """
     Detect the domain based on keywords.
@@ -204,7 +145,7 @@ def detect_domain(user_prompt: str) -> str:
     math_keywords = ["solve", "integral", "derivative", "equation", "proof", "calculate", "sum", "product"]
     writing_keywords = ["write", "story", "essay", "novel", "poem", "article", "narrative", "creative"]
     coding_keywords = ["code", "program", "debug", "compile", "algorithm", "function"]
-    
+
     if any(kw in prompt_lower for kw in math_keywords):
         logging.info("Domain detected as: math")
         return "math"
@@ -286,7 +227,7 @@ class MultiRoundAgent:
         r1 = generate_response(self.model, self.tokenizer, prompt_r1, params.get("max_new_tokens"), params.get("temp"), params.get("top_p"))
         self.memory_manager.store(f"Round 1 Response: {r1}")
 
-        # Round 2: Secondary Generation (strategy/outline/code)
+        # Round 2: Secondary Generation (detailed reasoning/strategy)
         logging.info("--- Round 2 ---")
         prompt_r2 = self.prompt_templates["round2"].format(brainstorm_response=r1, user_prompt=user_prompt)
         r2 = generate_response(self.model, self.tokenizer, prompt_r2, params.get("max_new_tokens") + 100, params.get("temp"), params.get("top_p"))
@@ -343,7 +284,7 @@ class MultiRoundAgent:
 def swarm_agent_iterative(user_prompt: str, temp: float, top_p: float, max_new_tokens: int, memory_top_k: int,
                           prompt_templates: Dict[str, str], domain: str, show_raw: bool) -> Generator[str, None, None]:
     """
-    Wraps the multi-round agent functionality. Depending on the detected domain,
+    Wraps the multi-round agent functionality. Depending on the detected or selected domain,
     it runs the 4-round pipeline.
     """
     model, tokenizer = get_model_and_tokenizer()
@@ -364,12 +305,12 @@ def handle_explanation_request(user_prompt: str, history: List) -> str:
             explanation_prompt += f"- {item}\n"
     else:
         explanation_prompt += "No stored final output found.\n"
-    
+
     explanation_prompt += "\nRecent related exchanges:\n"
     for chat in history:
         if ("explain" in chat[0].lower()) or (chat[1] and "explain" in chat[1].lower()):
             explanation_prompt += f"User: {chat[0]}\nAssistant: {chat[1]}\n"
-    
+
     explanation_prompt += "\nBased on the above context, please provide a detailed explanation of the creative choices."
     model, tokenizer = get_model_and_tokenizer()
     explanation = generate_response(model, tokenizer, explanation_prompt, 300, 0.7, 0.9)
@@ -392,10 +333,11 @@ def format_history(history: List) -> List[Dict[str, str]]:
     return messages
 
 # --- Gradio Chat Interface Function ---
-def gradio_interface(message: str, history: List, param_state: Dict, prompt_state: Dict) -> Generator[List[Dict[str, str]], None, None]:
+def gradio_interface(message: str, history: List, param_state: Dict, prompt_state: Dict, mode: str) -> Generator[List[Dict[str, str]], None, None]:
     """
     Called by Gradio's ChatInterface. Uses current generation parameters and preset prompt templates.
     If the user asks for an explanation, routes the request accordingly.
+    The selected mode (coding, math, or writing) overrides automatic domain detection.
     """
     if "explain" in message.lower():
         explanation = handle_explanation_request(message, history)
@@ -413,8 +355,9 @@ def gradio_interface(message: str, history: List, param_state: Dict, prompt_stat
         logging.error(f"Parameter conversion error: {e}")
         temp, top_p, max_new_tokens, memory_top_k, show_raw = 0.5, 0.9, 300, 2, False
 
-    domain = detect_domain(message)
-    # Get the prompt templates for the detected domain; default to coding if not set.
+    # Use selected mode if provided; otherwise, fallback to domain detection.
+    domain = mode if mode in default_prompts else detect_domain(message)
+    # Get the prompt templates for the chosen domain.
     prompt_templates = prompt_state.get(domain, default_prompts.get(domain, default_prompts["coding"]))
 
     history = history + [[message, ""]]
@@ -436,7 +379,7 @@ ui_description = '''
 <div>
   <h1 style="text-align: center;">DeepSeek Agent Swarm Chat</h1>
   <p style="text-align: center;">
-    Multi-round agent with 4-round prompt chaining for three presets:
+    Multi-round agent with 4-round prompt chaining, supporting three modes:
     <br>- Coding
     <br>- Math
     <br>- Writing
@@ -488,16 +431,18 @@ with gr.Blocks(css=css, title="DeepSeek Agent Swarm Chat") as demo:
 
     with gr.Tabs():
         with gr.Tab("Chat"):
+            # Add a mode selector for explicit domain selection.
+            mode_selector = gr.Radio(choices=["coding", "math", "writing"], value="coding", label="Select Mode")
             chatbot = gr.Chatbot(height=450, placeholder=ui_placeholder, label="Agent Swarm Output", type="messages")
             gr.ChatInterface(
                 fn=gradio_interface,
                 chatbot=chatbot,
-                additional_inputs=[param_state, prompt_state],
+                additional_inputs=[param_state, prompt_state, mode_selector],
                 examples=[
                     ['How can we build a robust web service that scales efficiently under load?'],
                     ['Solve the integral of x^2 from 0 to 1.'],
                     ['Write a short story about a mysterious writer in a busy city.'],
-                    ['Create a pun-filled birthday message with a coding twist.']
+                    ['Create a creative and reflective solution for a coding challenge.']
                 ],
                 cache_examples=False,
                 type="messages",
@@ -619,4 +564,4 @@ with gr.Blocks(css=css, title="DeepSeek Agent Swarm Chat") as demo:
     gr.Markdown(ui_license)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)