Update app.py

app.py

@@ -1,19 +1,16 @@
-# Multi-agent UI generator using direct model.generate() instead of pipeline()
-
 import streamlit as st
+import os
 import time
-import base64
+import gc
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 from typing import Dict, List, TypedDict
 from langgraph.graph import StateGraph, END
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
-import torch
-import os
-
-st.set_page_config(page_title="Multi-Agent Collaboration", layout="wide")
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+# Agent model config — all use Gemma
 AGENT_MODEL_CONFIG = {
     "product_manager": {
         "base": "unsloth/gemma-3-1b-it",
@@ -23,168 +20,115 @@ AGENT_MODEL_CONFIG = {
         "base": "unsloth/gemma-3-1b-it",
         "adapter": "spandana30/project-manager-gemma"
     },
-    "software_architect": {
-        "base": "unsloth/c4ai-command-r-08-2024-bnb-4bit",
-        "adapter": "spandana30/software-architect-cohere"
-    },
     "software_engineer": {
-        "base": "codellama/CodeLLaMA-7b-hf",
-        "adapter": "spandana30/software-engineer-codellama"
-    },
-    "qa": {
-        "base": "codellama/CodeLLaMA-7b-hf",
-        "adapter": "spandana30/software-engineer-codellama"
+        "base": "unsloth/gemma-3-1b-it",
+        "adapter": "spandana30/project-manager-gemma"
     },
+    "qa_engineer": {
+        "base": "unsloth/gemma-3-1b-it",
+        "adapter": "spandana30/project-manager-gemma"
+    }
 }
 
 @st.cache_resource
-
-def get_model_and_tokenizer(role: str):
-    cfg = AGENT_MODEL_CONFIG[role]
-    try:
-        st.write(f"🔍 Loading model for {role}: {cfg['base']} + {cfg['adapter']}")
-        base_model = AutoModelForCausalLM.from_pretrained(
-            cfg["base"], torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto", token=HF_TOKEN
-        )
-        model = PeftModel.from_pretrained(base_model, cfg["adapter"], token=HF_TOKEN)
-        tokenizer = AutoTokenizer.from_pretrained(cfg["adapter"], token=HF_TOKEN)
-        return model, tokenizer
-    except Exception as e:
-        st.error(f"❌ Failed to load model for {role}\nError: {e}")
-        raise
+def load_agent_model(base_id, adapter_id):
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_id,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+        token=HF_TOKEN
+    )
+    model = PeftModel.from_pretrained(base_model, adapter_id, token=HF_TOKEN)
+    tokenizer = AutoTokenizer.from_pretrained(adapter_id, token=HF_TOKEN)
+    return model.eval(), tokenizer
+
+def call_model(prompt: str, model, tokenizer) -> str:
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=512,
+        do_sample=False,
+        temperature=0.3
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 class AgentState(TypedDict):
     messages: List[Dict[str, str]]
-    user_request: str
-    refined_request: str
-    scoped_request: str
-    design_specs: str
     html: str
     feedback: str
     iteration: int
     done: bool
     timings: Dict[str, float]
 
-def run_pipeline(role: str, prompt: str):
-    model, tokenizer = get_model_and_tokenizer(role)
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=1024,
-            do_sample=False
-        )
-    return tokenizer.decode(output[0], skip_special_tokens=True).strip()
+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)
+    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 and clarify this request:
-{user_request}
-Ensure it's clear, feasible, and user-focused. Output the revised request only.""",
-    "project_manager": """You're a Project Manager. Given this refined request:
-{refined_request}
-Break it down into scope and constraints. Output the scoped request only.""",
-    "designer": """You're a UI designer. Create design specs for:
-{scoped_request}
-Include color palette, font, layout, and component styles. No code.""",
-    "software_engineer": """Create a full HTML page with embedded CSS for:
-{design_specs}
-Requirements:
-- Semantic, responsive HTML
-- Embedded CSS in <style> tag
-- Output complete HTML only.""",
-    "qa": """Review this webpage:
-{html}
-Is it visually appealing, responsive, and functional? Reply "APPROVED" or suggest improvements."""
+    "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]}",
+    "qa_engineer": "You're a QA Engineer. Review this HTML:\n{html}\nGive feedback or reply APPROVED."
 }
 
-def time_agent(agent_func, state: AgentState, label: str):
-    start = time.time()
-    result = agent_func(state)
-    result["timings"] = state["timings"]
-    result["timings"][label] = time.time() - start
-    return result
-
-def product_manager_agent(state: AgentState):
-    revised = run_pipeline("product_manager", PROMPTS["product_manager"].format(user_request=state["user_request"]))
-    return {"refined_request": revised, "messages": state["messages"] + [{"role": "product_manager", "content": revised}]}
-
-def project_manager_agent(state: AgentState):
-    scoped = run_pipeline("project_manager", PROMPTS["project_manager"].format(refined_request=state["refined_request"]))
-    return {"scoped_request": scoped, "messages": state["messages"] + [{"role": "project_manager", "content": scoped}]}
-
-def designer_agent(state: AgentState):
-    specs = run_pipeline("product_manager", PROMPTS["designer"].format(scoped_request=state["scoped_request"]))
-    return {"design_specs": specs, "messages": state["messages"] + [{"role": "designer", "content": specs}]}
-
-def engineer_agent(state: AgentState):
-    html = run_pipeline("software_engineer", PROMPTS["software_engineer"].format(design_specs=state["design_specs"]))
-    return {"html": html, "messages": state["messages"] + [{"role": "software_engineer", "content": html}]}
-
-def qa_agent(state: AgentState, max_iter: int):
-    feedback = run_pipeline("qa", PROMPTS["qa"].format(html=state["html"]))
-    done = "APPROVED" in feedback or state["iteration"] >= max_iter
-    return {"feedback": feedback, "done": done, "iteration": state["iteration"] + 1,
-            "messages": state["messages"] + [{"role": "qa", "content": feedback}]}
-
-def generate_ui(user_request: str, max_iter: int):
-    state = {"messages": [{"role": "user", "content": user_request}],
-             "user_request": user_request,
-             "refined_request": "", "scoped_request": "", "design_specs": "",
-             "html": "", "feedback": "", "iteration": 0, "done": False, "timings": {}}
+def generate_ui(user_prompt: str, max_iter: int):
+    state: AgentState = {
+        "messages": [{"role": "user", "content": user_prompt}],
+        "html": "",
+        "feedback": "",
+        "iteration": 0,
+        "done": False,
+        "timings": {}
+    }
 
     workflow = StateGraph(AgentState)
-    workflow.add_node("product_manager", lambda s: time_agent(product_manager_agent, s, "product_manager"))
-    workflow.add_node("project_manager", lambda s: time_agent(project_manager_agent, s, "project_manager"))
-    workflow.add_node("designer", lambda s: time_agent(designer_agent, s, "designer"))
-    workflow.add_node("software_engineer", lambda s: time_agent(engineer_agent, s, "software_engineer"))
-    workflow.add_node("qa", lambda s: time_agent(lambda x: qa_agent(x, max_iter), s, "qa"))
+    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"]}]
+    })
+    def qa_fn(s): 
+        feedback = agent(PROMPTS["qa_engineer"], s, "qa_engineer", "qa_engineer")
+        done = "APPROVED" in feedback or s["iteration"] >= max_iter
+        return {
+            "feedback": feedback,
+            "done": done,
+            "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")
-    workflow.add_edge("project_manager", "designer")
-    workflow.add_edge("designer", "software_engineer")
-    workflow.add_edge("software_engineer", "qa")
-    workflow.add_conditional_edges("qa", lambda s: END if s["done"] else "software_engineer")
+    workflow.add_edge("project_manager", "software_engineer")
+    workflow.add_edge("software_engineer", "qa_engineer")
+    workflow.add_conditional_edges("qa_engineer", lambda s: END if s["done"] else "software_engineer")
     workflow.set_entry_point("product_manager")
 
     app = workflow.compile()
-    total_start = time.time()
     final_state = app.invoke(state)
-    return final_state["html"], final_state, time.time() - total_start
+    return final_state
 
 def main():
-    st.title("🤖 Multi-Agent UI Generator")
-    with st.sidebar:
-        max_iter = st.slider("Max QA Iterations", 1, 5, 2)
-
-    prompt = st.text_area("📝 Describe the UI you want:", "A coffee shop landing page with hero, menu, and contact form.", height=150)
-
+    st.set_page_config(page_title="Multi-Agent UI Generator", layout="wide")
+    st.title(" Multi-Agent Collaboration")
+    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"):
         with st.spinner("Agents working..."):
-            html, final_state, total_time = generate_ui(prompt, max_iter)
-            st.success("✅ UI Generated Successfully!")
-            st.components.v1.html(html, height=600, scrolling=True)
-
-            b64 = base64.b64encode(html.encode()).decode()
-            st.markdown(f'<a href="data:file/html;base64,{b64}" download="ui.html">📥 Download HTML</a>', unsafe_allow_html=True)
-
-            st.subheader("🧠 Agent Communication Log")
-            history_text = ""
-            for msg in final_state["messages"]:
-                role = msg["role"].replace("_", " ").title()
-                content = msg["content"]
-                history_text += f"---\n{role}:\n{content}\n\n"
-            st.text_area("Agent Dialogue", value=history_text, height=300)
-
-            b64_hist = base64.b64encode(history_text.encode()).decode()
-            st.markdown(
-                f'<a href="data:file/txt;base64,{b64_hist}" download="agent_communication.txt">📥 Download Communication Log</a>',
-                unsafe_allow_html=True)
-
-            st.subheader("📊 Performance")
-            st.write(f"⏱️ Total Time: {total_time:.2f} seconds")
-            st.write(f"🔁 Iterations: {final_state['iteration']}")
-            for stage in ["product_manager", "project_manager", "designer", "software_engineer", "qa"]:
-                st.write(f"🧩 {stage.replace('_', ' ').title()} Time: {final_state['timings'].get(stage, 0):.2f}s")
+            final = generate_ui(prompt, max_iter)
+            st.success("✅ UI Generated")
+            st.subheader("🔍 Output HTML")
+            st.components.v1.html(final["html"], height=600, scrolling=True)
+            st.subheader("🧠 Agent Messages")
+            for msg in final["messages"]:
+                st.markdown(f"**{msg['role'].title()}**:\n```\n{msg['content']}\n```")
 
 if __name__ == "__main__":
     main()