Basic web loop works

sia/web_agent.py

@@ -1,4 +1,5 @@
 from enum import Enum, auto
+from threading import Thread, Lock
 from typing import Callable, List, Optional
 
 from .base_agent import BaseAgent
@@ -13,17 +14,15 @@ from .xml_validator import XMLValidator
 class WebAgentState(Enum):
     """
     States for the web agent state machine.
-    
-    States:
-        UPDATE: Updating system metrics and working memory entries
-        CONTEXT_APPROVAL: Waiting for human approval of context
-        INFERENCE: Processing context through LLM
-        RESPONSE_APPROVAL: Waiting for human approval of LLM response
     """
     UPDATE = auto()
+    """Updating system metrics and working memory entries"""
     CONTEXT_APPROVAL = auto()
+    """Waiting for human approval of context"""
     INFERENCE = auto()
+    """Processing context through LLM"""
     RESPONSE_APPROVAL = auto()
+    """Waiting for human approval of LLM response"""
     STOPPED = auto()
 
 class WebAgent(BaseAgent):
@@ -48,17 +47,19 @@ class WebAgent(BaseAgent):
         super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser)
         self._response = ""
         self._system_prompt = system_prompt
-        self._current_state = WebAgentState.CONTEXT_APPROVAL
+        self._state = WebAgentState.CONTEXT_APPROVAL
         self._validation_error: Optional[str] = None
         self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
+        self._context_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()
+        self._context = self._compile_context()
+        self._state_lock = Lock()
 
     @property
-    def current_state(self) -> WebAgentState:
+    def state(self) -> WebAgentState:
         """Get the current state of the agent."""
-        return self._current_state
+        return self._state
     
     @property
     def command_result(self) -> Optional[CommandResult]:
@@ -74,6 +75,11 @@ class WebAgent(BaseAgent):
     def response(self) -> str:
         """Get the current response."""
         return self._response
+    
+    @property
+    def context(self) -> str:
+        """Get the current context."""
+        return self._context
         
     def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
         """
@@ -85,6 +91,16 @@ class WebAgent(BaseAgent):
         if handler not in self._state_change_handlers:
             self._state_change_handlers.append(handler)
     
+    def add_context_change_handler(self, handler: Callable[[str], None]) -> None:
+        """
+        Add a callback for context changes.
+
+        Args:
+            handler: Function to call with new context
+        """
+        if handler not in self._context_change_handlers:
+            self._context_change_handlers.append(handler)
+    
     def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None:
         """
         Add a callback for response changes.
@@ -94,53 +110,54 @@ class WebAgent(BaseAgent):
         """
         if handler not in self._response_change_handlers:
             self._response_change_handlers.append(handler)
-            
-    def _notify_state_change(self) -> None:
-        """Notify all handlers of state change."""
-        for handler in self._state_change_handlers:
-            handler(self._current_state)
     
-    def set_response(self, response: str) -> None:
-        """Set response and notify handlers."""
-        self._response = response
-        validation_error = self._validator.validate(response)
-        self._validation_error = validation_error
-        for handler in self._response_change_handlers:
-            handler(response, validation_error)
-
-    def _set_state(self, new_state: WebAgentState) -> None:
-        """
-        Set the current state and notify handlers.
-
-        Args:
-            new_state: New state to set
-        """
-        self._current_state = new_state
-        self._notify_state_change()
+    def modify_context(self, context: str) -> None:
+        with self._state_lock:
+            if self._state != WebAgentState.CONTEXT_APPROVAL:
+                error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
+                raise Exception(error_msg)
+        Thread(target=self._set_context, args=(context,)).start()
+    
+    def modify_response(self, response: str) -> None:
+        started = False
+        try:
+            self._state_lock.acquire()
+            if self._state != WebAgentState.RESPONSE_APPROVAL:
+                error_msg = f"Not in RESPONSE_APPROVAL state (current: {self._state})"
+                raise Exception(error_msg)
+            Thread(target=self._set_response, args=(response, self._state_lock)).start()
+            started = True
+        finally:
+            if not started:
+                self._state_lock.release()
             
     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)
+        with self._state_lock:
+            if self._state != WebAgentState.CONTEXT_APPROVAL:
+                error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
+                raise Exception(error_msg)
+            self._set_state(WebAgentState.INFERENCE)
+        Thread(target=self._approve_context_thread).start()
+
+    def approve_response(self) -> None:
+        with self._state_lock:
+            if self._state != WebAgentState.RESPONSE_APPROVAL:
+                raise Exception("Not in RESPONSE_APPROVAL state")
+            self._set_state(WebAgentState.UPDATE)
+            Thread(target=self._approve_response_thread).start()
                 
-        self._set_state(WebAgentState.INFERENCE)
-        self.set_response("")
-        
+    def _approve_context_thread(self) -> None:
+        self._set_response("")
         response_token_iter = self._llm.infer(f"{self._system_prompt}\n{self._action_schema}", self.context)
         response = ""
         for token in response_token_iter:
             response += token
-            self.set_response(response)
+            self._set_response(response)
             print(f"{token}")
         self._set_state(WebAgentState.RESPONSE_APPROVAL)
-
-    def approve_response(self) -> None:
-        if self._current_state != WebAgentState.RESPONSE_APPROVAL:
-            raise Exception("Not in RESPONSE_APPROVAL state")
-
-        self._set_state(WebAgentState.UPDATE)
-
-        parse_result = self._parser.parse(self.response)
+    
+    def _approve_response_thread(self) -> None:
+        parse_result = self._parser.parse(self._response)
         if isinstance(parse_result, Command):
             result = parse_result.execute(self._working_memory)
             self._command_result = result
@@ -151,5 +168,29 @@ class WebAgent(BaseAgent):
             self._working_memory.add_entry(parse_result)
         
         self._working_memory.update()
-        self.context = self._compile_context()
+        self._set_context(self._compile_context())
         self._set_state(WebAgentState.CONTEXT_APPROVAL)
+            
+    def _set_state(self, state) -> None:
+        """Notify all handlers of state change."""
+        self._state = state
+        for handler in self._state_change_handlers:
+            handler(self._state)
+    
+    def _set_context(self, context) -> None:
+        """Notify all handlers of context change."""
+        self._context = context
+        for handler in self._context_change_handlers:
+            handler(context)
+    
+    def _set_response(self, response: str, lock: Optional[Lock] = None) -> None:
+        """Set response and notify handlers."""
+        try:
+            self._response = response
+            validation_error = self._validator.validate(response)
+            self._validation_error = validation_error
+            for handler in self._response_change_handlers:
+                handler(response, validation_error)
+        finally:
+            if lock is not None:
+                lock.release()