Update app.py

app.py

@@ -54,6 +54,7 @@ def call_model(prompt: str, model, tokenizer) -> str:
 class AgentState(TypedDict):
     messages: List[Dict[str, str]]
     html: str
+    final_prompt: str
     feedback: str
     iteration: int
     done: bool
@@ -63,12 +64,10 @@ def agent(prompt_template, state: AgentState, agent_key: str, timing_label: str)
     start = time.time()
     model, tokenizer = load_agent_model(**AGENT_MODEL_CONFIG[agent_key])
 
-    latest_input = state["messages"][-1]["content"]
-    html_content = state.get("html", "")
+    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", ""))
     
-    prompt = prompt_template.format(user_input=latest_input, html=html_content)
     response = call_model(prompt, model, tokenizer)
-    
     state["messages"].append({"role": agent_key, "content": response})
     state["timings"][timing_label] = time.time() - start
     gc.collect()
@@ -76,8 +75,8 @@ def agent(prompt_template, state: AgentState, agent_key: str, timing_label: str)
 
 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:\n{user_input}",
-    "software_engineer": "You're a Software Engineer. Generate HTML+CSS code for:\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."
 }
 
@@ -85,6 +84,7 @@ def generate_ui(user_prompt: str, max_iter: int):
     state: AgentState = {
         "messages": [{"role": "user", "content": user_prompt}],
         "html": "",
+        "final_prompt": "",
         "feedback": "",
         "iteration": 0,
         "done": False,
@@ -101,15 +101,13 @@ def generate_ui(user_prompt: str, max_iter: int):
     workflow.add_node("project_manager", lambda s: {
         "messages": s["messages"] + [{
             "role": "project_manager",
-            "content": agent(PROMPTS["project_manager"], s, "project_manager", "project_manager")
-        }]
+            "content": (pm_out := agent(PROMPTS["project_manager"], s, "project_manager", "project_manager"))
+        }],
+        "final_prompt": pm_out
     })
     workflow.add_node("software_engineer", lambda s: {
-        "html": agent(PROMPTS["software_engineer"], s, "software_engineer", "software_engineer"),
-        "messages": s["messages"] + [{
-            "role": "software_engineer",
-            "content": s["html"]
-        }]
+        "html": (html_out := agent(PROMPTS["software_engineer"], s, "software_engineer", "software_engineer")),
+        "messages": s["messages"] + [{"role": "software_engineer", "content": html_out}]
     })
     def qa_fn(s):
         feedback = agent(PROMPTS["qa_engineer"], s, "qa_engineer", "qa_engineer")
@@ -148,4 +146,4 @@ def main():
                 st.markdown(f"**{msg['role'].title()}**:\n```\n{msg['content']}\n```")
 
 if __name__ == "__main__":
-    main()
+    main()