fixes

Author: AngeloGiacco

pyproject.toml

@@ -3,7 +3,7 @@ name = "elevenlabs"
 
 [tool.poetry]
 name = "elevenlabs"
-version = "v2.10.0"
+version = "v2.11.0"
 description = ""
 readme = "README.md"
 authors = []

src/elevenlabs/conversational_ai/conversation.py

@@ -1,18 +1,17 @@
-from abc import ABC, abstractmethod
+import asyncio
 import base64
 import json
 import threading
-from typing import Callable, Optional, Awaitable, Union, Any, Literal, Dict, Tuple
-import asyncio
+from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
+from typing import Any, Awaitable, Callable, Dict, Literal, Optional, Protocol, Tuple, Union
 
-from websockets.sync.client import connect, Connection
 import websockets
-from websockets.exceptions import ConnectionClosedOK
-
 from ..base_client import BaseElevenLabs
 from ..version import __version__
+from websockets.exceptions import ConnectionClosedOK
+from websockets.sync.client import Connection, connect
 
 
 class ClientToOrchestratorEvent(str, Enum):
@@ -276,6 +275,71 @@ def __init__(
         self.user_id = user_id
 
 
+class MessageHandler(Protocol):
+    """Protocol defining the interface for message handlers."""
+    
+    callback_agent_response: Optional[Callable]
+    callback_agent_response_correction: Optional[Callable] 
+    callback_user_transcript: Optional[Callable]
+    callback_latency_measurement: Optional[Callable]
+    
+    def handle_audio_output(self, audio: bytes) -> Union[None, Awaitable[None]]:
+        """Handle audio output."""
+        ...
+    
+    def handle_agent_response(self, response: str) -> Union[None, Awaitable[None]]:
+        """Handle agent response."""
+        ...
+    
+    def handle_agent_response_correction(self, original: str, corrected: str) -> Union[None, Awaitable[None]]:
+        """Handle agent response correction."""
+        ...
+        
+    def handle_user_transcript(self, transcript: str) -> Union[None, Awaitable[None]]:
+        """Handle user transcript."""
+        ...
+        
+    def handle_interruption(self) -> Union[None, Awaitable[None]]:
+        """Handle interruption."""
+        ...
+        
+    def handle_ping(self, event: Dict[str, Any]) -> Union[None, Awaitable[None]]:
+        """Handle ping event."""
+        ...
+        
+    def handle_latency_measurement(self, latency: int) -> Union[None, Awaitable[None]]:
+        """Handle latency measurement."""
+        ...
+        
+    def handle_client_tool_call(self, tool_name: str, parameters: Dict[str, Any]) -> None:
+        """Handle client tool call."""
+        ...
+
+
+class BaseMessageHandler:
+    """Base implementation for message handlers with common functionality."""
+    
+    def __init__(self, conversation, ws_or_websocket):
+        self.conversation = conversation
+        self.ws = ws_or_websocket
+        self.callback_agent_response = conversation.callback_agent_response
+        self.callback_agent_response_correction = conversation.callback_agent_response_correction
+        self.callback_user_transcript = conversation.callback_user_transcript
+        self.callback_latency_measurement = conversation.callback_latency_measurement
+    
+    def handle_client_tool_call(self, tool_name: str, parameters: Dict[str, Any]) -> None:
+        """Handle client tool call - common implementation for both sync and async."""
+        def send_response(response):
+            if not self.conversation._should_stop.is_set():
+                self._send_response(response)
+        
+        self.conversation.client_tools.execute_tool(tool_name, parameters, send_response)
+    
+    def _send_response(self, response: Dict[str, Any]) -> None:
+        """Send response - to be implemented by subclasses."""
+        raise NotImplementedError
+
+
 class BaseConversation:
     """Base class for conversation implementations with shared parameters and logic."""
 
@@ -300,6 +364,42 @@ def __init__(
 
         self._conversation_id = None
         self._last_interrupt_id = 0
+    
+    def _create_sync_audio_callback(self, ws) -> Callable[[bytes], None]:
+        """Create sync audio input callback."""
+        def callback(audio: bytes) -> None:
+            try:
+                ws.send(
+                    json.dumps({
+                        "user_audio_chunk": base64.b64encode(audio).decode(),
+                    })
+                )
+            except ConnectionClosedOK:
+                self.end_session()
+            except Exception as e:
+                print(f"Error sending user audio chunk: {e}")
+                self.end_session()
+        return callback
+    
+    def _create_async_audio_callback(self, ws) -> Callable[[bytes], Awaitable[None]]:
+        """Create async audio input callback."""
+        async def callback(audio: bytes) -> None:
+            try:
+                await ws.send(
+                    json.dumps({
+                        "user_audio_chunk": base64.b64encode(audio).decode(),
+                    })
+                )
+            except ConnectionClosedOK:
+                await self.end_session()
+            except Exception as e:
+                print(f"Error sending user audio chunk: {e}")
+                await self.end_session()
+        return callback
+    
+    def _handle_connection_closed(self) -> Union[None, Awaitable[None]]:
+        """Handle WebSocket connection closed - to be implemented by subclasses."""
+        raise NotImplementedError
 
     def _get_wss_url(self):
         base_ws_url = self.client._client_wrapper.get_environment().wss
@@ -327,7 +427,7 @@ def _create_initiation_message(self):
             }
         )
 
-    def _handle_message_core(self, message, message_handler):
+    def _handle_message_core(self, message: Dict[str, Any], message_handler: MessageHandler) -> None:
         """Core message handling logic shared between sync and async implementations.
         
         Args:
@@ -383,7 +483,7 @@ def _handle_message_core(self, message, message_handler):
         else:
             pass  # Ignore all other message types.
 
-    async def _handle_message_core_async(self, message, message_handler):
+    async def _handle_message_core_async(self, message: Dict[str, Any], message_handler: MessageHandler) -> None:
         """Async wrapper for core message handling logic."""
         if message["type"] == "conversation_initiation_metadata":
             event = message["conversation_initiation_metadata_event"]
@@ -592,35 +692,24 @@ def send_contextual_update(self, text: str):
             print(f"Error sending contextual update: {e}")
             raise
 
+    def _handle_connection_closed(self) -> None:
+        self.end_session()
+    
     def _run(self, ws_url: str):
         with connect(ws_url, max_size=16 * 1024 * 1024) as ws:
             self._ws = ws
             ws.send(self._create_initiation_message())
-            self._ws = ws
-
-            def input_callback(audio):
-                try:
-                    ws.send(
-                        json.dumps(
-                            {
-                                "user_audio_chunk": base64.b64encode(audio).decode(),
-                            }
-                        )
-                    )
-                except ConnectionClosedOK:
-                    self.end_session()
-                except Exception as e:
-                    print(f"Error sending user audio chunk: {e}")
-                    self.end_session()
-
+            
+            input_callback = self._create_sync_audio_callback(ws)
             self.audio_interface.start(input_callback)
+            
             while not self._should_stop.is_set():
                 try:
                     message = json.loads(ws.recv(timeout=0.5))
                     if self._should_stop.is_set():
                         return
                     self._handle_message(message, ws)
-                except ConnectionClosedOK as e:
+                except ConnectionClosedOK:
                     self.end_session()
                 except TimeoutError:
                     pass
@@ -631,31 +720,23 @@ def input_callback(audio):
             self._ws = None
 
     def _handle_message(self, message, ws):
-        class SyncMessageHandler:
-            def __init__(self, conversation, ws):
-                self.conversation = conversation
-                self.ws = ws
-                self.callback_agent_response = conversation.callback_agent_response
-                self.callback_agent_response_correction = conversation.callback_agent_response_correction
-                self.callback_user_transcript = conversation.callback_user_transcript
-                self.callback_latency_measurement = conversation.callback_latency_measurement
-            
-            def handle_audio_output(self, audio):
+        class SyncMessageHandler(BaseMessageHandler):
+            def handle_audio_output(self, audio: bytes) -> None:
                 self.conversation.audio_interface.output(audio)
 
-            def handle_agent_response(self, response):
+            def handle_agent_response(self, response: str) -> None:
                 self.conversation.callback_agent_response(response)
 
-            def handle_agent_response_correction(self, original, corrected):
+            def handle_agent_response_correction(self, original: str, corrected: str) -> None:
                 self.conversation.callback_agent_response_correction(original, corrected)
 
-            def handle_user_transcript(self, transcript):
+            def handle_user_transcript(self, transcript: str) -> None:
                 self.conversation.callback_user_transcript(transcript)
 
-            def handle_interruption(self):
+            def handle_interruption(self) -> None:
                 self.conversation.audio_interface.interrupt()
 
-            def handle_ping(self, event):
+            def handle_ping(self, event: Dict[str, Any]) -> None:
                 self.ws.send(
                     json.dumps(
                         {
@@ -665,15 +746,11 @@ def handle_ping(self, event):
                     )
                 )
 
-            def handle_latency_measurement(self, latency):
+            def handle_latency_measurement(self, latency: int) -> None:
                 self.conversation.callback_latency_measurement(latency)
 
-            def handle_client_tool_call(self, tool_name, parameters):
-                def send_response(response):
-                    if not self.conversation._should_stop.is_set():
-                        self.ws.send(json.dumps(response))
-                
-                self.conversation.client_tools.execute_tool(tool_name, parameters, send_response)
+            def _send_response(self, response: Dict[str, Any]) -> None:
+                self.ws.send(json.dumps(response))
 
         handler = SyncMessageHandler(self, ws)
         self._handle_message_core(message, handler)
@@ -761,6 +838,14 @@ async def end_session(self):
         self.client_tools.stop()
         self._ws = None
         self._should_stop.set()
+        
+        # Cleanup the background task
+        if self._task and not self._task.done():
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
 
         if self.callback_end_session:
             await self.callback_end_session()
@@ -836,26 +921,15 @@ async def send_contextual_update(self, text: str):
             print(f"Error sending contextual update: {e}")
             raise
 
+    async def _handle_connection_closed(self) -> None:
+        await self.end_session()
+    
     async def _run(self, ws_url: str):
         async with websockets.connect(ws_url, max_size=16 * 1024 * 1024) as ws:
             self._ws = ws
             await ws.send(self._create_initiation_message())
-
-            async def input_callback(audio):
-                try:
-                    await ws.send(
-                        json.dumps(
-                            {
-                                "user_audio_chunk": base64.b64encode(audio).decode(),
-                            }
-                        )
-                    )
-                except ConnectionClosedOK:
-                    await self.end_session()
-                except Exception as e:
-                    print(f"Error sending user audio chunk: {e}")
-                    await self.end_session()
-
+            
+            input_callback = self._create_async_audio_callback(ws)
             await self.audio_interface.start(input_callback)
 
             try:
@@ -879,31 +953,23 @@ async def input_callback(audio):
                 self._ws = None
 
     async def _handle_message(self, message, ws):
-        class AsyncMessageHandler:
-            def __init__(self, conversation, ws):
-                self.conversation = conversation
-                self.ws = ws
-                self.callback_agent_response = conversation.callback_agent_response
-                self.callback_agent_response_correction = conversation.callback_agent_response_correction
-                self.callback_user_transcript = conversation.callback_user_transcript
-                self.callback_latency_measurement = conversation.callback_latency_measurement
-            
-            async def handle_audio_output(self, audio):
+        class AsyncMessageHandler(BaseMessageHandler):
+            async def handle_audio_output(self, audio: bytes) -> None:
                 await self.conversation.audio_interface.output(audio)
 
-            async def handle_agent_response(self, response):
+            async def handle_agent_response(self, response: str) -> None:
                 await self.conversation.callback_agent_response(response)
 
-            async def handle_agent_response_correction(self, original, corrected):
+            async def handle_agent_response_correction(self, original: str, corrected: str) -> None:
                 await self.conversation.callback_agent_response_correction(original, corrected)
 
-            async def handle_user_transcript(self, transcript):
+            async def handle_user_transcript(self, transcript: str) -> None:
                 await self.conversation.callback_user_transcript(transcript)
 
-            async def handle_interruption(self):
+            async def handle_interruption(self) -> None:
                 await self.conversation.audio_interface.interrupt()
 
-            async def handle_ping(self, event):
+            async def handle_ping(self, event: Dict[str, Any]) -> None:
                 await self.ws.send(
                     json.dumps(
                         {
@@ -913,15 +979,11 @@ async def handle_ping(self, event):
                     )
                 )
 
-            async def handle_latency_measurement(self, latency):
+            async def handle_latency_measurement(self, latency: int) -> None:
                 await self.conversation.callback_latency_measurement(latency)
 
-            def handle_client_tool_call(self, tool_name, parameters):
-                def send_response(response):
-                    if not self.conversation._should_stop.is_set():
-                        asyncio.create_task(self.ws.send(json.dumps(response)))
-                
-                self.conversation.client_tools.execute_tool(tool_name, parameters, send_response)
+            def _send_response(self, response: Dict[str, Any]) -> None:
+                asyncio.create_task(self.ws.send(json.dumps(response)))
 
         handler = AsyncMessageHandler(self, ws)