""" LangGraph workflow for form generation and refine. Mirrors patterns from resume_workflow and visualization_workflow. """ from typing import TypedDict, List, Dict, Any, Optional, Literal import json import logging import time import traceback import os logger = logging.getLogger(__name__) from langgraph.graph import StateGraph, END from langchain_openai import ChatOpenAI from prompts import GENERATE_SYSTEM, REFINE_SYSTEM, PLAN_QUESTIONS_SYSTEM from openai_client import extract_json_from_response from .docs_client import fetch_docs from .context_pack_builder import detect_intent, build_context_packs class PlanState(TypedDict): """State for clarification planning (0-5 questions).""" description: Optional[str] current_fields: Optional[List[Dict[str, Any]]] current_title: Optional[str] should_ask_questions: bool questions: List[Dict[str, Any]] reasoning_summary: Optional[str] error: Optional[str] success: bool class FormGeneratorState(TypedDict, total=False): """State for form generator workflow.""" # Mode: which path to take mode: Literal["generate", "refine"] # Inputs for generate description: Optional[str] # Inputs for refine user_request: Optional[str] current_fields: Optional[List[Dict[str, Any]]] current_title: Optional[str] # Optional context hints (from chat) conversation_context: Optional[List[str]] user_goals: Optional[List[str]] preferred_field_types: Optional[List[str]] must_have_logic: Optional[bool] # Results form_json: Optional[Dict[str, Any]] diff: Optional[Dict[str, Any]] metadata: Optional[Dict[str, Any]] # Status error: Optional[str] success: bool processing_time: float usage: Optional[Dict[str, int]] processing_complete: bool class FormGeneratorWorkflow: """LangGraph workflow for form generation and surgical refine.""" def __init__(self): self.workflow = self._build_workflow() def _build_workflow(self) -> StateGraph: workflow = StateGraph(FormGeneratorState) workflow.add_node("validate_input", self.validate_input_node) workflow.add_node("generate_form_node", self.generate_form_node) workflow.add_node("refine_form_node", self.refine_form_node) workflow.add_node("finalize_results", self.finalize_results_node) workflow.add_conditional_edges( "validate_input", self._route_by_mode, {"generate": "generate_form_node", "refine": "refine_form_node", "error": "finalize_results"}, ) workflow.add_edge("generate_form_node", "finalize_results") workflow.add_edge("refine_form_node", "finalize_results") workflow.add_edge("finalize_results", END) workflow.set_entry_point("validate_input") return workflow.compile() def _route_by_mode(self, state: FormGeneratorState) -> str: if state.get("error"): return "error" return state.get("mode", "error") def validate_input_node(self, state: FormGeneratorState) -> FormGeneratorState: """Validate input and set mode (generate vs refine).""" try: if state.get("description") and str(state["description"]).strip(): state["mode"] = "generate" state["success"] = True return state if ( state.get("user_request") and str(state["user_request"]).strip() and state.get("current_fields") is not None ): state["mode"] = "refine" state["success"] = True return state state["error"] = "Missing input: provide description (generate) or user_request + current_fields (refine)" state["success"] = False except Exception as e: state["error"] = f"Validation error: {str(e)}" state["success"] = False return state def generate_form_node(self, state: FormGeneratorState) -> FormGeneratorState: """Generate full form JSON from description using LLM.""" try: description = (state.get("description") or "").strip() user_content = f"USER REQUEST:\n{description}\n\nGenerate the form JSON only (no markdown)." docs = fetch_docs() if docs: intent = detect_intent( description=description, conversation_context=state.get("conversation_context"), ) if state.get("must_have_logic"): intent["needs_conditional_logic"] = True context_packs = build_context_packs(docs, intent, log_usage=True) if context_packs: logger.info("generate_form using dynamic context intent=%s", intent) user_content = ( f"USER REQUEST:\n{description}\n\n" f"REFERENCE (use for field types, conditional logic, document-extraction):\n{context_packs}\n\n" "Generate the form JSON only (no markdown)." ) else: logger.debug("generate_form using static prompt only (no docs)") client = ChatOpenAI( api_key=os.getenv("OPENAI_API_KEY"), model="gpt-4.1-nano", temperature=0.3, max_tokens=4000, model_kwargs={"response_format": {"type": "json_object"}}, ) start = time.time() response = client.invoke( [{"role": "system", "content": GENERATE_SYSTEM}, {"role": "user", "content": user_content}] ) elapsed = time.time() - start raw = response.content if hasattr(response, "content") else str(response) json_str = extract_json_from_response(raw) data = json.loads(json_str) form_json, metadata = self._validate_form_json(data) state["form_json"] = form_json state["metadata"] = metadata state["processing_time"] = elapsed state["success"] = True if hasattr(response, "response_metadata") and response.response_metadata.get("usage"): usage = response.response_metadata["usage"] state["usage"] = { "input_tokens": usage.get("input_tokens", 0), "output_tokens": usage.get("output_tokens", 0), "total_tokens": usage.get("total_tokens", 0), } except json.JSONDecodeError as e: state["error"] = f"Generated form JSON is invalid: {e}" state["success"] = False except Exception as e: state["error"] = str(e) state["success"] = False traceback.print_exc() return state def refine_form_node(self, state: FormGeneratorState) -> FormGeneratorState: """Produce surgical diff (changes, additions, removals) using LLM.""" try: user_request = (state.get("user_request") or "").strip() current_fields = state.get("current_fields") or [] current_title = state.get("current_title") or "Untitled" fields_json = json.dumps(current_fields, indent=2) max_fields_chars = 12000 if len(fields_json) > max_fields_chars: fields_json = fields_json[: max_fields_chars - 3] + "..." user_content = ( f"Current form title: {current_title}\n\n" f"Current fields (JSON):\n{fields_json}\n\n" f"User request: {user_request}\n\n" "Return ONLY the diff JSON (changes, additions, removals). No markdown." ) docs = fetch_docs() if docs: intent = detect_intent( user_request=user_request, current_fields=current_fields, conversation_context=state.get("conversation_context"), ) context_packs = build_context_packs(docs, intent, max_total_chars=4000, log_usage=True) if context_packs: logger.info("refine_form using dynamic context intent=%s", intent) user_content = ( f"Current form title: {current_title}\n\n" f"Current fields (JSON):\n{fields_json}\n\n" f"REFERENCE (for new/changed fields):\n{context_packs}\n\n" f"User request: {user_request}\n\n" "Return ONLY the diff JSON (changes, additions, removals). No markdown." ) else: logger.debug("refine_form using static prompt only (no docs)") client = ChatOpenAI( api_key=os.getenv("OPENAI_API_KEY"), model="gpt-4.1-nano", temperature=0.3, max_tokens=4000, model_kwargs={"response_format": {"type": "json_object"}}, ) start = time.time() response = client.invoke( [{"role": "system", "content": REFINE_SYSTEM}, {"role": "user", "content": user_content}] ) elapsed = time.time() - start raw = response.content if hasattr(response, "content") else str(response) json_str = extract_json_from_response(raw) data = json.loads(json_str) changes = data.get("changes") additions = data.get("additions") removals = data.get("removals") if not isinstance(changes, list): changes = [] if not isinstance(additions, list): additions = [] if not isinstance(removals, list): removals = [] state["diff"] = {"changes": changes, "additions": additions, "removals": removals} state["processing_time"] = elapsed state["success"] = True if hasattr(response, "response_metadata") and response.response_metadata.get("usage"): usage = response.response_metadata["usage"] state["usage"] = { "input_tokens": usage.get("input_tokens", 0), "output_tokens": usage.get("output_tokens", 0), "total_tokens": usage.get("total_tokens", 0), } except json.JSONDecodeError as e: state["error"] = f"Refine diff JSON is invalid: {e}" state["success"] = False except Exception as e: state["error"] = str(e) state["success"] = False traceback.print_exc() return state def finalize_results_node(self, state: FormGeneratorState) -> FormGeneratorState: """Mark processing complete.""" state["processing_complete"] = True return state @staticmethod def _validate_form_json(data: Dict[str, Any]) -> tuple: """Ensure form has title and fields; return (form_json, metadata).""" if not isinstance(data, dict): raise ValueError("Form must be a JSON object") title = data.get("title") or "Untitled Form" fields = data.get("fields") if not isinstance(fields, list): raise ValueError("Form must have a 'fields' array") form_json = {"title": title, "fields": fields} if data.get("description") is not None: form_json["description"] = data["description"] field_types = list({f.get("type") for f in fields if isinstance(f, dict) and f.get("type")}) metadata = { "fieldCount": len(fields), "fieldTypes": field_types, "hasConditionalLogic": any( (f.get("conditionalLogic") or {}).get("enabled") for f in fields if isinstance(f, dict) ), "hasAIFields": any(f.get("type") == "document-extraction" for f in fields if isinstance(f, dict)), } return form_json, metadata def generate_form( self, description: str, conversation_context: Optional[List[str]] = None, user_goals: Optional[List[str]] = None, preferred_field_types: Optional[List[str]] = None, must_have_logic: Optional[bool] = None, ) -> Dict[str, Any]: """Run workflow in generate mode and return API-shaped result.""" initial = FormGeneratorState( mode="generate", description=description.strip(), user_request=None, current_fields=None, current_title=None, conversation_context=conversation_context, user_goals=user_goals, preferred_field_types=preferred_field_types, must_have_logic=must_have_logic, form_json=None, diff=None, metadata=None, error=None, success=False, processing_time=0.0, usage=None, processing_complete=False, ) try: final = self.workflow.invoke(initial) return { "formJSON": final.get("form_json"), "metadata": final.get("metadata") or {}, "warnings": None, "success": final.get("success", False), "error": final.get("error"), "processing_time": final.get("processing_time", 0.0), "usage": final.get("usage"), } except Exception as e: return { "formJSON": None, "metadata": {}, "warnings": None, "success": False, "error": str(e), "processing_time": 0.0, "usage": None, } def refine_form( self, user_request: str, current_fields: List[Dict[str, Any]], current_title: str, conversation_context: Optional[List[str]] = None, ) -> Dict[str, Any]: """Run workflow in refine mode and return API-shaped result.""" initial = FormGeneratorState( mode="refine", description=None, user_request=user_request.strip(), current_fields=current_fields, current_title=current_title or "", conversation_context=conversation_context, form_json=None, diff=None, metadata=None, error=None, success=False, processing_time=0.0, usage=None, processing_complete=False, ) try: final = self.workflow.invoke(initial) return { "diff": final.get("diff") or {"changes": [], "additions": [], "removals": []}, "success": final.get("success", False), "error": final.get("error"), "processing_time": final.get("processing_time", 0.0), "usage": final.get("usage"), } except Exception as e: return { "diff": {"changes": [], "additions": [], "removals": []}, "success": False, "error": str(e), "processing_time": 0.0, "usage": None, } def _build_plan_workflow() -> StateGraph: """Build small graph for plan-form-request (clarification questions).""" workflow = StateGraph(PlanState) workflow.add_node("plan_questions", _plan_questions_node) workflow.add_edge("plan_questions", END) workflow.set_entry_point("plan_questions") return workflow.compile() def _plan_questions_node(state: PlanState) -> PlanState: """Single node: call LLM to decide if we need clarifying questions.""" try: description = (state.get("description") or "").strip() current_fields = state.get("current_fields") current_title = state.get("current_title") or "" client = ChatOpenAI( api_key=os.getenv("OPENAI_API_KEY"), model="gpt-4.1-nano", temperature=0.2, max_tokens=1500, model_kwargs={"response_format": {"type": "json_object"}}, ) user_content = f"User request: {description}" if current_fields or current_title: user_content += f"\nCurrent form title: {current_title}\nCurrent fields (count): {len(current_fields or [])}" user_content += "\n\nReturn the JSON only (should_ask_questions, reasoning_summary, questions)." response = client.invoke( [{"role": "system", "content": PLAN_QUESTIONS_SYSTEM}, {"role": "user", "content": user_content}] ) raw = response.content if hasattr(response, "content") else str(response) json_str = extract_json_from_response(raw) data = json.loads(json_str) state["should_ask_questions"] = bool(data.get("should_ask_questions", False)) state["reasoning_summary"] = str(data.get("reasoning_summary", "")).strip() or "Ready." qs = data.get("questions") if isinstance(qs, list): state["questions"] = [ { "id": q.get("id", f"q{i}"), "label": q.get("label", ""), "type": q.get("type", "text"), "required": bool(q.get("required", True)), "options": q.get("options") if isinstance(q.get("options"), list) else None, } for i, q in enumerate(qs) if isinstance(q, dict) ][:5] else: state["questions"] = [] state["success"] = True except Exception as e: state["error"] = str(e) state["success"] = False state["should_ask_questions"] = False state["questions"] = [] state["reasoning_summary"] = "" traceback.print_exc() return state plan_workflow = _build_plan_workflow() def plan_form_request( description: str, current_fields: Optional[List[Dict[str, Any]]] = None, current_title: str = "", ) -> Dict[str, Any]: """Run plan workflow and return should_ask_questions, questions, reasoning_summary.""" initial: PlanState = { "description": description.strip(), "current_fields": current_fields, "current_title": current_title or "", "should_ask_questions": False, "questions": [], "reasoning_summary": "", "error": None, "success": False, } try: final = plan_workflow.invoke(initial) return { "should_ask_questions": final.get("should_ask_questions", False), "questions": final.get("questions", []), "reasoning_summary": final.get("reasoning_summary", "") or "Ready.", "success": final.get("success", False), "error": final.get("error"), } except Exception as e: return { "should_ask_questions": False, "questions": [], "reasoning_summary": "", "success": False, "error": str(e), } workflow_instance = FormGeneratorWorkflow()