Update app.py

app.py

@@ -54,19 +54,23 @@ def call_model(prompt: str, model, tokenizer) -> str:
 class AgentState(TypedDict):
     messages: List[Dict[str, str]]
     html: str
+    refined_request: str
     final_prompt: str
     feedback: str
     iteration: int
     done: bool
     timings: Dict[str, float]
 
-def agent(prompt_template, state: AgentState, agent_key: str, timing_label: str):
+def agent(template: str, state: AgentState, agent_key: str, timing_label: str):
     start = time.time()
     model, tokenizer = load_agent_model(**AGENT_MODEL_CONFIG[agent_key])
 
-    latest_input = state.get("final_prompt", state["messages"][-1]["content"])
-    prompt = prompt_template.format(user_input=latest_input, html=state.get("html", ""), final_prompt=state.get("final_prompt", ""))
-    
+    latest_input = (
+        state.get("final_prompt")
+        or state.get("refined_request")
+        or state["messages"][-1]["content"]
+    )
+    prompt = template.format(user_input=latest_input, html=state.get("html", ""), final_prompt=state.get("final_prompt", ""))
     response = call_model(prompt, model, tokenizer)
     state["messages"].append({"role": agent_key, "content": response})
     state["timings"][timing_label] = time.time() - start
@@ -74,16 +78,17 @@ def agent(prompt_template, state: AgentState, agent_key: str, timing_label: str)
     return response
 
 PROMPTS = {
-    "product_manager": "You're a Product Manager. Refine this user request:\n{user_input}",
-    "project_manager": "You're a Project Manager. Break down this refined request into a clear spec:\n{user_input}",
-    "software_engineer": "You're a Software Engineer. Generate clean HTML+CSS code for this spec:\n{final_prompt}",
-    "qa_engineer": "You're a QA Engineer. Review this HTML:\n{html}\nGive feedback or reply APPROVED."
+    "product_manager": "{user_input}",
+    "project_manager": "{user_input}",
+    "software_engineer": "{final_prompt}",
+    "qa_engineer": "{html}"
 }
 
 def generate_ui(user_prompt: str, max_iter: int):
     state: AgentState = {
         "messages": [{"role": "user", "content": user_prompt}],
         "html": "",
+        "refined_request": "",
         "final_prompt": "",
         "feedback": "",
         "iteration": 0,
@@ -92,23 +97,28 @@ def generate_ui(user_prompt: str, max_iter: int):
     }
 
     workflow = StateGraph(AgentState)
+
     workflow.add_node("product_manager", lambda s: {
         "messages": s["messages"] + [{
             "role": "product_manager",
-            "content": agent(PROMPTS["product_manager"], s, "product_manager", "product_manager")
-        }]
+            "content": (pm := agent(PROMPTS["product_manager"], s, "product_manager", "product_manager"))
+        }],
+        "refined_request": pm
     })
+
     workflow.add_node("project_manager", lambda s: {
         "messages": s["messages"] + [{
             "role": "project_manager",
-            "content": (pm_out := agent(PROMPTS["project_manager"], s, "project_manager", "project_manager"))
+            "content": (pr := agent(PROMPTS["project_manager"], s, "project_manager", "project_manager"))
         }],
-        "final_prompt": pm_out
+        "final_prompt": pr
     })
+
     workflow.add_node("software_engineer", lambda s: {
-        "html": (html_out := agent(PROMPTS["software_engineer"], s, "software_engineer", "software_engineer")),
-        "messages": s["messages"] + [{"role": "software_engineer", "content": html_out}]
+        "html": (html := agent(PROMPTS["software_engineer"], s, "software_engineer", "software_engineer")),
+        "messages": s["messages"] + [{"role": "software_engineer", "content": html}]
     })
+
     def qa_fn(s):
         feedback = agent(PROMPTS["qa_engineer"], s, "qa_engineer", "qa_engineer")
         done = "APPROVED" in feedback or s["iteration"] >= max_iter
@@ -118,6 +128,7 @@ def generate_ui(user_prompt: str, max_iter: int):
             "iteration": s["iteration"] + 1,
             "messages": s["messages"] + [{"role": "qa_engineer", "content": feedback}]
         }
+
     workflow.add_node("qa_engineer", qa_fn)
 
     workflow.add_edge("product_manager", "project_manager")