Separate and tested web_socket_manager

sia/web_agent.py

@@ -24,6 +24,7 @@ class WebAgentState(Enum):
     CONTEXT_APPROVAL = auto()
     INFERENCE = auto()
     RESPONSE_APPROVAL = auto()
+    STOPPED = auto()
 
 class WebAgent(BaseAgent):
     """
@@ -45,12 +46,12 @@ class WebAgent(BaseAgent):
         Initialize web agent with required components.
         """
         super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser)
-        self.response = ""
+        self._response = ""
         self._system_prompt = system_prompt
         self._current_state = WebAgentState.CONTEXT_APPROVAL
         self._validation_error: Optional[str] = None
         self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
-        self._response_change_handlers: List[Callable[[str], None]] = []
+        self._response_change_handlers: List[Callable[[str, str], None]] = []
         self._command_result: Optional[CommandResult] = None
         self.context = self._compile_context()
 
@@ -63,6 +64,16 @@ class WebAgent(BaseAgent):
     def command_result(self) -> Optional[CommandResult]:
         """Get the result of the last command execution."""
         return self._command_result
+
+    @property
+    def validation_error(self) -> Optional[str]:
+        """Get the validation error, if any."""
+        return self._validation_error
+    
+    @property
+    def response(self) -> str:
+        """Get the current response."""
+        return self._response
         
     def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
         """
@@ -74,7 +85,7 @@ class WebAgent(BaseAgent):
         if handler not in self._state_change_handlers:
             self._state_change_handlers.append(handler)
     
-    def add_response_change_handler(self, handler: Callable[[str], None]) -> None:
+    def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None:
         """
         Add a callback for response changes.
 
@@ -89,20 +100,13 @@ class WebAgent(BaseAgent):
         for handler in self._state_change_handlers:
             handler(self._current_state)
     
-    def _notify_response_change(self) -> None:
-        """Notify all handlers of response change."""
-        for handler in self._response_change_handlers:
-            handler(self.response)
-    
-    def _set_response(self, response: str) -> None:
+    def set_response(self, response: str) -> None:
         """Set response and notify handlers."""
-        self.response = response
-        self._validation_error = self._validator.validate(response)
+        self._response = response
+        validation_error = self._validator.validate(response)
+        self._validation_error = validation_error
         for handler in self._response_change_handlers:
-            try:
-                handler(response)
-            except Exception as e:
-                print(f"Error in response handler: {e}")
+            handler(response, validation_error)
 
     def _set_state(self, new_state: WebAgentState) -> None:
         """
@@ -114,39 +118,22 @@ class WebAgent(BaseAgent):
         self._current_state = new_state
         self._notify_state_change()
             
-    async def approve_context(self) -> None:
+    def approve_context(self) -> None:
         if self._current_state != WebAgentState.CONTEXT_APPROVAL:
             error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._current_state})"
             raise Exception(error_msg)
                 
         self._set_state(WebAgentState.INFERENCE)
-        self._set_response("")
+        self.set_response("")
         
         response_token_iter = self._llm.infer(self._system_prompt, self.context)
         response = ""
-        try:
-            for token in response_token_iter:
-                response += token
-                await self._set_response(response)
-                print(f"{token}")
-        except Exception as e:
-            print(f"Error during inference: {e}")
-            self._set_state(WebAgentState.CONTEXT_APPROVAL)
-            return
-        
-        await self._set_response(response)
+        for token in response_token_iter:
+            response += token
+            self.set_response(response)
+            print(f"{token}")
         self._set_state(WebAgentState.RESPONSE_APPROVAL)
 
-    async def _set_response(self, response: str) -> None:
-        """Set response and notify handlers asynchronously."""
-        self.response = response
-        self._validation_error = self._validator.validate(response)
-        for handler in self._response_change_handlers:
-            try:
-                await handler(response)
-            except Exception as e:
-                print(f"Error in response handler: {e}")
-
     def approve_response(self) -> None:
         if self._current_state != WebAgentState.RESPONSE_APPROVAL:
             raise Exception("Not in RESPONSE_APPROVAL state")
@@ -157,10 +144,12 @@ class WebAgent(BaseAgent):
         if isinstance(parse_result, Command):
             result = parse_result.execute(self._working_memory)
             self._command_result = result
+            if result.should_stop:
+                self._set_state(WebAgentState.STOPPED)
+                return
         else:
             self._working_memory.add_entry(parse_result)
         
         self._working_memory.update()
         self.context = self._compile_context()
-
         self._set_state(WebAgentState.CONTEXT_APPROVAL)