Create agents/graph_builder_langgraph.py

agents/graph_builder_langgraph.py

@@ -0,0 +1,118 @@
+from typing import List, Dict, Any, TypedDict, Optional
+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 .ui_qa import run_ui_qa
+from .ui_repair import repair_ui_html
+from .report_generator import generate_pdf_report
+import os
+import uuid
+import asyncio
+
+# ----- shared state -----
+class State(TypedDict, total=False):
+    file_paths: List[str]
+    texts: str
+    files_meta: List[Dict[str, Any]]
+    requirements: Dict[str, Any]
+    html: str
+    qa: Dict[str, Any]
+    attempts: int
+    pdf_path: str
+    public_url: str
+
+# ----- nodes -----
+async def node_parse(state: State, config=None, runtime=None) -> State:
+    res = await parse_documents(state["file_paths"])
+    return {"texts": res["texts"], "files_meta": res["files"]}
+
+async def node_extract(state: State, config=None, runtime=None) -> State:
+    reqs = await extract_requirements(state["texts"])
+    return {"requirements": reqs}
+
+async def node_generate(state: State, config=None, runtime=None) -> State:
+    reference_html = runtime.config.get("reference_html") if runtime else os.environ.get("REFERENCE_HTML", "")
+    if not reference_html and os.path.exists("templates/demo_qms_design.html"):
+        with open("templates/demo_qms_design.html", "r", encoding="utf-8") as f:
+            reference_html = f.read()
+    res = await generate_ui_html(state["requirements"], reference_html)
+    return {"html": res["html"], "attempts": state.get("attempts", 0) + 1}
+
+async def node_qa(state: State, config=None, runtime=None) -> State:
+    qa = await run_ui_qa(state["html"])
+    return {"qa": qa}
+
+async def node_repair(state: State, config=None, runtime=None) -> State:
+    repaired = await repair_ui_html(state["html"], state["qa"]["comment"])
+    return {"html": repaired["html"]}
+
+async def node_report(state: State, config=None, runtime=None) -> State:
+    html = state["html"]
+    session = uuid.uuid4().hex
+    out_dir = os.path.join("static", "outputs")
+    os.makedirs(out_dir, exist_ok=True)
+    html_path = os.path.join(out_dir, f"{session}.html")
+    with open(html_path, "w", encoding="utf-8") as f:
+        f.write(html)
+
+    base_url = state.get("public_url") or os.environ.get("BASE_PUBLIC_URL", "")
+    if base_url:
+        base_url = base_url.rstrip("/")
+        public_url = f"{base_url}/static/outputs/{session}.html"
+    else:
+        public_url = f"/static/outputs/{session}.html"
+
+    pdf_info = await generate_pdf_report(public_url)
+    return {"pdf_path": pdf_info["pdf_path"], "public_url": public_url}
+
+# ----- conditional guard -----
+def qa_guard(state: State) -> str:
+    """Return next node key based on QA score."""
+    score = int(state.get("qa", {}).get("score", 0))
+    attempts = int(state.get("attempts", 0))
+    max_retries = int(os.environ.get("MAX_RETRIES", "2"))
+    if score >= int(os.environ.get("QA_THRESHOLD", "85")) or attempts >= max_retries:
+        return "report"
+    return "repair"
+
+# ----- builder -----
+def build_langgraph_runner(reference_html: Optional[str] = None) -> StateGraph:
+    g = StateGraph(State)
+    g.add_node("parse", node_parse)
+    g.add_node("extract", node_extract)
+    g.add_node("generate", node_generate)
+    g.add_node("qa", node_qa)
+    g.add_node("repair", node_repair)
+    g.add_node("report", node_report)
+
+    g.set_entry_point("parse")
+    g.set_finish_point("report")
+
+    # linear part
+    g.add_edge(START, "parse")
+    g.add_edge("parse", "extract")
+    g.add_edge("extract", "generate")
+    g.add_edge("generate", "qa")
+
+    # explicit conditional edges
+    g.add_conditional_edges(
+        "qa",
+        qa_guard,
+        {"report": "report", "repair": "repair"},
+    )
+
+    # loop
+    g.add_edge("repair", "generate")
+
+    graph = g.compile()
+    if reference_html:
+        graph.config = {"reference_html": reference_html}
+    return graph
+
+# ----- runner -----
+async def run_graph_once(graph, file_paths: List[str], base_public_url: str | None = None) -> Dict[str, Any]:
+    if base_public_url:
+        os.environ["BASE_PUBLIC_URL"] = base_public_url
+    state: State = {"file_paths": file_paths, "attempts": 0}
+    return await graph.ainvoke(state)