Fix client side message handling

sia/web_agent.py

@@ -126,7 +126,7 @@ class WebAgent(BaseAgent):
         self._set_state(WebAgentState.INFERENCE)
         self.set_response("")
         
-        response_token_iter = self._llm.infer(self._system_prompt, self.context)
+        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

sia/web_socket_manager.py

@@ -86,11 +86,8 @@ class WebSocketManager:
             })
             async for msg in ws:
                 if msg.type == WSMsgType.TEXT:
-                    #try:
                     data = json.loads(msg.data)
                     await self._handle_client_message(request, data)
-                    #except Exception as e:
-                    #    print(f"Error handling incoming websocket message: {msg}\n{e}")
                 elif msg.type == WSMsgType.ERROR:
                     print(f"WebSocket connection closed with error: {ws.exception()}")
         finally:

system_prompt.md

@@ -85,4 +85,4 @@ By maintaining a dynamic relationship between context and action, SIA can tackle
 # User interaction
 Be a helpful assistant to the user.
 Get to know them and make notes about them.
-Open the relevant user notes when you interact with them.
+Open the relevant user notes when you interact with them.

web/src/components/App.jsx

@@ -102,15 +102,15 @@ const SIAInterface = () => {
   const handleWebSocketMessage = (message) => {
     switch (message.type) {
       case MessageType.STATE_CHANGE:
-        setAgentState(message.data);
+        setAgentState(message.state);
         break;
       case MessageType.CONTEXT_UPDATE:
-        setContext(message.data);
+        setContext(message.context);
         break;
       case MessageType.RESPONSE_UPDATE:
-        setOriginalResponse(message.data);
-        setModifiedResponse(message.data);
-        setValidationError(message.validation_error || null);
+        setOriginalResponse(message.response);
+        setModifiedResponse(message.response);
+        setValidationError(message.validation_error);
         break;
       case MessageType.OUTPUT_UPDATE:
         setOutput(prev => prev + message.data);