Spaces:

Rahul-8799
/

MAC

Runtime error

App Files Files Community

Rahul-8799 commited on May 9

Commit

b60ae2d

verified ·

1 Parent(s): 6678303

Update utils/langgraph_pipeline.py

Browse files

Files changed (1) hide show

utils/langgraph_pipeline.py +100 -84

utils/langgraph_pipeline.py CHANGED Viewed

@@ -1,6 +1,4 @@
-import uuid
-import zipfile
-import re
 from pathlib import Path
 from typing import TypedDict, List, Dict, Any, Tuple
@@ -13,9 +11,12 @@ from agents import (
     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]]
@@ -25,34 +26,33 @@ class OutputState(TypedDict):
     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"
@@ -67,91 +67,107 @@ def bridge_to_pm(state: Dict[str, Any]) -> Dict[str, Any]:
         "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)

+import uuid, zipfile, re
 from pathlib import Path
 from typing import TypedDict, List, Dict, Any, Tuple
     project_manager_agent,
     software_architect_agent,
     software_engineer_agent,
+    quality_assurance_agent,
 )
+# ——————————————
+# 1) State definitions
+# ——————————————
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: List[Dict[str, Any]]
     proj_output: str
     arch_output: str
     dev_output: str
+    qa_output: str
     chat_log: List[Dict[str, Any]]
+# ——————————————
+# 2) Wrap agents so they see 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})
         return {
+            "messages": history + result["messages"],
+            "chat_log":  result["chat_log"],
+            output_key:  result[output_key],
         }
     return node
+# ——————————————
+# 3) Bridge → 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 a HumanMessage at history end")
     prompt = history[-1].content
     spec_prompt = (
         f"# Stakeholder Prompt\n\n"
         f"\"{prompt}\"\n\n"
         "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) Parse spec into sections
+# ——————————————
+def parse_spec(spec: str) -> Dict[str, List[str]]:
+    sections: Dict[str, List[str]] = {}
+    for m in re.finditer(r"##\s*(.+?)\n((?:- .+\n?)+)", spec):
+        name = m.group(1).strip()
+        items = [line[2:].strip() for line in m.group(2).splitlines() if line.startswith("- ")]
+        sections[name] = items
+    return sections
+# ——————————————
+# 6) Run pipeline, generate site, zip, return (chat_log, zip_path)
+# ——————————————
 def run_pipeline_and_save(prompt: str) -> Tuple[List[Dict[str, Any]], str]:
+    # a) invoke agents
     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 spec
+    spec = parse_spec(qa_output)
+    features     = spec.get("Key features", [])
+    testimonials = spec.get("User stories", [])
+    # c) build HTML
+    title = prompt.title()
+    domain = prompt.replace(" ", "").lower() + ".com"
+    cards_html = "\n".join(f"<div class='card'><h3>{f}</h3></div>" for f in features)
+    test_html  = "\n".join(f"<blockquote>{t}</blockquote>" for t in testimonials)
+    html_code = f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width,initial-scale=1">
+  <title>{title}</title>
+  <link rel="stylesheet" href="styles.css">
+</head>
+<body>
+  <header><h1>{title}</h1></header>
+  <section id="features">
+    <h2>Features</h2>
+    <div class="cards">
+      {cards_html}
+    </div>
+  </section>
+  <section id="testimonials">
+    <h2>Testimonials</h2>
+    {test_html or '<p>No testimonials provided.</p>'}
+  </section>
+  <section id="contact">
+    <h2>Contact Us</h2>
+    <p>Email: info@{domain}</p>
+  </section>
+</body>
+</html>"""
+    # d) basic CSS
+    css_code = """
+body { font-family: Arial, sans-serif; margin: 1em; line-height: 1.5; }
+header { text-align: center; margin-bottom: 2em; }
+.cards { display: grid; grid-template-columns: repeat(auto-fit,minmax(150px,1fr)); gap: 1em; }
+.card { background: #f9f9f9; padding: 1em; border-radius: 8px; box-shadow: 0 2px 5px rgba(0,0,0,0.1); text-align: center; }
+blockquote { font-style: italic; margin: 1em; padding: 0.5em; background: #eef; border-left: 4px solid #99f; }
+"""
+    # e) write & 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)