Fixed auto approver

sia/auto_approver.py

@@ -1,12 +1,13 @@
 from threading import Thread, Event
 from typing import Callable, TypeAlias, TypedDict
-from .web_agent import WebAgent, WebAgentState
+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]
 
@@ -20,6 +21,7 @@ class AutoApprover:
         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
@@ -29,7 +31,8 @@ class AutoApprover:
         self._response_thread: Thread | None = None
         self._config_change_handlers: list[ConfigChangeHandler] = []
         
-        self.agent.add_llm_change_handler(self._handle_state_change)
+        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:
@@ -37,25 +40,29 @@ class AutoApprover:
             context_enabled=self._context_enabled,
             response_enabled=self._response_enabled,
             context_timeout=self._context_timeout,
-            response_timeout=self._response_timeout
+            response_timeout=self._response_timeout,
+            llm_name=self._llm_name
         )
 
     def set_config(self, config: AutoApproverConfig) -> None:
-        self._context_enabled = config['context_enabled']
-        self._response_enabled = config['response_enabled'] 
-        self._context_timeout = config['context_timeout']
-        self._response_timeout = config['response_timeout']
+        if config['llm_name'] not in self.agent.llms:
+            raise ValueError(f"Unknown LLM: {config['llm_name']}")
         
-        if self._context_enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
-            self._start_context_thread()
-        else:
-            self._stop_context_thread()
+        notify_config_change = self._notify_config_change
+        self._notify_config_change = lambda: None
+            
+        try:
+            self.context_enabled = False
+            self.response_enabled = False
 
-        if self._response_enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
-            self._start_response_thread()
-        else:
-            self._stop_response_thread()
-        
+            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:
@@ -98,7 +105,7 @@ class AutoApprover:
             return
         self._context_enabled = enabled
         self._stop_context_thread()
-        if enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
+        if enabled and self.agent.llms[self._llm_name] == LlmState.NO_OUTPUT:
             self._start_context_thread()
         self._notify_config_change()
 
@@ -112,21 +119,36 @@ class AutoApprover:
             return
         self._response_enabled = enabled
         self._stop_response_thread()
-        if enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
+        if enabled and self.agent.llms[self._llm_name] == LlmState.OUTPUT:
             self._start_response_thread()
         self._notify_config_change()
 
-    def _handle_state_change(self, new_state: WebAgentState) -> None:
-        if new_state == WebAgentState.CONTEXT_APPROVAL:
-            self._start_context_thread()
-        elif self._context_enabled:
-            self._stop_context_thread()
+    @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()
 
-        if new_state == WebAgentState.RESPONSE_APPROVAL:
+    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()
-        elif self._response_enabled:
+        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()
@@ -150,13 +172,12 @@ class AutoApprover:
     def _context_approval_thread(self) -> None:
         if self._stop_event.wait(self._context_timeout):
             return
-        if (self._context_enabled and 
-            self.agent.state == WebAgentState.CONTEXT_APPROVAL):
-            self.agent.approve_context()
+        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.state == WebAgentState.RESPONSE_APPROVAL):
-            self.agent.approve_response()
+            self.agent.llms[self._llm_name] == LlmState.OUTPUT):
+            self.agent.approve_response(self._llm_name, self.agent.get_output(self._llm_name))