Update graph.py

graph.py

@@ -411,61 +411,78 @@ def run_intent_agent(state: AgentState):
 
 def run_pm_agent(state: AgentState):
     log.info("--- PM ---")
-    current_cycles = ensure_int(state, 'rework_cycles', 0) + 1
-    max_loops_val = ensure_int(state, 'max_loops', 0)
+    # Ensure keys
+    current_rework = ensure_int(state, 'rework_cycles', 0)
+    max_loops_val = ensure_int(state, 'max_loops', INITIAL_MAX_REWORK_CYCLES)
+
+    # If we've exhausted loops, short-circuit and produce fallback plan with a note
+    if current_rework > max_loops_val:
+        path = ensure_list(state, 'execution_path') + ["PM"]
+        fallback_plan = {
+            "plan_steps": ["Rework limit exceeded. Manual review required."],
+            "experiment_needed": False,
+            "experiment_type": "word",
+            "experiment_goal": state.get('coreObjectivePrompt', state.get('userInput',''))
+        }
+        return {"pmPlan": fallback_plan, "execution_path": path, "rework_cycles": current_rework, "status_update": "Rework limit hit - manual review"}
+
+    # Normal behavior: increment rework count for this pass
+    current_cycles = current_rework + 1
     path = ensure_list(state, 'execution_path') + ["PM"]
-    
+
+    # (rest of your original PM prompt & parse flow, but ensure the output sets rework_cycles and max_loops)
+    # --- build full_context like before ---
     context_parts = [
         f"=== USER REQUEST ===\n{state.get('userInput', '')}",
         f"\n=== OBJECTIVE ===\n{state.get('coreObjectivePrompt', '')}",
         f"\n=== MEMORY ===\n{state.get('retrievedMemory', 'None')}",
     ]
-    
     if state.get('qaFeedback'):
         context_parts.append(f"\n=== QA FEEDBACK (MUST FIX) ===\n{state.get('qaFeedback')}")
         context_parts.append(f"\n=== PREVIOUS PLAN ===\n{json.dumps(state.get('pmPlan', {}).get('plan_steps', []), indent=2)}")
-    
     full_context = "\n".join(context_parts)
-    
+
     prompt = f"""Create DETAILED, EXECUTABLE plan.
 
 {full_context}
 
-Each step must be SPECIFIC and ACTIONABLE:
-- State EXACTLY what will be created/analyzed
-- Specify WHAT data/information will be used
-- Define WHAT methods will be applied
-
-JSON format:
+Return JSON with:
 {{
-    "plan_steps": ["Specific step 1...", "Specific step 2..."],
+    "plan_steps": [...],
     "experiment_needed": true/false,
     "experiment_type": "notebook|script|excel|word|pdf|repo",
-    "experiment_goal": "Detailed artifact description",
-    "key_requirements": ["Critical requirements"]
+    "experiment_goal": "...",
+    "key_requirements": [...]
 }}
+Be concrete.
+"""
+    try:
+        response = llm.invoke(prompt)
+        plan = parse_json_from_llm(getattr(response, "content", "") or "")
+    except Exception as e:
+        log.warning("PM LLM failed: %s", e)
+        plan = None
 
-Be specific about using uploaded files, implementing algorithms, creating schemas."""
-    
-    response = llm.invoke(prompt)
-    plan = parse_json_from_llm(getattr(response, "content", "") or "")
-    
     if not plan:
         detection = detect_requested_output_types(state.get('userInput', ''))
         plan = {
             "plan_steps": ["Analyze request", "Process information", "Create deliverable", "Review"],
             "experiment_needed": detection.get('requires_artifact', False),
             "experiment_type": detection.get('artifact_type', 'word'),
-            "experiment_goal": state.get('coreObjectivePrompt', state.get('userInput', ''))
+            "experiment_goal": state.get('coreObjectivePrompt', state.get('userInput', '')),
+            "key_requirements": []
         }
-    
+
+    # Normalize experiment_type
     exp_type = normalize_experiment_type(plan.get('experiment_type'), plan.get('experiment_goal',''))
     plan['experiment_type'] = exp_type
-    
     if plan.get('experiment_needed') and not plan.get('experiment_goal'):
         plan['experiment_goal'] = state.get('userInput','')
-    
-    return {"pmPlan": plan, "execution_path": path, "rework_cycles": current_cycles, "status_update": f"Plan created ({len(plan.get('plan_steps', []))} steps)"}
+
+    # Attach loop control info
+    plan['max_loops_initial'] = max_loops_val
+    plan['estimated_cost_usd'] = plan.get('estimated_cost_usd', plan.get('estimated_cost_usd', 0.0))
+    return {"pmPlan": plan, "execution_path": path, "rework_cycles": current_cycles, "max_loops": max_loops_val, "status_update": f"Plan created ({len(plan.get('plan_steps', []))} steps)"}
 
 def _extract_code_blocks(text: str, lang_hint: Optional[str]=None) -> List[str]:
     if lang_hint and "python" in (lang_hint or "").lower():