Update graph_upgraded.py

graph_upgraded.py

@@ -505,129 +505,138 @@ def run_knowledge_curator_agent(state: AgentState) -> Dict[str, Any]:
 # --- Wiring / injection into existing main_workflow ---
 
 def apply_upgrades():
-    log.info("Applying graph upgrades (flexible pragmatist/governance)")
+    """
+    Rebuild the main workflow graph with upgraded routing.
+    CRITICAL: Creates a NEW graph instead of modifying the compiled one.
+    """
+    log.info("Applying graph upgrades (rebuilding graph with proper routing)")
+    
     try:
-        mw = getattr(base_graph, "main_workflow", None)
-        if mw is None:
-            raise RuntimeError("base_graph.main_workflow not found. Ensure graph.py created main_workflow before applying upgrades.")
-
-        node_map = {
-            "pragmatist_agent": run_pragmatist_agent,
-            "governance_agent": run_governance_agent,
-            "compliance_agent": run_compliance_agent,
-            "observer_agent": run_observer_agent,
-            "knowledge_curator_agent": run_knowledge_curator_agent
-        }
-
-        # Add nodes defensively
-        for nm, fn in node_map.items():
-            try:
-                mw.add_node(nm, fn)
-                log.info("Added node: %s", nm)
-            except Exception as e:
-                log.debug("Could not add node %s (may already exist): %s", nm, e)
-
-        # Wire edges: pm -> pragmatist -> governance
-        try:
-            mw.add_edge("pm_agent", "pragmatist_agent")
-            log.info("Edge: pm_agent -> pragmatist_agent")
-        except Exception as e:
-            log.debug("pm_agent -> pragmatist_agent: %s", e)
-
-        try:
-            mw.add_edge("pragmatist_agent", "governance_agent")
-            log.info("Edge: pragmatist_agent -> governance_agent")
-        except Exception as e:
-            log.debug("pragmatist_agent -> governance_agent: %s", e)
-
-        # Conditional edges from governance
+        from langgraph.graph import StateGraph, END
+        
+        # Create BRAND NEW graph
+        new_workflow = StateGraph(AgentState)
+        
+        # Add all nodes (reusing original functions where appropriate)
+        new_workflow.add_node("memory_retriever", run_memory_retrieval)
+        new_workflow.add_node("intent_agent", run_intent_agent)
+        new_workflow.add_node("pm_agent", run_pm_agent)
+        new_workflow.add_node("pragmatist_agent", run_pragmatist_agent)
+        new_workflow.add_node("governance_agent", run_governance_agent)
+        new_workflow.add_node("experimenter_agent", run_experimenter_agent)
+        new_workflow.add_node("compliance_agent", run_compliance_agent)
+        new_workflow.add_node("synthesis_agent", run_synthesis_agent)
+        new_workflow.add_node("qa_agent", run_qa_agent)
+        new_workflow.add_node("observer_agent", run_observer_agent)
+        new_workflow.add_node("archivist_agent", run_archivist_agent)
+        new_workflow.add_node("knowledge_curator_agent", run_knowledge_curator_agent)
+        new_workflow.add_node("disclaimer_agent", run_disclaimer_agent)
+        
+        log.info("✅ All nodes added to new graph")
+        
+        # Set entry point
+        new_workflow.set_entry_point("memory_retriever")
+        
+        # Standard flow: Memory → Intent → PM
+        new_workflow.add_edge("memory_retriever", "intent_agent")
+        new_workflow.add_edge("intent_agent", "pm_agent")
+        
+        # NEW ROUTING: PM → Pragmatist → Governance
+        new_workflow.add_edge("pm_agent", "pragmatist_agent")
+        new_workflow.add_edge("pragmatist_agent", "governance_agent")
+        log.info("✅ New routing added: PM → Pragmatist → Governance")
+        
+        # Governance conditional: approved → Experimenter, rejected → PM
         def governance_decider(state: AgentState):
+            """Decide next step based on governance decision"""
             gov = state.get("governanceReport", {}) or {}
-            if gov.get("approved_for_experiment", True) is True:
+            decision = gov.get("governanceDecision", "approve")
+            approved = gov.get("approved_for_experiment", True)
+            
+            log.info(f"Governance decision: {decision}, approved: {approved}")
+            
+            if approved and decision in ("approve", "approve_with_warning"):
                 return "experimenter_agent"
-            return "pm_agent"
-
-        try:
-            mw.add_conditional_edges("governance_agent", governance_decider, {
+            else:
+                # Rejected or requires escalation - loop back to PM
+                return "pm_agent"
+        
+        new_workflow.add_conditional_edges(
+            "governance_agent",
+            governance_decider,
+            {
                 "experimenter_agent": "experimenter_agent",
                 "pm_agent": "pm_agent"
-            })
-            log.info("Added conditional edges for governance_agent")
-        except Exception as e:
-            log.debug("Could not add conditional edges for governance_agent: %s", e)
-
-        # experimenter -> compliance -> synthesis
-        try:
-            mw.add_edge("experimenter_agent", "compliance_agent")
-            log.info("Edge: experimenter_agent -> compliance_agent")
-        except Exception as e:
-            log.debug("experimenter_agent -> compliance_agent: %s", e)
-
-        try:
-            mw.add_edge("compliance_agent", "synthesis_agent")
-            log.info("Edge: compliance_agent -> synthesis_agent")
-        except Exception as e:
-            log.debug("compliance_agent -> synthesis_agent: %s", e)
-
-        # synthesis -> qa -> observer -> archivist -> knowledge_curator -> END
-        try:
-            mw.add_edge("synthesis_agent", "qa_agent")
-            log.info("Edge: synthesis_agent -> qa_agent ensured/added")
-        except Exception as e:
-            log.debug("synthesis_agent -> qa_agent: %s", e)
-
-        try:
-            mw.add_edge("qa_agent", "observer_agent")
-            log.info("Edge: qa_agent -> observer_agent")
-        except Exception as e:
-            log.debug("qa_agent -> observer_agent: %s", e)
-
-        try:
-            mw.add_edge("observer_agent", "archivist_agent")
-            log.info("Edge: observer_agent -> archivist_agent")
-        except Exception as e:
-            log.debug("observer_agent -> archivist_agent: %s", e)
-
-        try:
-            mw.add_edge("archivist_agent", "knowledge_curator_agent")
-            log.info("Edge: archivist_agent -> knowledge_curator_agent")
-        except Exception as e:
-            log.debug("archivist_agent -> knowledge_curator_agent: %s", e)
-
-        try:
-            mw.add_edge("knowledge_curator_agent", base_graph.END)
-            log.info("Edge: knowledge_curator_agent -> END")
-        except Exception as e:
-            log.debug("knowledge_curator_agent -> END: %s", e)
-
-        # keep QA -> disclaimer if present
-        try:
-            mw.add_edge("qa_agent", "disclaimer_agent")
-            log.info("Edge: qa_agent -> disclaimer_agent ensured/added")
-        except Exception as e:
-            log.debug("qa_agent -> disclaimer_agent: %s", e)
-
-        # Feedback agent wiring (best effort)
-        try:
-            from feedback_agent import run_feedback_agent
-            mw.add_node("feedback_agent", run_feedback_agent)
-            mw.add_edge("synthesis_agent", "feedback_agent")
-            mw.add_edge("archivist_agent", "feedback_agent")
-            mw.add_edge("feedback_agent", "archivist_agent")
-            log.info("Feedback agent wired.")
-        except Exception:
-            log.debug("Feedback agent wiring skipped or failed.")
-
-        # Recompile main_app defensively
-        try:
-            base_graph.main_app = mw.compile()
-            log.info("Recompiled main_workflow -> main_app")
-        except Exception as e:
-            log.warning("Could not recompile main_workflow: %s", e)
-
-        log.info("Graph upgrades applied successfully.")
+            }
+        )
+        log.info("✅ Governance conditional routing added")
+        
+        # Continue standard flow: Experimenter → Compliance → Synthesis → QA
+        new_workflow.add_edge("experimenter_agent", "compliance_agent")
+        new_workflow.add_edge("compliance_agent", "synthesis_agent")
+        new_workflow.add_edge("synthesis_agent", "qa_agent")
+        
+        # QA conditional routing (from original graph)
+        def should_continue(state: AgentState):
+            """Decide next step after QA"""
+            # Budget check first
+            if state.get("budget_exceeded"):
+                return "disclaimer_agent"
+            
+            # Parse rework cycles
+            try:
+                rework = int(state.get("rework_cycles", 0))
+                max_loops = int(state.get("max_loops", 0))
+            except Exception:
+                rework = state.get("rework_cycles", 0) or 0
+                max_loops = state.get("max_loops", 0) or 0
+            
+            # If approved → success path
+            if state.get("approved"):
+                return "observer_agent"
+            
+            # If exceeded rework limit → disclaimer
+            if rework > max_loops:
+                return "disclaimer_agent"
+            
+            # Otherwise → loop back to PM for revision
+            return "pm_agent"
+        
+        new_workflow.add_conditional_edges(
+            "qa_agent",
+            should_continue,
+            {
+                "observer_agent": "observer_agent",
+                "pm_agent": "pm_agent",
+                "disclaimer_agent": "disclaimer_agent"
+            }
+        )
+        log.info("✅ QA conditional routing added")
+        
+        # Final success path: Observer → Archivist → Knowledge Curator → END
+        new_workflow.add_edge("observer_agent", "archivist_agent")
+        new_workflow.add_edge("archivist_agent", "knowledge_curator_agent")
+        new_workflow.add_edge("knowledge_curator_agent", END)
+        
+        # Disclaimer path (failure/limit reached)
+        new_workflow.add_edge("disclaimer_agent", END)
+        
+        log.info("�� Final flow edges added")
+        
+        # CRITICAL: Compile NEW graph and REPLACE old one
+        base_graph.main_app = new_workflow.compile()
+        base_graph.main_workflow = new_workflow  # Also update workflow reference
+        
+        log.info("=" * 60)
+        log.info("✅ GRAPH REBUILD SUCCESSFUL")
+        log.info("=" * 60)
+        log.info("New flow: Memory → Intent → PM → Pragmatist → Governance")
+        log.info("          → Experimenter → Compliance → Synthesis → QA")
+        log.info("          → Observer → Archivist → Knowledge Curator → END")
+        log.info("=" * 60)
+        
         return True
-
+        
     except Exception as e:
-        log.exception("Failed to apply graph upgrades: %s", e)
+        log.exception(f"❌ Failed to rebuild graph: {e}")
         return False