Spaces:
Build error
Build error
| 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}): {str(e)}") | |
| st.write("\n\n**Full Error:**", e) | |
| st.stop() # Force exit to show error | |
| 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) | |
| st.write("\n🔐 HF_TOKEN found:", st.secrets.get("HF_TOKEN", "❌ Missing!")) | |
| 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() | |