Update agents/graph_builder_langgraph.py

agents/graph_builder_langgraph.py

@@ -1,86 +1,44 @@
 import os
-import json
-import uuid
-from typing import Dict, Any
-from langgraph.graph import StateGraph, START, END
 from .document_parser import parse_documents
 from .requirements_extractor import extract_requirements
 from .ui_generator import generate_ui_html
+from langgraph.graph import StateGraph
 
-# ======================================================================
-# 🧠 State Schema
-# ======================================================================
-
-class State(Dict[str, Any]):
-    pass
-
-# ======================================================================
-# 🚀 Individual Agent Nodes
-# ======================================================================
-
-async def node_parse(state: State, config=None, runtime=None) -> State:
-    """Extract plain text content from uploaded documents."""
-    print("📄 [Document Parser] Extracting text...")
+async def node_parse_docs(state, config=None):
+    """First agent node — extracts raw text from uploaded documents."""
     files = state.get("files", [])
-    all_text = await parse_documents(files)
-    return {"text": all_text}
-
+    if not files:
+        raise ValueError("No input files provided to document parser.")
+    print("📄 [Document Parser] Extracting text...")
+    text = await parse_documents(files)
+    if not text:
+        raise ValueError("Document parser returned empty text.")
+    return {"text": text}
 
-async def node_requirements(state: State, config=None, runtime=None) -> State:
-    """Extract structured requirements (entities, workflows, etc.) from text."""
+async def node_requirements(state, config=None):
+    """Second agent — extracts structured business requirements."""
     print("🧩 [Requirements Agent] Extracting structured info...")
-    text = state["text"]
+    text = state["text"]   # <-- this line will now work
     requirements = await extract_requirements(text)
     return {"requirements": requirements}
 
-
-async def node_ui(state: State, config=None, runtime=None) -> State:
-    """Generate PowerApps-style UI HTML from extracted requirements."""
-    print("🎨 [UI Generator] Creating HTML...")
-    requirements = state["requirements"]
-
-    reference_path = os.path.join("templates", "demo_qms_design.html")
-    if not os.path.exists(reference_path):
-        raise FileNotFoundError("Reference design (demo_qms_design.html) not found")
-
-    with open(reference_path, "r", encoding="utf-8") as f:
+async def node_ui_generator(state, config=None):
+    """Third agent — generates final HTML UI."""
+    print("🎨 [UI Generator] Creating PowerApps-style interface...")
+    reqs = state["requirements"]
+    reference_html_path = os.path.join("templates", "demo_qms_design.html")
+    with open(reference_html_path, "r", encoding="utf-8") as f:
         reference_html = f.read()
-
-    result = await generate_ui_html(requirements, reference_html)
-    html = result["html"]
-
-    # Save generated UI file
-    session_id = uuid.uuid4().hex
-    output_dir = os.path.join("static", "outputs")
-    os.makedirs(output_dir, exist_ok=True)
-    output_path = os.path.join(output_dir, f"{session_id}.html")
-
-    with open(output_path, "w", encoding="utf-8") as f:
-        f.write(html)
-
-    public_url = f"/static/outputs/{session_id}.html"
-
-    return {
-        "html": html,
-        "public_url": public_url,
-    }
-
-# ======================================================================
-# 🕸️ Graph Definition
-# ======================================================================
-
-def build_graph():
-    graph = StateGraph(State)
-
-    graph.add_node("Document Parser Agent", node_parse)
-    graph.add_node("Requirements Extraction Agent", node_requirements)
-    graph.add_node("UI Generator Agent", node_ui)
-
-    graph.add_edge(START, "Document Parser Agent")
-    graph.add_edge("Document Parser Agent", "Requirements Extraction Agent")
-    graph.add_edge("Requirements Extraction Agent", "UI Generator Agent")
-    graph.add_edge("UI Generator Agent", END)
-
-    return graph.compile()
-
-graph = build_graph()
+    ui = await generate_ui_html(reqs, reference_html)
+    return {"html": ui["html"]}
+
+# Build pipeline
+graph = StateGraph()
+graph.add_node("parse_docs", node_parse_docs)
+graph.add_node("requirements", node_requirements)
+graph.add_node("ui_generator", node_ui_generator)
+
+graph.set_entry_point("parse_docs")
+graph.add_edge("parse_docs", "requirements")
+graph.add_edge("requirements", "ui_generator")
+graph.set_finish_point("ui_generator")