Update app.py

app.py

@@ -10,7 +10,6 @@ from langgraph.graph import StateGraph, END
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# ✅ Corrected keys: base_id and adapter_id
 AGENT_MODEL_CONFIG = {
     "product_manager": {
         "base_id": "unsloth/gemma-3-1b-it",
@@ -63,17 +62,22 @@ class AgentState(TypedDict):
 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])
-    prompt = prompt_template.format(**state)
+
+    latest_input = state["messages"][-1]["content"]
+    html_content = state.get("html", "")
+    
+    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()
     return response
 
 PROMPTS = {
-    "product_manager": "You're a Product Manager. Refine this user request:\n{messages[-1][content]}",
-    "project_manager": "You're a Project Manager. Break down this refined request:\n{messages[-1][content]}",
-    "software_engineer": "You're a Software Engineer. Generate HTML+CSS code for:\n{messages[-1][content]}",
+    "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}",
     "qa_engineer": "You're a QA Engineer. Review this HTML:\n{html}\nGive feedback or reply APPROVED."
 }