From b9e13a3c7ab755fe1d03f22d3773a5cf848ca42b Mon Sep 17 00:00:00 2001
From: jerhadf <jhadfield129@gmail.com>
Date: Thu, 31 Oct 2024 02:03:19 -0400
Subject: [PATCH] Use connect_with_callbacks

---
 computer_use_demo/voice_interface.py | 66 +++++++++++++++++++---------
 1 file changed, 46 insertions(+), 20 deletions(-)

diff --git a/computer_use_demo/voice_interface.py b/computer_use_demo/voice_interface.py
index e10358d..ae0177e 100644
--- a/computer_use_demo/voice_interface.py
+++ b/computer_use_demo/voice_interface.py
@@ -65,36 +65,62 @@ def render_voice_controls(self):
                 st.write("Click 🎤 to start speaking")
 
     async def start_voice_connection(self):
-        options = ChatConnectOptions(
-            config_id=os.getenv("HUME_CONFIG_ID"),
-            secret_key=os.getenv("HUME_SECRET_KEY")
-        )
-
         try:
-            async with self.hume_client.empathic_voice.chat.connect() as socket:
+            # Initialize options for WebSocket connection
+            options = ChatConnectOptions(
+                config_id=os.getenv("HUME_CONFIG_ID"),
+                secret_key=os.getenv("HUME_SECRET_KEY")
+            )
+
+            # Connect with callbacks pattern
+            async with self.hume_client.empathic_voice.chat.connect_with_callbacks(
+                options=options,
+                on_open=self._on_socket_open,
+                on_message=self._on_socket_message,
+                on_close=self._on_socket_close,
+                on_error=self._on_socket_error
+            ) as socket:
                 self.socket = socket
 
-                # Set up message handler
-                async def on_message(msg):
-                    if msg.type == "user_message":
-                        await self.handle_voice_input(msg.message.content)
-
-                # Subscribe to socket messages
-                await socket.subscribe(on_message)
-
-                # Start microphone interface when recording
+                # Start microphone interface when recording is enabled
                 while True:
                     if st.session_state.is_recording:
-                        await MicrophoneInterface.start(
-                            socket,
-                            allow_user_interrupt=True,
-                            byte_stream=self.byte_stream
-                        )
+                        try:
+                            await MicrophoneInterface.start(
+                                socket,
+                                allow_user_interrupt=True,
+                                byte_stream=self.byte_stream
+                            )
+                        except Exception as e:
+                            st.error(f"Microphone error: {str(e)}")
+                            st.session_state.is_recording = False
                     await asyncio.sleep(0.1)
+
         except Exception as e:
             st.error(f"Voice connection error: {str(e)}")
             st.session_state.is_recording = False
 
+    async def _on_socket_open(self):
+        """Called when WebSocket connection opens"""
+        print("Voice connection opened")
+
+    async def _on_socket_message(self, message):
+        """Handle incoming WebSocket messages"""
+        if message.type == "user_message":
+            await self.handle_voice_input(message.message.content)
+        elif message.type == "audio_output":
+            # Handle audio output if needed
+            pass
+
+    async def _on_socket_close(self):
+        """Called when WebSocket connection closes"""
+        print("Voice connection closed")
+
+    async def _on_socket_error(self, error):
+        """Handle WebSocket errors"""
+        st.error(f"Voice connection error: {str(error)}")
+        st.session_state.is_recording = False
+
     async def start(self):
         self.init_streamlit_state()
         await self.start_voice_connection()
\ No newline at end of file