Update agents/graph_builder_langgraph.py

agents/graph_builder_langgraph.py

@@ -1,109 +1,109 @@
 import os
+import uuid
+from typing import List, Dict, Optional
+from pydantic import BaseModel
 from langgraph.graph import StateGraph
 from .document_parser import parse_documents
 from .requirements_extractor import extract_requirements
 from .ui_generator import generate_ui_html
 
+
 # ==========================================================
-# 🧩 Agent Node 1: Document Parser
+# 🧱 Define the State Schema (Pydantic model)
 # ==========================================================
-async def node_parse_docs(state, config=None):
-    """Extract text content from uploaded documents."""
-    files = state.get("files", [])
-    if not files:
-        raise ValueError("No input files provided to document parser.")
+class PipelineState(BaseModel):
+    files: Optional[List[str]] = None
+    text: Optional[str] = None
+    requirements: Optional[Dict] = None
+    html: Optional[str] = None
+    public_url: Optional[str] = None
+
 
+# ==========================================================
+# 📄 Node 1 — Document Parser
+# ==========================================================
+async def node_parse_docs(state: PipelineState, config=None) -> Dict:
     print("📄 [Document Parser] Extracting text...")
-    text = await parse_documents(files)
 
-    if not text or not text.strip():
+    if not state.files:
+        raise ValueError("No input files provided to document parser.")
+
+    text = await parse_documents(state.files)
+    if not text.strip():
         raise ValueError("Document parser returned empty text.")
 
-    # Return as state update
     return {"text": text}
 
 
 # ==========================================================
-# 🧠 Agent Node 2: Requirements Extractor
+# 🧠 Node 2 — Requirements Extractor
 # ==========================================================
-async def node_requirements(state, config=None):
-    """Extract structured business requirements from text."""
+async def node_requirements(state: PipelineState, config=None) -> Dict:
     print("🧩 [Requirements Agent] Extracting structured info...")
 
-    text = state.get("text")
-    if not text:
-        raise ValueError("Missing input text for requirements extraction.")
-
-    requirements = await extract_requirements(text)
+    if not state.text:
+        raise ValueError("Missing 'text' field in state for requirement extraction.")
 
+    requirements = await extract_requirements(state.text)
     if not requirements:
-        raise ValueError("Requirements extraction failed or returned empty.")
+        raise ValueError("Requirements extraction returned empty.")
 
     return {"requirements": requirements}
 
 
 # ==========================================================
-# 🎨 Agent Node 3: UI Generator
+# 🎨 Node 3 — UI Generator
 # ==========================================================
-async def node_ui_generator(state, config=None):
-    """Generate HTML UI using Gemini (PowerApps-style layout)."""
+async def node_ui_generator(state: PipelineState, config=None) -> Dict:
     print("🎨 [UI Generator] Creating PowerApps-style interface...")
 
-    requirements = state.get("requirements")
-    if not requirements:
-        raise ValueError("Missing requirements for UI generation.")
+    if not state.requirements:
+        raise ValueError("Missing 'requirements' in state for UI generation.")
 
     # Load reference design
-    reference_html_path = os.path.join("templates", "demo_qms_design.html")
-    if not os.path.exists(reference_html_path):
-        raise FileNotFoundError(f"Reference HTML not found at {reference_html_path}")
+    ref_path = os.path.join("templates", "demo_qms_design.html")
+    if not os.path.exists(ref_path):
+        raise FileNotFoundError(f"Reference HTML not found at {ref_path}")
 
-    with open(reference_html_path, "r", encoding="utf-8") as f:
+    with open(ref_path, "r", encoding="utf-8") as f:
         reference_html = f.read()
 
-    ui = await generate_ui_html(requirements, reference_html)
+    ui = await generate_ui_html(state.requirements, reference_html)
     html = ui.get("html", "")
 
     if not html or "<html" not in html:
         raise ValueError("Generated HTML invalid or empty.")
 
-    # Save to static/outputs
+    # Save to static/outputs/
     os.makedirs("static/outputs", exist_ok=True)
-    import uuid
     session_id = uuid.uuid4().hex
-    html_filename = f"{session_id}.html"
-    html_path = os.path.join("static/outputs", html_filename)
+    html_path = os.path.join("static/outputs", f"{session_id}.html")
     with open(html_path, "w", encoding="utf-8") as f:
         f.write(html)
 
-    # Construct public link
     base_url = os.environ.get("BASE_PUBLIC_URL", "").rstrip("/")
     if not base_url:
         base_url = "http://localhost:7860"
-    public_url = f"{base_url}/static/outputs/{html_filename}"
 
-    print(f"✅ [UI Generator] HTML generated at: {public_url}")
+    public_url = f"{base_url}/static/outputs/{session_id}.html"
+
+    print(f"✅ [UI Generator] Generated UI available at: {public_url}")
 
-    return {
-        "html": html,
-        "public_url": public_url
-    }
+    return {"html": html, "public_url": public_url}
 
 
 # ==========================================================
-# 🔗 LangGraph Flow Definition
+# 🔗 Build the LangGraph
 # ==========================================================
-graph = StateGraph()
+graph = StateGraph(PipelineState)
 
-# Add nodes (agents)
 graph.add_node("parse_docs", node_parse_docs)
 graph.add_node("requirements", node_requirements)
 graph.add_node("ui_generator", node_ui_generator)
 
-# Define edges (data flow)
 graph.set_entry_point("parse_docs")
 graph.add_edge("parse_docs", "requirements")
 graph.add_edge("requirements", "ui_generator")
 graph.set_finish_point("ui_generator")
 
-print("✅ LangGraph pipeline initialized: [Document Parser → Requirements → UI Generator]")
+print("✅ LangGraph pipeline loaded successfully: Document → Requirements → UI Generator")