Spaces:

Rahul-8799
/

MAC

Runtime error

App Files Files Community

Rahul-8799 commited on May 9

Commit

6678303

verified ·

1 Parent(s): 5261f65

Update utils/langgraph_pipeline.py

Browse files

Files changed (1) hide show

utils/langgraph_pipeline.py +84 -89

utils/langgraph_pipeline.py CHANGED Viewed

@@ -13,12 +13,9 @@ from agents import (
     project_manager_agent,
     software_architect_agent,
     software_engineer_agent,
-    quality_assurance_agent,
 )
-# ——————————————
-# 1) State shapes
-# ——————————————
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: List[Dict[str, Any]]
@@ -28,137 +25,135 @@ class OutputState(TypedDict):
     proj_output: str
     arch_output: str
     dev_output: str
-    qa_output: str
     chat_log: List[Dict[str, Any]]
-# ——————————————
-# 2) Wrap each agent so it sees the full conversation history
-# ——————————————
 def wrap_agent(agent_run, output_key: str):
     def node(state: Dict[str, Any]) -> Dict[str, Any]:
         history = state["messages"]
         log     = state["chat_log"]
         result  = agent_run({"messages": history, "chat_log": log})
-        # append any AIMessage(s) that the agent returned
         new_history = history + result["messages"]
         return {
             "messages": new_history,
-            "chat_log":  result["chat_log"],
-            output_key:  result[output_key],
         }
     return node
-# ——————————————
-# 3) Bridge to seed the first AIMessage for ProductManager
-# ——————————————
 def bridge_to_pm(state: Dict[str, Any]) -> Dict[str, Any]:
     history = state["messages"]
     log     = state["chat_log"]
     if not history or not isinstance(history[-1], HumanMessage):
-        raise ValueError("bridge_to_pm expected last message to be HumanMessage")
     prompt = history[-1].content
     spec_prompt = (
         f"# Stakeholder Prompt\n\n"
         f"\"{prompt}\"\n\n"
-        "Please generate a structured product specification including:\n"
         "- Goals\n"
         "- Key features\n"
         "- User stories\n"
         "- Success metrics\n"
     )
-    ai = AIMessage(content=spec_prompt)
     return {
-        "messages": [ai],
         "chat_log": log + [{"role": "System", "content": spec_prompt}],
     }
-# ——————————————
-# 4) Build & compile the LangGraph
-# ——————————————
 graph = StateGraph(input=InputState, output=OutputState)
 graph.add_node("BridgePM",        bridge_to_pm)
-graph.add_node("ProductManager",  wrap_agent(product_manager_agent.run,   "pm_output"))
-graph.add_node("ProjectManager",  wrap_agent(project_manager_agent.run,   "proj_output"))
-graph.add_node("SoftwareArchitect", wrap_agent(software_architect_agent.run, "arch_output"))
-graph.add_node("SoftwareEngineer",  wrap_agent(software_engineer_agent.run,  "dev_output"))
-graph.add_node("QualityAssurance",  wrap_agent(quality_assurance_agent.run,  "qa_output"))
 graph.set_entry_point("BridgePM")
-graph.add_edge("BridgePM",         "ProductManager")
-graph.add_edge("ProductManager",   "ProjectManager")
-graph.add_edge("ProjectManager",   "SoftwareArchitect")
-graph.add_edge("SoftwareArchitect","SoftwareEngineer")
-graph.add_edge("SoftwareEngineer", "QualityAssurance")
-graph.add_edge("QualityAssurance", END)
 compiled_graph = graph.compile()
-# ——————————————
-# 5) Run, parse QA spec, generate HTML/CSS, zip, and return
-# ——————————————
 def run_pipeline_and_save(prompt: str) -> Tuple[List[Dict[str, Any]], str]:
-    # a) Invoke the multi-agent pipeline
-    initial_state = {
-        "messages": [HumanMessage(content=prompt)],
-        "chat_log": [],
-    }
-    final_state = compiled_graph.invoke(initial_state)
-    chat_log  = final_state["chat_log"]
-    qa_output = final_state["qa_output"]
-    # b) Parse the QA spec into sections
-    sections = {}
-    for match in re.finditer(r"##\s*(.+?)\n((?:- .+\n?)+)", qa_output):
-        title = match.group(1).strip()
-        items = [line.strip()[2:].strip() for line in match.group(2).splitlines() if line.startswith("- ")]
-        sections[title] = items
-    # c) Build HTML
-    title = prompt.title()
-    html_parts = [
-        "<!DOCTYPE html>",
-        "<html lang='en'><head>",
-        "  <meta charset='UTF-8'>",
-        "  <meta name='viewport' content='width=device-width,initial-scale=1'>",
-        f"  <title>{title}</title>",
-        "  <link rel='stylesheet' href='styles.css'>",
-        "</head><body>",
-        f"  <header><h1>{title}</h1></header>",
-    ]
-    # for each spec section, render as a <section>
-    for sec in ["Goals", "Key features", "User stories", "Success metrics"]:
-        items = sections.get(sec, [])
-        html_parts.append(f"  <section id='{sec.lower().replace(' ', '-')}'><h2>{sec}</h2><ul>")
-        for it in items:
-            html_parts.append(f"    <li>{it}</li>")
-        html_parts.append("  </ul></section>")
-    html_parts.append("</body></html>")
-    html_code = "\n".join(html_parts)
-    # d) Simple CSS
-    css_code = """
-body { font-family: Arial, sans-serif; margin: 2em; line-height: 1.5; }
-header { text-align: center; margin-bottom: 2em; }
-section { margin-bottom: 1.5em; }
-section h2 { color: #333; }
-ul { list-style: disc inside; }
-"""
-    # e) Write files & zip
     site_id  = uuid.uuid4().hex
-    out_dir  = Path("output")
-    site_dir = out_dir / f"site_{site_id}"
     site_dir.mkdir(parents=True, exist_ok=True)
     (site_dir / "index.html").write_text(html_code, encoding="utf-8")
-    (site_dir / "styles.css").write_text(css_code, encoding="utf-8")
-    zip_path = out_dir / f"site_{site_id}.zip"
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
         for f in site_dir.iterdir():
             zf.write(f, arcname=f.name)
-    return chat_log, str(zip_path)

     project_manager_agent,
     software_architect_agent,
     software_engineer_agent,
 )
+# 1) Define the shapes of our state
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: List[Dict[str, Any]]
     proj_output: str
     arch_output: str
     dev_output: str
     chat_log: List[Dict[str, Any]]
+# 2) Helper to wrap any agent so it sees full history
 def wrap_agent(agent_run, output_key: str):
     def node(state: Dict[str, Any]) -> Dict[str, Any]:
         history = state["messages"]
         log     = state["chat_log"]
         result  = agent_run({"messages": history, "chat_log": log})
+        # append the AIMessage(s) returned by the agent
         new_history = history + result["messages"]
         return {
             "messages": new_history,
+            "chat_log": result["chat_log"],
+            output_key: result[output_key],
         }
     return node
+# 3) Seed the chain by turning the user's prompt into a spec for PM
 def bridge_to_pm(state: Dict[str, Any]) -> Dict[str, Any]:
     history = state["messages"]
     log     = state["chat_log"]
+    # last message must be the HumanMessage
     if not history or not isinstance(history[-1], HumanMessage):
+        raise ValueError("bridge_to_pm expects a HumanMessage at the end of history")
     prompt = history[-1].content
     spec_prompt = (
         f"# Stakeholder Prompt\n\n"
         f"\"{prompt}\"\n\n"
+        "Generate a structured product specification including:\n"
         "- Goals\n"
         "- Key features\n"
         "- User stories\n"
         "- Success metrics\n"
     )
     return {
+        "messages": [AIMessage(content=spec_prompt)],
         "chat_log": log + [{"role": "System", "content": spec_prompt}],
     }
+# 4) After architecture, ask the Software Engineer for two files
+def bridge_to_code(state: Dict[str, Any]) -> Dict[str, Any]:
+    history = state["messages"]
+    log     = state["chat_log"]
+    # Ensure the architect has spoken
+    if not history or not isinstance(history[-1], AIMessage):
+        raise ValueError("bridge_to_code expects an AIMessage from the architect")
+    # Grab the original user prompt for titling
+    user_prompt = history[0].content
+    title = user_prompt.title()
+    instruction = (
+        f"Now generate two files for a static website based on the above architecture:\n\n"
+        f"===index.html===\n"
+        f"<!DOCTYPE html>\n"
+        f"<html lang=\"en\">\n"
+        f"<head>\n"
+        f"  <meta charset=\"UTF-8\">\n"
+        f"  <meta name=\"viewport\" content=\"width=device-width,initial-scale=1\">\n"
+        f"  <title>{title}</title>\n"
+        f"  <link rel=\"stylesheet\" href=\"styles.css\">\n"
+        f"</head>\n"
+        f"<body>\n"
+        f"  <!-- Your site UI goes here -->\n"
+        f"</body>\n"
+        f"</html>\n\n"
+        f"===styles.css===\n"
+        f"/* Your CSS styles go here */\n"
+    )
+    return {
+        "messages": history + [AIMessage(content=instruction)],
+        "chat_log": log + [{"role": "System", "content": instruction}],
+    }
+# 5) Build and compile the LangGraph
 graph = StateGraph(input=InputState, output=OutputState)
+# Nodes
 graph.add_node("BridgePM",        bridge_to_pm)
+graph.add_node("ProductManager",  wrap_agent(product_manager_agent.run,  "pm_output"))
+graph.add_node("ProjectManager",  wrap_agent(project_manager_agent.run,  "proj_output"))
+graph.add_node("SoftwareArchitect", wrap_agent(software_architect_agent.run,"arch_output"))
+graph.add_node("BridgeCode",      bridge_to_code)
+graph.add_node("SoftwareEngineer", wrap_agent(software_engineer_agent.run,"dev_output"))
+# Flow
 graph.set_entry_point("BridgePM")
+graph.add_edge("BridgePM",          "ProductManager")
+graph.add_edge("ProductManager",    "ProjectManager")
+graph.add_edge("ProjectManager",    "SoftwareArchitect")
+graph.add_edge("SoftwareArchitect", "BridgeCode")
+graph.add_edge("BridgeCode",        "SoftwareEngineer")
+graph.add_edge("SoftwareEngineer",  END)
 compiled_graph = graph.compile()
+# 6) Run the pipeline, parse out the two files, write & zip them
 def run_pipeline_and_save(prompt: str) -> Tuple[List[Dict[str, Any]], str]:
+    # a) Run through all agents
+    initial_state = {"messages": [HumanMessage(content=prompt)], "chat_log": []}
+    final_state   = compiled_graph.invoke(initial_state)
+    chat_log = final_state["chat_log"]
+    code_blob = final_state["dev_output"]
+    # b) Extract HTML and CSS between our markers
+    html_match = re.search(r"===index\.html===\s*(.*?)\s*===styles\.css===", code_blob, re.S)
+    css_match  = re.search(r"===styles\.css===\s*(.*)", code_blob, re.S)
+    html_code = html_match.group(1).strip() if html_match else ""
+    css_code  = css_match.group(1).strip()  if css_match  else ""
+    # c) Write files to disk
     site_id  = uuid.uuid4().hex
+    site_dir = Path("output") / f"site_{site_id}"
     site_dir.mkdir(parents=True, exist_ok=True)
     (site_dir / "index.html").write_text(html_code, encoding="utf-8")
+    (site_dir / "styles.css").write_text(css_code,  encoding="utf-8")
+    # d) Zip them up
+    zip_path = Path("output") / f"site_{site_id}.zip"
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
         for f in site_dir.iterdir():
             zf.write(f, arcname=f.name)
+    return chat_log, str(zip_path)