Update graph_upgraded.py

graph_upgraded.py

@@ -514,22 +514,36 @@ def apply_upgrades():
     try:
         from langgraph.graph import StateGraph, END
         
+        # Import the missing agent functions from base graph
+        from graph import (
+            run_memory_retrieval,
+            run_intent_agent,
+            run_pm_agent,
+            run_experimenter_agent,
+            run_synthesis_agent,
+            run_qa_agent,
+            run_archivist_agent,
+            run_disclaimer_agent,
+            should_continue,
+            should_run_experiment
+        )
+        
         # Create BRAND NEW graph
         new_workflow = StateGraph(AgentState)
         
-        # Add all nodes (reusing original functions where appropriate)
+        # Add all nodes (using imported functions and local upgraded ones)
         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("pragmatist_agent", run_pragmatist_agent)  # Local upgraded
+        new_workflow.add_node("governance_agent", run_governance_agent)    # Local upgraded
         new_workflow.add_node("experimenter_agent", run_experimenter_agent)
-        new_workflow.add_node("compliance_agent", run_compliance_agent)
+        new_workflow.add_node("compliance_agent", run_compliance_agent)    # Local upgraded
         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("observer_agent", run_observer_agent)        # Local upgraded
         new_workflow.add_node("archivist_agent", run_archivist_agent)
-        new_workflow.add_node("knowledge_curator_agent", run_knowledge_curator_agent)
+        new_workflow.add_node("knowledge_curator_agent", run_knowledge_curator_agent)  # Local upgraded
         new_workflow.add_node("disclaimer_agent", run_disclaimer_agent)
         
         log.info("✅ All nodes added to new graph")
@@ -576,32 +590,7 @@ def apply_upgrades():
         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"
-        
+        # QA conditional routing (use imported should_continue)
         new_workflow.add_conditional_edges(
             "qa_agent",
             should_continue,