Update app.py

app.py

@@ -10,23 +10,23 @@ from langgraph.graph import StateGraph, END
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Agent model config — all use Gemma
+# ✅ Corrected keys: base_id and adapter_id
 AGENT_MODEL_CONFIG = {
     "product_manager": {
-        "base": "unsloth/gemma-3-1b-it",
-        "adapter": "spandana30/project-manager-gemma"
+        "base_id": "unsloth/gemma-3-1b-it",
+        "adapter_id": "spandana30/project-manager-gemma"
     },
     "project_manager": {
-        "base": "unsloth/gemma-3-1b-it",
-        "adapter": "spandana30/project-manager-gemma"
+        "base_id": "unsloth/gemma-3-1b-it",
+        "adapter_id": "spandana30/project-manager-gemma"
     },
     "software_engineer": {
-        "base": "unsloth/gemma-3-1b-it",
-        "adapter": "spandana30/project-manager-gemma"
+        "base_id": "unsloth/gemma-3-1b-it",
+        "adapter_id": "spandana30/project-manager-gemma"
     },
     "qa_engineer": {
-        "base": "unsloth/gemma-3-1b-it",
-        "adapter": "spandana30/project-manager-gemma"
+        "base_id": "unsloth/gemma-3-1b-it",
+        "adapter_id": "spandana30/project-manager-gemma"
     }
 }
 
@@ -88,13 +88,26 @@ 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")}]})
-    workflow.add_node("project_manager", lambda s: {"messages": s["messages"] + [{"role": "project_manager", "content": agent(PROMPTS["project_manager"], s, "project_manager", "project_manager")}]})
+    workflow.add_node("product_manager", lambda s: {
+        "messages": s["messages"] + [{
+            "role": "product_manager",
+            "content": agent(PROMPTS["product_manager"], s, "product_manager", "product_manager")
+        }]
+    })
+    workflow.add_node("project_manager", lambda s: {
+        "messages": s["messages"] + [{
+            "role": "project_manager",
+            "content": agent(PROMPTS["project_manager"], s, "project_manager", "project_manager")
+        }]
+    })
     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"]}]
+        "messages": s["messages"] + [{
+            "role": "software_engineer",
+            "content": s["html"]
+        }]
     })
-    def qa_fn(s): 
+    def qa_fn(s):
         feedback = agent(PROMPTS["qa_engineer"], s, "qa_engineer", "qa_engineer")
         done = "APPROVED" in feedback or s["iteration"] >= max_iter
         return {
@@ -117,7 +130,7 @@ def generate_ui(user_prompt: str, max_iter: int):
 
 def main():
     st.set_page_config(page_title="Multi-Agent UI Generator", layout="wide")
-    st.title(" Multi-Agent Collaboration")
+    st.title("🤖 Multi-Agent Collaboration (Gemma Only)")
     max_iter = st.sidebar.slider("Max QA Iterations", 1, 5, 2)
     prompt = st.text_area("Describe your UI:", "A landing page for a coffee shop with a hero image, menu, and contact form.", height=150)
     if st.button("🚀 Generate UI"):