atlas/home-voice-agent/safety/confirmations/flow.py

"""
Confirmation flow manager.

Orchestrates the confirmation process for high-risk actions.
"""

import logging
from typing import Dict, Any, Optional, Tuple
from safety.confirmations.risk import get_classifier, RiskLevel
from safety.confirmations.confirmation_token import get_token_system

logger = logging.getLogger(__name__)


class ConfirmationFlow:
    """Manages confirmation flows for tool calls."""
    
    def __init__(self):
        """Initialize confirmation flow."""
        self.classifier = get_classifier()
        self.token_system = get_token_system()
    
    def check_confirmation_required(self, tool_name: str, **kwargs) -> Tuple[bool, Optional[str]]:
        """
        Check if confirmation is required and get message.
        
        Args:
            tool_name: Name of the tool
            **kwargs: Tool parameters
            
        Returns:
            (requires_confirmation, confirmation_message)
        """
        requires = self.classifier.requires_confirmation(tool_name, **kwargs)
        if requires:
            message = self.classifier.get_confirmation_message(tool_name, **kwargs)
            return True, message
        return False, None
    
    def generate_confirmation_token(self,
                                   tool_name: str,
                                   parameters: Dict[str, Any],
                                   session_id: Optional[str] = None,
                                   user_id: Optional[str] = None) -> str:
        """
        Generate a confirmation token for an action.
        
        Args:
            tool_name: Name of the tool
            parameters: Tool parameters
            session_id: Session ID
            user_id: User ID
            
        Returns:
            Confirmation token
        """
        return self.token_system.generate_token(
            tool_name=tool_name,
            parameters=parameters,
            session_id=session_id,
            user_id=user_id
        )
    
    def validate_confirmation(self,
                              token: str,
                              tool_name: str,
                              parameters: Dict[str, Any]) -> Tuple[bool, Optional[str]]:
        """
        Validate a confirmation token.
        
        Args:
            token: Confirmation token
            tool_name: Expected tool name
            parameters: Expected parameters
            
        Returns:
            (is_valid, error_message)
        """
        return self.token_system.verify_action(token, tool_name, parameters)
    
    def process_confirmation_request(self,
                                   tool_name: str,
                                   parameters: Dict[str, Any],
                                   session_id: Optional[str] = None,
                                   user_id: Optional[str] = None) -> Dict[str, Any]:
        """
        Process a confirmation request.
        
        Args:
            tool_name: Name of the tool
            parameters: Tool parameters
            session_id: Session ID
            user_id: User ID
            
        Returns:
            Response dict with confirmation_required, message, and token
        """
        requires, message = self.check_confirmation_required(tool_name, **parameters)
        
        if requires:
            token = self.generate_confirmation_token(
                tool_name=tool_name,
                parameters=parameters,
                session_id=session_id,
                user_id=user_id
            )
            
            return {
                "confirmation_required": True,
                "message": message,
                "token": token,
                "risk_level": self.classifier.classify_risk(tool_name, **parameters).value
            }
        else:
            return {
                "confirmation_required": False,
                "message": None,
                "token": None,
                "risk_level": "low"
            }


# Global flow instance
_flow = ConfirmationFlow()


def get_flow() -> ConfirmationFlow:
    """Get the global confirmation flow instance."""
    return _flow