Spaces:

spandana30
/

Multi_agent

Build error

App Files Files Community

spandana30 commited on May 7

Commit

50e0854

verified ·

1 Parent(s): 9ccfc13

Update app.py

Browse files

Files changed (1) hide show

app.py +142 -155

app.py CHANGED Viewed

@@ -1,185 +1,172 @@
 import streamlit as st
 import os
 import time
-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
-HF_TOKEN = os.getenv("HF_TOKEN")
-AGENT_MODEL_CONFIG = {
-    "product_manager": {
-        "base_id": "unsloth/gemma-3-1b-it",
-        "adapter_id": "spandana30/project-manager-gemma"
-    },
-    "project_manager": {
-        "base_id": "unsloth/gemma-3-1b-it",
-        "adapter_id": "spandana30/project-manager-gemma"
-    },
-    "software_engineer": {
-        "base_id": "unsloth/gemma-3-1b-it",
-        "adapter_id": "spandana30/project-manager-gemma"
-    },
-    "qa_engineer": {
-        "base_id": "unsloth/gemma-3-1b-it",
-        "adapter_id": "spandana30/project-manager-gemma"
-    }
-}
-@st.cache_resource
-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]]
     html: str
-    refined_request: str
-    final_prompt: str
     feedback: str
     iteration: int
     done: bool
     timings: Dict[str, float]
-def agent(template: str, state: AgentState, agent_key: str, timing_label: str):
-    st.write(f'🛠 Running agent: {agent_key}')
     start = time.time()
-    model, tokenizer = load_agent_model(**AGENT_MODEL_CONFIG[agent_key])
-    latest_input = (
-        state.get("final_prompt")
-        or state.get("refined_request")
-        or state["messages"][-1]["content"]
-    )
-    prompt = template.format(user_input=latest_input, html=state.get("html", ""), final_prompt=state.get("final_prompt", ""))
-    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 are a Product Manager. Your role is to gather and prioritize user requirements "
-        "to ensure the project aligns with business goals. "
-        "Take the user's input and rewrite it as a clear, well-scoped requirement. "
-        "**Do not generate any code or implementation details.**\n\n"
-        "User input:\n{user_input}"
-    ),
-    "project_manager": (
-        "You are a Project Manager. Based on the refined requirement below, outline a realistic and actionable "
-        "design specification. Include necessary UI components and their functionality, but **do not generate code**. "
-        "Your goal is to organize the work in a way that supports efficient development.\n\n"
-        "Refined requirement:\n{user_input}"
-    ),
-    "software_engineer": (
-        "You are a Software Engineer. Based on the design specification below, generate a complete HTML page "
-        "with embedded CSS. The output should include a DOCTYPE declaration, semantic structure, and responsive layout. "
-        "**Only output the HTML code.**\n\nDesign spec:\n{final_prompt}"
-    ),
-    "qa_engineer": (
-        "You are a QA Engineer. Your responsibility is to test and evaluate the HTML below for correctness, responsiveness, "
-        "and quality. Provide clear feedback on any issues you find. If everything looks perfect, respond with 'APPROVED'.\n\n"
-        "HTML code:\n{html}"
-    )
-}
-def generate_ui(user_prompt: str, max_iter: int):
-    state: AgentState = {
-        "messages": [{"role": "user", "content": user_prompt}],
-        "html": "",
-        "refined_request": "",
-        "final_prompt": "",
-        "feedback": "",
-        "iteration": 0,
-        "done": False,
-        "timings": {}
-    }
     workflow = StateGraph(AgentState)
-    workflow.add_node("product_manager", lambda s: {
-        "messages": s["messages"] + [{
-            "role": "product_manager",
-            "content": (pm := agent(PROMPTS["product_manager"], s, "product_manager", "product_manager"))
-        }],
-        "refined_request": pm
-    })
-    workflow.add_node("project_manager", lambda s: {
-        "messages": s["messages"] + [{
-            "role": "project_manager",
-            "content": (pr := agent(PROMPTS["project_manager"], s, "project_manager", "project_manager"))
-        }],
-        "final_prompt": pr
-    })
-    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", "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()
     final_state = app.invoke(state)
-    return final_state
 def main():
-    st.set_page_config(page_title="Multi-Agent UI Generator", layout="wide")
-    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"):
         with st.spinner("Agents working..."):
-            final = generate_ui(prompt, max_iter)
-            st.success("✅ UI Generated")
-            st.write("🧠 Final state:", final)
-            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()

 import streamlit as st
 import os
 import time
+import base64
 from typing import Dict, List, TypedDict
 from langgraph.graph import StateGraph, END
+from huggingface_hub import InferenceClient
+client = InferenceClient(
+    model="mistralai/Mistral-7B-Instruct-v0.2",
+    token=st.secrets["HF_TOKEN"]
+)
 class AgentState(TypedDict):
     messages: List[Dict[str, 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. Given the user request:
+{user_request}
+Break it down into clear features and priorities."""
+PROJECT_MANAGER_PROMPT = """You're a project manager. Based on these features:
+{features}
+Draft a quick development plan with key tasks and timeline."""
+ARCHITECT_PROMPT = """You're a software architect. Create design specs for:
+{user_request}
+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 call_model(prompt: str, max_retries=3) -> str:
+    for attempt in range(max_retries):
+        try:
+            return client.text_generation(
+                prompt,
+                max_new_tokens=3000,
+                temperature=0.3,
+                return_full_text=False
+            )
+        except Exception as e:
+            st.error(f"Model call failed (attempt {attempt+1}): {e}")
+            time.sleep(2)
+    return "<html><body><h1>Error generating UI</h1></body></html>"
+def time_agent(agent_func, state: AgentState, label: str):
     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):
+    features = call_model(PRODUCT_MANAGER_PROMPT.format(user_request=state["messages"][-1]["content"]))
+    return {"messages": state["messages"] + [{"role": "product_manager", "content": features}]}
+def project_manager_agent(state: AgentState):
+    features_msg = next((m["content"] for m in state["messages"] if m["role"] == "product_manager"), "")
+    plan = call_model(PROJECT_MANAGER_PROMPT.format(features=features_msg))
+    return {"messages": state["messages"] + [{"role": "project_manager", "content": plan}]}
+def software_architect_agent(state: AgentState):
+    specs = call_model(ARCHITECT_PROMPT.format(user_request=state["messages"][-1]["content"]))
+    return {"design_specs": specs, "messages": state["messages"] + [{"role": "software_architect", "content": specs}]}
+def engineer_agent(state: AgentState):
+    html = call_model(ENGINEER_PROMPT.format(design_specs=state["design_specs"]))
+    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 = call_model(QA_PROMPT.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}], "design_specs": "", "html": "",
+             "css": "", "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("software_architect", lambda s: time_agent(software_architect_agent, s, "software_architect"))
+    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_edge("product_manager", "project_manager")
+    workflow.add_edge("project_manager", "software_architect")
+    workflow.add_edge("software_architect", "software_engineer")
+    workflow.add_edge("software_engineer", "qa")
+    workflow.add_conditional_edges("qa", 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
 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"):
         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", "software_architect", "software_engineer", "qa"]:
+                st.write(f"🤖 {stage.title().replace('_', ' ')} Time: {final_state['timings'].get(stage, 0):.2f}s")
 if __name__ == "__main__":
+    main()