From b4c70cd5b81fdcfeab5685b1f129e8c350e385d2 Mon Sep 17 00:00:00 2001
From: geens <niels.geens@gmail.com>
Date: Sun, 3 Nov 2024 22:44:43 +0100
Subject: [PATCH] Fix client side message handling

---
 sia/web_agent.py           |  2 +-
 sia/web_socket_manager.py  |  3 ---
 system_prompt.md           |  2 +-
 web/src/components/App.jsx | 10 +++++-----
 4 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/sia/web_agent.py b/sia/web_agent.py
index 02794e9..2a10ba1 100644
--- a/sia/web_agent.py
+++ b/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
diff --git a/sia/web_socket_manager.py b/sia/web_socket_manager.py
index bd7f3b8..ddf5332 100644
--- a/sia/web_socket_manager.py
+++ b/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:
diff --git a/system_prompt.md b/system_prompt.md
index 880d008..faeaff7 100644
--- a/system_prompt.md
+++ b/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.
\ No newline at end of file
+Open the relevant user notes when you interact with them.
diff --git a/web/src/components/App.jsx b/web/src/components/App.jsx
index b115813..d59118c 100644
--- a/web/src/components/App.jsx
+++ b/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);