action_router.py

# action_router.py
import logging
from typing import Dict, Any
from datetime import datetime

class ActionRouter:
    def __init__(self):
        self.logger = logging.getLogger('ActionRouter')
        self._init_action_handlers()
        
    def _init_action_handlers(self):
        """Initialize all action handler configurations"""
        self.handlers = {
            'email': {
                'high_urgency_complaint': {
                    'condition': lambda data: (
                        data.get("analysis", {}).get("tone") == "urgent" and 
                        "complaint" in data.get("analysis", {}).get("category", "")
                    ),
                    'action': self._escalate_to_crm
                },
                'management_alert': {
                    'condition': lambda data: "management" in data.get("route_to", ""),
                    'action': self._notify_management
                }
            },
            'json': {
                'payment_received': {
                    'condition': lambda data: (
                        data.get("extracted", {}).get("event_type") == "payment_received"
                    ),
                    'action': self._route_to_accounting
                },
                'user_created': {
                    'condition': lambda data: (
                        data.get("extracted", {}).get("event_type") == "user_created"
                    ),
                    'action': self._route_to_crm
                },
                'high_risk_transaction': {
                    'condition': lambda data: (
                        data.get("extracted", {}).get("event_type") == "high_risk_transaction" or
                        float(data.get("extracted", {}).get("payload", {}).get("amount", 0)) > 10000
                    ),
                    'action': self._trigger_fraud_alert
                }
            },
            'pdf': {
                'high_value_invoice': {
                    'condition': lambda data: (
                        data.get("type") == "invoice" and
                        float(data.get("extractions", {}).get("total_amount", 0)) > 10000
                    ),
                    'action': self._flag_high_value_transaction
                },
                'compliance_alert': {
                    'condition': lambda data: bool(
                        data.get("extractions", {}).get("keywords")
                    ),
                    'action': self._handle_compliance_issues
                },
                'sensitive_document': {
                    'condition': lambda data: (
                        data.get("type") == "policy" and
                        any(kw in data.get("extractions", {}).get("keywords", []) 
                            for kw in ["GDPR", "HIPAA"])
                    ),
                    'action': self._restrict_document_access
                }
            }
        }

    def handle_agent_output(self, agent_output: Dict[str, Any]):
        """Main routing method for all agent outputs"""
        try:
            source_type = self._determine_source_type(agent_output)
            
            if not source_type:
                self.logger.warning(f"Unknown source type for output: {agent_output}")
                return

            for handler_name, handler_config in self.handlers.get(source_type, {}).items():
                if handler_config['condition'](agent_output):
                    self.logger.info(f"Triggering {handler_name} for {source_type}")
                    handler_config['action'](agent_output)
                    
        except Exception as e:
            self.logger.error(f"Error in action routing: {str(e)}")
            raise

    def _determine_source_type(self, data: Dict) -> str:
        """Determine the source agent type from output structure"""
        if "analysis" in data and "route_to" in data:
            return "email"
        elif "extracted" in data and "status" in data:
            return "json"
        elif "extractions" in data and "type" in data:
            return "pdf"
        return ""

    # Email Agent Actions
    def _escalate_to_crm(self, data: Dict):
        """Handle high-priority customer complaints"""
        ticket_data = {
            "customer": data.get("metadata", {}).get("sender"),
            "subject": data.get("metadata", {}).get("subject"),
            "priority": "P1",
            "timestamp": datetime.now().isoformat()
        }
        print(f"🚨 [CRM] Creating high-priority ticket: {ticket_data}")
        # Implement actual CRM API call here

    def _notify_management(self, data: Dict):
        """Alert management team about important emails"""
        alert_msg = (
            f"Management attention required for email from {data.get('metadata', {}).get('sender')}\n"
            f"Subject: {data.get('metadata', {}).get('subject')}\n"
            f"Category: {data.get('analysis', {}).get('category')}"
        )
        print(f"📩 [MANAGEMENT] {alert_msg}")
        # Implement management notification here

    # JSON Agent Actions
    def _route_to_accounting(self, data: Dict):
        """Process payment events"""
        payment_data = {
            "amount": data.get("extracted", {}).get("payload", {}).get("amount"),
            "currency": data.get("extracted", {}).get("payload", {}).get("currency"),
            "reference": data.get("extracted", {}).get("payload", {}).get("transaction_id")
        }
        print(f"💰 [ACCOUNTING] Processing payment: {payment_data}")
        # Implement accounting system integration here

    def _route_to_crm(self, data: Dict):
        """Process new user events"""
        user_data = {
            "user_id": data.get("extracted", {}).get("payload", {}).get("user_id"),
            "email": data.get("extracted", {}).get("payload", {}).get("email"),
            "signup_date": data.get("extracted", {}).get("timestamp")
        }
        print(f"👤 [CRM] Creating user profile: {user_data}")
        # Implement CRM system integration here

    def _trigger_fraud_alert(self, data: Dict):
        """Flag suspicious transactions"""
        alert_data = {
            "event_type": data.get("extracted", {}).get("event_type"),
            "amount": data.get("extracted", {}).get("payload", {}).get("amount"),
            "risk_score": "high",
            "timestamp": data.get("extracted", {}).get("timestamp")
        }
        print(f"🛑 [FRAUD] Alert triggered: {alert_data}")
        # Implement fraud detection system integration here

    # PDF Agent Actions
    def _flag_high_value_transaction(self, data: Dict):
        """Handle large financial transactions"""
        transaction_data = {
            "amount": data.get("extractions", {}).get("total_amount"),
            "vendor": data.get("extractions", {}).get("vendor"),
            "document": data.get("source"),
            "flags": data.get("flags", [])
        }
        print(f"💸 [FINANCE] High value transaction: {transaction_data}")
        # Implement financial control system integration here

    def _handle_compliance_issues(self, data: Dict):
        """Process compliance-related findings"""
        compliance_data = {
            "document_type": data.get("type"),
            "keywords": data.get("extractions", {}).get("keywords"),
            "source": data.get("source"),
            "timestamp": datetime.now().isoformat()
        }
        print(f"⚖️ [COMPLIANCE] Issues found: {compliance_data}")
        # Implement compliance system integration here

    def _restrict_document_access(self, data: Dict):
        """Apply security controls to sensitive documents"""
        security_data = {
            "document": data.get("source"),
            "sensitivity": "high",
            "access_control": "legal-team-only",
            "keywords": data.get("extractions", {}).get("keywords")
        }
        print(f"🔒 [SECURITY] Applying restrictions: {security_data}")
        # Implement document security controls here