TradingAgents/tradingagents/graph/enhanced_conditional_logic.py

"""
Enhanced Conditional Logic with Rejection Loops

Adds backward edges to send proposals back to agents if they fail validation.
"""

from tradingagents.agents.utils.agent_states import AgentState


class EnhancedConditionalLogic:
    """Handles conditional logic with rejection loops and quality checks."""
    
    def __init__(self, max_debate_rounds=1, max_risk_discuss_rounds=1):
        """Initialize with configuration parameters."""
        self.max_debate_rounds = max_debate_rounds
        self.max_risk_discuss_rounds = max_risk_discuss_rounds
    
    # ... (keep existing analyst conditional methods) ...
    
    def should_continue_debate_with_validation(self, state: AgentState) -> str:
        """
        Determine if debate should continue WITH QUALITY CHECKS.
        
        This replaces the naive round-robin with actual validation.
        """
        debate_state = state["investment_debate_state"]
        
        # Check 1: Was last argument fact-checked and rejected?
        if debate_state.get("last_argument_invalid", False):
            # Send back to same agent to revise
            print(f"❌ REJECTED: {debate_state.get('rejection_reason', 'Invalid argument')}")
            print(f"   Sending back to {debate_state['latest_speaker']} for revision")
            
            # Route back to the agent that made the bad argument
            if debate_state["latest_speaker"] == "Bull":
                return "Bull Researcher"
            else:
                return "Bear Researcher"
        
        # Check 2: Has consensus been reached?
        if debate_state.get("consensus_reached", False):
            print("✅ CONSENSUS REACHED: Proceeding to Research Manager")
            return "Research Manager"
        
        # Check 3: Max rounds exceeded
        if debate_state["count"] >= 2 * self.max_debate_rounds:
            print(f"⏱️  MAX ROUNDS REACHED: {debate_state['count']} rounds")
            return "Research Manager"
        
        # Check 4: Confidence too low (force another round)
        if debate_state.get("confidence", 1.0) < 0.5:
            print(f"⚠️  LOW CONFIDENCE ({debate_state['confidence']:.1%}): Continuing debate")
            # Continue round-robin
            if debate_state["current_response"].startswith("Bull"):
                return "Bear Researcher"
            return "Bull Researcher"
        
        # Default: Round-robin
        if debate_state["current_response"].startswith("Bull"):
            return "Bear Researcher"
        return "Bull Researcher"
    
    def should_proceed_after_risk_gate(self, state: AgentState) -> str:
        """
        Determine next step after deterministic risk gate validation.
        
        This is a NEW node that checks mathematical risk validation.
        """
        risk_validation = state.get("risk_gate_result", {})
        
        # Check 1: Was trade rejected by risk gate?
        if not risk_validation.get("approved", False):
            rejection_reason = risk_validation.get("rejection_reason", "Unknown")
            
            # Determine severity
            if "CIRCUIT BREAKER" in rejection_reason:
                # Critical failure - halt trading
                print(f"🚨 CIRCUIT BREAKER TRIGGERED: {rejection_reason}")
                return "END"
            
            elif "DATA QUALITY" in rejection_reason:
                # Data issue - send back to analysts
                print(f"📊 DATA QUALITY FAILURE: {rejection_reason}")
                print("   Routing back to Market Analyst for data refresh")
                return "Market Analyst"
            
            elif "PORTFOLIO HEAT" in rejection_reason or "POSITION RISK" in rejection_reason:
                # Risk limit exceeded - send to Risk Manager for review
                print(f"⚠️  RISK LIMIT EXCEEDED: {rejection_reason}")
                print("   Routing to Risk Manager for position adjustment")
                return "Risk Manager Revision"
            
            else:
                # Generic rejection - log and hold
                print(f"❌ TRADE REJECTED: {rejection_reason}")
                return "END"
        
        # Check 2: Was position size overridden?
        if risk_validation.get("override_message"):
            print(f"🔧 {risk_validation['override_message']}")
        
        # Approved - proceed to execution
        print("✅ RISK GATE PASSED: Trade approved")
        return "Execute Trade"
    
    def should_continue_risk_analysis_with_validation(self, state: AgentState) -> str:
        """
        Enhanced risk analysis routing with validation.
        """
        risk_state = state["risk_debate_state"]
        
        # Check 1: Did any analyst provide mathematically invalid reasoning?
        if risk_state.get("invalid_reasoning_detected", False):
            # Send back to the analyst who made the error
            print(f"❌ INVALID REASONING: {risk_state.get('error_message', '')}")
            return risk_state["latest_speaker"]
        
        # Check 2: Max rounds
        if risk_state["count"] >= 3 * self.max_risk_discuss_rounds:
            return "Deterministic Risk Gate"  # NEW: Route to math validation
        
        # Round-robin
        if risk_state["latest_speaker"].startswith("Risky"):
            return "Safe Analyst"
        if risk_state["latest_speaker"].startswith("Safe"):
            return "Neutral Analyst"
        return "Risky Analyst"


# Integration example for trading_graph.py
"""
To integrate this into your graph:

1. Add the Deterministic Risk Gate node:
   workflow.add_node("Deterministic Risk Gate", deterministic_risk_gate_node)

2. Replace the edge from "Risk Judge" to END:
   # OLD:
   workflow.add_edge("Risk Judge", END)
   
   # NEW:
   workflow.add_conditional_edges(
       "Risk Judge",
       enhanced_logic.should_proceed_after_risk_gate,
       {
           "END": END,
           "Market Analyst": "Market Analyst",  # Data quality failure
           "Risk Manager Revision": "Risk Manager Revision",  # Risk limit exceeded
           "Execute Trade": "Execute Trade"  # Approved
       }
   )

3. Add backward edge for debate rejection:
   workflow.add_conditional_edges(
       "Bull Researcher",
       enhanced_logic.should_continue_debate_with_validation,
       {
           "Bear Researcher": "Bear Researcher",
           "Bull Researcher": "Bull Researcher",  # NEW: Rejection loop
           "Research Manager": "Research Manager",
       }
   )
"""