Update app.py

app.py

@@ -1,17 +1,60 @@
 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 huggingface_hub import InferenceClient
 
-# Individual clients per role
-product_manager_client = InferenceClient("unsloth/mistral-7b-bnb-4bit", token=st.secrets["HF_TOKEN"])
-project_manager_client = InferenceClient("unsloth/gemma-3-1b-it", token=st.secrets["HF_TOKEN"])
-software_architect_client = InferenceClient("unsloth/c4ai-command-r-08-2024-bnb-4bit", token=st.secrets["HF_TOKEN"])
-software_engineer_client = InferenceClient("codellama/CodeLlama-7b-hf", token=st.secrets["HF_TOKEN"])
-qa_client = InferenceClient("codellama/CodeLlama-7b-hf", token=st.secrets["HF_TOKEN"])
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+AGENT_MODEL_CONFIG = {
+    "product_manager": {
+        "base_id": "unsloth/mistral-7b-bnb-4bit",
+        "adapter_id": "spandana30/product-manager-mistral"
+    },
+    "project_manager": {
+        "base_id": "unsloth/gemma-3-1b-it",
+        "adapter_id": "spandana30/project-manager-gemma"
+    },
+    "designer": {
+        "base_id": "unsloth/gemma-3-1b-it",
+        "adapter_id": "spandana30/project-manager-gemma"
+    },
+    "software_engineer": {
+        "base_id": "codellama/CodeLLaMA-7b-hf",
+        "adapter_id": "spandana30/software-engineer-codellama"
+    },
+    "qa_engineer": {
+        "base_id": "codellama/CodeLLaMA-7b-hf",
+        "adapter_id": "spandana30/software-engineer-codellama"
+    }
+}
+
+@st.cache_resource
+def load_agent_model(base_id, adapter_id):
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_id,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        load_in_4bit=True,
+        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=1024,
+        do_sample=False,
+        temperature=0.3
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 class AgentState(TypedDict):
     messages: List[Dict[str, str]]
@@ -19,156 +62,151 @@ class AgentState(TypedDict):
     project_plan: str
     design_specs: str
     html: str
-    css: str
     feedback: str
     iteration: int
     done: bool
     timings: Dict[str, float]
 
-PRODUCT_MANAGER_PROMPT = """You're a Product Manager. Interpret this user request:
-{user_request}
-Define the high-level product goals, features, and user stories."""
-
-PROJECT_MANAGER_PROMPT = """You're a Project Manager. Based on this feature list:
-{product_vision}
-Create a project plan with key milestones and task assignments."""
-
-DESIGNER_PROMPT = """You're a UI designer. Create design specs for:
-{project_plan}
-Include:
-1. Color palette (primary, secondary, accent)
-2. Font choices
-3. Layout structure
-4. Component styles
-Don't write code - just design guidance."""
-
-ENGINEER_PROMPT = """Create a complete HTML page with embedded CSS for:
-{design_specs}
-Requirements:
-1. Full HTML document with <!DOCTYPE>
-2. CSS inside <style> tags in head
-3. Mobile-responsive
-4. Semantic HTML
-5. Ready-to-use (will work when saved as .html)
-Output JUST the complete HTML file content:"""
-
-QA_PROMPT = """Review this website:
-{html}
-Check for:
-1. Visual quality
-2. Responsiveness
-3. Functionality
-Reply \"APPROVED\" if perfect, or suggest improvements."""
-
-def time_agent(agent_func, state: AgentState, label: str):
+def agent(template: str, state: AgentState, agent_key: str, timing_label: str):
+    st.write(f'🛠 Running agent: {agent_key}')
     start = time.time()
-    result = agent_func(state)
-    duration = time.time() - start
-    result["timings"] = state["timings"]
-    result["timings"][label] = duration
-    return result
-
-def product_manager_agent(state: AgentState):
-    vision = product_manager_client.text_generation(
-        PRODUCT_MANAGER_PROMPT.format(user_request=state["messages"][-1]["content"]),
-        max_new_tokens=1000, temperature=0.3, return_full_text=False
-    )
-    return {"product_vision": vision, "messages": state["messages"] + [{"role": "product_manager", "content": vision}]}
-
-def project_manager_agent(state: AgentState):
-    plan = project_manager_client.text_generation(
-        PROJECT_MANAGER_PROMPT.format(product_vision=state["product_vision"]),
-        max_new_tokens=1000, temperature=0.3, return_full_text=False
+    model, tokenizer = load_agent_model(**AGENT_MODEL_CONFIG[agent_key])
+
+    prompt = template.format(
+        user_request=state["messages"][0]["content"],
+        product_vision=state.get("product_vision", ""),
+        project_plan=state.get("project_plan", ""),
+        design_specs=state.get("design_specs", ""),
+        html=state.get("html", "")
     )
-    return {"project_plan": plan, "messages": state["messages"] + [{"role": "project_manager", "content": plan}]}
-
-def designer_agent(state: AgentState):
-    specs = software_architect_client.text_generation(
-        DESIGNER_PROMPT.format(project_plan=state["project_plan"]),
-        max_new_tokens=1000, temperature=0.3, return_full_text=False
+    st.write(f'📤 Prompt for {agent_key}:', prompt)
+
+    response = call_model(prompt, model, tokenizer)
+    st.write(f'📥 Response from {agent_key}:', response[:500])
+
+    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. Interpret this user request:\\n"
+        "{user_request}\\n"
+        "Define the high-level product goals, features, and user stories."
+    ),
+    "project_manager": (
+        "You're a Project Manager. Based on this feature list:\\n"
+        "{product_vision}\\n"
+        "Create a project plan with key milestones and task assignments."
+    ),
+    "designer": (
+        "You're a UI designer. Create design specs for:\\n"
+        "{project_plan}\\n"
+        "Include:\\n"
+        "1. Color palette (primary, secondary, accent)\\n"
+        "2. Font choices\\n"
+        "3. Layout structure\\n"
+        "4. Component styles\\n"
+        "Don't write code - just design guidance."
+    ),
+    "software_engineer": (
+        "Create a complete HTML page with embedded CSS for:\\n"
+        "{design_specs}\\n"
+        "Requirements:\\n"
+        "1. Full HTML document with <!DOCTYPE>\\n"
+        "2. CSS inside <style> tags in head\\n"
+        "3. Mobile-responsive\\n"
+        "4. Semantic HTML\\n"
+        "5. Ready-to-use (will work when saved as .html)\\n"
+        "Output JUST the complete HTML file content:"
+    ),
+    "qa_engineer": (
+        "Review this website:\\n"
+        "{html}\\n"
+        "Check for:\\n"
+        "1. Visual quality\\n"
+        "2. Responsiveness\\n"
+        "3. Functionality\\n"
+        "Reply \\"APPROVED\\" if perfect, or suggest improvements."
     )
-    return {"design_specs": specs, "messages": state["messages"] + [{"role": "designer", "content": specs}]}
-
-def engineer_agent(state: AgentState):
-    html = software_engineer_client.text_generation(
-        ENGINEER_PROMPT.format(design_specs=state["design_specs"]),
-        max_new_tokens=3000, temperature=0.3, return_full_text=False
-    )
-    if not html.strip().startswith("<!DOCTYPE"):
-        html = f"""<!DOCTYPE html><html><head><meta charset='UTF-8'><meta name='viewport' content='width=device-width, initial-scale=1.0'><title>Generated UI</title></head><body>{html}</body></html>"""
-    return {"html": html, "messages": state["messages"] + [{"role": "software_engineer", "content": html}]}
-
-def qa_agent(state: AgentState, max_iter: int):
-    feedback = qa_client.text_generation(
-        QA_PROMPT.format(html=state["html"]),
-        max_new_tokens=1000, temperature=0.3, return_full_text=False
-    )
-    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}], "product_vision": "", "project_plan": "",
-             "design_specs": "", "html": "", "css": "", "feedback": "", "iteration": 0, "done": False, "timings": {}}
+}
+
+def generate_ui(user_prompt: str, max_iter: int):
+    state: AgentState = {
+        "messages": [{"role": "user", "content": user_prompt}],
+        "product_vision": "",
+        "project_plan": "",
+        "design_specs": "",
+        "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": (pv := agent(PROMPTS["product_manager"], s, "product_manager", "product_manager"))
+        }],
+        "product_vision": pv
+    })
+    workflow.add_node("project_manager", lambda s: {
+        "messages": s["messages"] + [{
+            "role": "project_manager",
+            "content": (pp := agent(PROMPTS["project_manager"], s, "project_manager", "project_manager"))
+        }],
+        "project_plan": pp
+    })
+    workflow.add_node("designer", lambda s: {
+        "messages": s["messages"] + [{
+            "role": "designer",
+            "content": (ds := agent(PROMPTS["designer"], s, "designer", "designer"))
+        }],
+        "design_specs": ds
+    })
+    workflow.add_node("software_engineer", lambda s: {
+        "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
+        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("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.set_page_config(page_title="Multi-Agent Collaboration", layout="wide")
-    st.title("🤝 Multi-Agent Collaboration")
-    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)
-
-    if st.button("🚀 Generate UI"):
+    st.set_page_config(page_title="Multi-Agent UI Generator", layout="wide")
+    st.title("🧠 Multi-Agent UI Generation System")
+    max_iter = st.sidebar.slider("Max QA Iterations", 1, 5, 2)
+    prompt = st.text_area("What UI do you want to build?", "A coffee shop landing page with a hero image, menu, and contact form.", height=150)
+    if st.button("🚀 Generate"):
         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)
-
-            st.subheader("📥 Download HTML")
-            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" '
-                'style="padding: 0.4em 1em; background: #4CAF50; color: white; border-radius: 0.3em; text-decoration: none;">'
-                '📥 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.title().replace('_', ' ')} Time: {final_state['timings'].get(stage, 0):.2f}s")
+            final = generate_ui(prompt, max_iter)
+            st.success("✅ UI Generated")
+            st.subheader("🔍 Final Output")
+            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()