Spaces:

spandana30
/

Multi_agent

Build error

App Files Files Community

spandana30 commited on May 8

Commit

2834a6c

verified ·

1 Parent(s): ed6f467

Update app.py

Browse files

Files changed (1) hide show

app.py +54 -48

app.py CHANGED Viewed

@@ -4,16 +4,19 @@ import time
 import base64
 from typing import Dict, List, TypedDict
 from langgraph.graph import StateGraph, END
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-# Load CodeLLaMA locally
-model_id = "codellama/CodeLlama-7b-hf"
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
 class AgentState(TypedDict):
     messages: List[Dict[str, str]]
     design_specs: str
     html: str
     css: str
@@ -22,16 +25,16 @@ class AgentState(TypedDict):
     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
@@ -55,16 +58,7 @@ 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:
-    try:
-        outputs = generator(prompt, max_new_tokens=1000, temperature=0.3)
-        return outputs[0]["generated_text"]
-    except Exception as e:
-        st.error(f"Local model call failed: {str(e)}")
-        st.stop()
-    return "<html><body><h1>Error generating UI</h1></body></html>"
 def time_agent(agent_func, state: AgentState, label: str):
     start = time.time()
@@ -75,46 +69,58 @@ def time_agent(agent_func, state: AgentState, label: str):
     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")
@@ -138,11 +144,11 @@ def main():
             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()
@@ -154,15 +160,15 @@ def main():
             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()

 import base64
 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"])
 class AgentState(TypedDict):
     messages: List[Dict[str, str]]
+    product_vision: str
+    project_plan: str
     design_specs: str
     html: str
     css: str
     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
 1. Visual quality
 2. Responsiveness
 3. Functionality
+Reply \"APPROVED\" if perfect, or suggest improvements."""
 def time_agent(agent_func, state: AgentState, label: str):
     start = time.time()
     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
+    )
+    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
+    )
+    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": {}}
     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_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.set_entry_point("product_manager")
             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()
             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")
 if __name__ == "__main__":
+    main()