SIA/sia/auto_approver.py

from threading import Thread, Event
from typing import Callable, TypeAlias, TypedDict
from .web_agent import WebAgent, LlmState

class AutoApproverConfig(TypedDict):
    context_enabled: bool
    response_enabled: bool
    context_timeout: float
    response_timeout: float
    llm_name: str

ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]

class AutoApprover:
    """
    Handles automatic approvals for WebAgent states after configurable timeouts.
    """
    
    def __init__(self, agent: WebAgent):
        """
        Initialize auto approver with a WebAgent instance.
        """
        self.agent = agent
        self._llm_name = next(iter(agent.llms.keys()))
        self._context_timeout = 5.0
        self._response_timeout = 10.0
        self._context_enabled = False
        self._response_enabled = False
        self._stop_event = Event()
        self._context_thread: Thread | None = None
        self._response_thread: Thread | None = None
        self._config_change_handlers: list[ConfigChangeHandler] = []
        
        self.agent.add_llm_change_handler(self._handle_llm_state_change)
        self.agent.add_context_change_handler(self._handle_context_change)

    @property
    def config(self) -> AutoApproverConfig:
        return AutoApproverConfig(
            context_enabled=self._context_enabled,
            response_enabled=self._response_enabled,
            context_timeout=self._context_timeout,
            response_timeout=self._response_timeout,
            llm_name=self._llm_name
        )

    def set_config(self, config: AutoApproverConfig) -> None:
        if config['llm_name'] not in self.agent.llms:
            raise ValueError(f"Unknown LLM: {config['llm_name']}")
        
        notify_config_change = self._notify_config_change
        self._notify_config_change = lambda: None
            
        try:
            self.context_enabled = False
            self.response_enabled = False

            self.context_timeout = config['context_timeout']
            self.response_timeout = config['response_timeout']
            self.llm_name = config['llm_name']

            self.context_enabled = config['context_enabled']
            self.response_enabled = config['response_enabled'] 
        finally:
            self._notify_config_change = notify_config_change
        self._notify_config_change()

    def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
        self._config_change_handlers.append(handler)

    def _notify_config_change(self) -> None:
        current_config = self.config
        for handler in self._config_change_handlers:
            handler(current_config)

    @property
    def context_timeout(self) -> float:
        return self._context_timeout
    
    @context_timeout.setter 
    def context_timeout(self, timeout: float) -> None:
        if self._context_enabled:
            raise ValueError("Cannot change timeout while auto-approval is enabled")
        self._context_timeout = timeout
        self._notify_config_change()

    @property
    def response_timeout(self) -> float:
        return self._response_timeout
    
    @response_timeout.setter
    def response_timeout(self, timeout: float) -> None:
        if self._response_enabled:
            raise ValueError("Cannot change timeout while auto-approval is enabled")
        self._response_timeout = timeout
        self._notify_config_change()

    @property
    def context_enabled(self) -> bool:
        return self._context_enabled

    @context_enabled.setter
    def context_enabled(self, enabled: bool) -> None:
        if enabled == self._context_enabled:
            return
        self._context_enabled = enabled
        self._stop_context_thread()
        if enabled and self.agent.llms[self._llm_name] == LlmState.NO_OUTPUT:
            self._start_context_thread()
        self._notify_config_change()

    @property 
    def response_enabled(self) -> bool:
        return self._response_enabled

    @response_enabled.setter
    def response_enabled(self, enabled: bool) -> None:
        if enabled == self._response_enabled:
            return
        self._response_enabled = enabled
        self._stop_response_thread()
        if enabled and self.agent.llms[self._llm_name] == LlmState.OUTPUT:
            self._start_response_thread()
        self._notify_config_change()

    @property
    def llm_name(self) -> str:
        return self._llm_name
    
    @llm_name.setter
    def llm_name(self, name: str) -> None:
        if name not in self.agent.llms:
            raise ValueError(f"Unknown LLM: {name}")
        self._llm_name = name
        self._notify_config_change()

    def _handle_llm_state_change(self, llm_name: str, state: LlmState) -> None:
        if llm_name != self._llm_name:
            return
            
        if state == LlmState.OUTPUT and self._response_enabled:
            self._start_response_thread()
        else:
            self._stop_response_thread()

    def _handle_context_change(self, context: str, generated: bool) -> None:
        if generated and self._context_enabled:
            self._start_context_thread()
        else:
            self._stop_context_thread()

    def _stop_context_thread(self) -> None:
        if self._context_thread:
            self._stop_event.set()
            self._context_thread = None
            self._stop_event.clear()

    def _stop_response_thread(self) -> None:
        if self._response_thread:
            self._stop_event.set()
            self._response_thread = None
            self._stop_event.clear()

    def _start_context_thread(self) -> None:
        self._context_thread = Thread(target=self._context_approval_thread)
        self._context_thread.start()

    def _start_response_thread(self) -> None:
        self._response_thread = Thread(target=self._response_approval_thread)
        self._response_thread.start()

    def _context_approval_thread(self) -> None:
        if self._stop_event.wait(self._context_timeout):
            return
        if self._context_enabled:
            self.agent.run_inference(self._llm_name)

    def _response_approval_thread(self) -> None:
        if self._stop_event.wait(self._response_timeout):
            return
        if (self._response_enabled and 
            self.agent.llms[self._llm_name] == LlmState.OUTPUT):
            self.agent.approve_response(self._llm_name, self.agent.get_output(self._llm_name))