ojin-client 0.1.7.dev8__tar.gz → 0.1.7.dev10__tar.gz

{ojin_client-0.1.7.dev8 → ojin_client-0.1.7.dev10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ojin-client
-Version: 0.1.7.dev8
+Version: 0.1.7.dev10
 Summary: Ojin platform services
 Author: Journee
 License: Apache-2.0

{ojin_client-0.1.7.dev8 → ojin_client-0.1.7.dev10}/ojin/ojin_persona_client.py

@@ -1,4 +1,4 @@
-"""WebSocket client for OJIN Persona service."""
+"""WebSocket client for OJIN Persona service with optimized cancellation handling."""
 
 import asyncio
 import contextlib
@@ -87,9 +87,12 @@ class OjinPersonaClient(IOjinPersonaClient):
         self._receive_task: Optional[asyncio.Task] = None
         self._inference_server_ready: bool = False
         self._cancelled: bool = False
-        self.active_interaction_id: str | None = None
+        self._active_interaction_id: str | None = None
         self._split_audio_task: Optional[asyncio.Task] = None
         self._audio_queue: asyncio.Queue[OjinPersonaInteractionInputMessage] = asyncio.Queue()
+        
+        # Add cancellation event for immediate stopping
+        self._cancel_event = asyncio.Event()
 
     async def connect(self) -> None:
         """Establish WebSocket connection and authenticate with the service."""
@@ -135,7 +138,8 @@ class OjinPersonaClient(IOjinPersonaClient):
             pass
 
         self._running = False
-        self.active_interaction_id = None
+        self._active_interaction_id = None
+        self._cancel_event.set()  # Signal cancellation to all tasks
 
         if self._ws:
             try:
@@ -156,7 +160,6 @@ class OjinPersonaClient(IOjinPersonaClient):
                 await self._receive_task
             self._receive_task = None   
 
-
         logger.info("Disconnected from OJIN Persona service")
 
     async def _receive_messages(self) -> None:
@@ -200,8 +203,9 @@ class OjinPersonaClient(IOjinPersonaClient):
                     )
                     logger.debug("Received InteractionResponse for id %s", interaction_response.interaction_id)
                     
-                    if interaction_response.interaction_id == self._cancelled_interaction_id:
-                        logger.warning("Message From old interaction")
+                    # TODO: Possibly want to delete
+                    if interaction_response.interaction_id != self._active_interaction_id:
+                        logger.warning("Message From other interaction")
                         return
                     await self._message_queue.put(interaction_response)
                     return
@@ -282,73 +286,55 @@ class OjinPersonaClient(IOjinPersonaClient):
             raise ConnectionError("Infernece Server is not ready to receive messsages")
 
         if isinstance(message, OjinPersonaCancelInteractionMessage):
-            logger.info("Interrupt")
+            logger.info("Interrupt - Processing cancellation immediately")
             
+            # Set cancellation flag and event immediately
             self._cancelled = True
+            self._cancel_event.set()
+            
+            # Send cancellation message with high priority
             cancel_input = CancelInteractionMessage(
                     payload=message.to_proxy_message()
             )
 
-            await self._ws.send(cancel_input.model_dump_json())
-
-            logger.info(f"Message sent {message.interaction_id}")
+            # Send immediately without waiting
+            try:
+                await self._ws.send(cancel_input.model_dump_json())
+                logger.info(f"Cancellation message sent immediately for {message.interaction_id}")
+            except Exception as e:
+                logger.error(f"Failed to send cancellation message: {e}")
 
-            while not self._message_queue.empty():
-                try:
-                    self._message_queue.get_nowait()
-                except asyncio.QueueEmpty:
-                    break
+            # Clear queues quickly without blocking
+            self._clear_queues_non_blocking()
 
+            # Reset cancellation state
             self._cancelled = False
+            self._cancel_event.clear()
 
             return
 
         if isinstance(message, StartInteractionMessage):
             interaction_id = str(uuid.uuid4())
-            self.active_interaction_id = interaction_id
+            self._active_interaction_id = interaction_id
             logger.info("Generate UUID %s", interaction_id)
             interaction_response = StartInteractionResponseMessage(
                 interaction_id=interaction_id
             )
-            while not self._message_queue.empty():
-                await self._message_queue.get()
+            # Clear queues non-blocking
+            self._clear_queues_non_blocking()
             self._message_queue.put_nowait(interaction_response)
             return
 
         if isinstance(message, OjinPersonaInteractionInputMessage):
             logger.info("InteractionMessage")            
             logger.info(f"Message sent {message.interaction_id}")
-            if message.interaction_id != self.active_interaction_id:
+            if message.interaction_id != self._active_interaction_id:
                 return
 
             if not message.audio_int16_bytes:
                 raise ValueError("Audio cannot be empty")
             
             await self._audio_queue.put(message)
-            # Split audio bytes into chunks of max 3200 samples
-            # max_chunk_size = 3200 * 2
-            # audio_chunks = [
-            #     message.audio_int16_bytes[i : i + max_chunk_size]
-            #     for i in range(0, len(message.audio_int16_bytes), max_chunk_size)
-            # ]
-            # logger.info(
-            #     "Split audio into %d chunks of max %d bytes",
-            #     len(audio_chunks), max_chunk_size
-            # )
-
-            # for i, chunk in enumerate(audio_chunks):
-            #     is_last = i == len(audio_chunks) - 1 and message.is_last_input
-            #
-            #     interaction_input = InteractionInput(
-            #         interaction_id=message.interaction_id,
-            #         is_final_input=is_last,
-            #         payload_type="audio",
-            #         payload=chunk,
-            #         timestamp=int(time.monotonic() * 1000),
-            #         params=message.params if i == 0 else None,
-            #     )
-            #     proxy_message = InteractionInputMessage(payload=interaction_input)
-            #     await self._ws.send(proxy_message.to_bytes())
             return
 
         logger.error("The message %s is Unknown", message)
@@ -364,18 +350,73 @@ class OjinPersonaClient(IOjinPersonaClient):
         )
         raise Exception(error)
 
-    async def _split_audio(self) -> None:
+    def _clear_queues_non_blocking(self) -> None:
+        """Clear all queues without blocking."""
+        # Clear message queue
         while True:
-            message_audio: OjinPersonaInteractionInputMessage| None = None
-            if self._cancelled:
-                continue
+            try:
+                self._message_queue.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+        
+        # Clear audio queue
+        while True:
+            try:
+                self._audio_queue.get_nowait()
+            except asyncio.QueueEmpty:
+                break
 
+    async def _split_audio(self) -> None:
+        """Split audio into chunks and send them, with cancellation support."""
+        while self._running:
+            message_audio: OjinPersonaInteractionInputMessage | None = None
+            
             try:
-                message_audio = self._audio_queue.get_nowait()
+                # Use wait_for with cancellation event to make this interruptible
+                wait_tasks = [
+                    asyncio.create_task(self._audio_queue.get()),
+                    asyncio.create_task(self._cancel_event.wait())
+                ]
+                
+                done, pending = await asyncio.wait(
+                    wait_tasks, 
+                    return_when=asyncio.FIRST_COMPLETED,
+                    timeout=0.1  # Short timeout to check cancellation frequently
+                )
+                
+                # Cancel pending tasks
+                for task in pending:
+                    task.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await task
+                
+                # Check if cancellation was triggered
+                if self._cancelled or self._cancel_event.is_set():
+                    logger.info("Audio splitting cancelled")
+                    continue
+                
+                # Check if we got a message
+                if done:
+                    completed_task = done.pop()
+                    if completed_task == wait_tasks[0]:  # Audio queue task completed
+                        message_audio = completed_task.result()
+                    else:  # Cancellation event was set
+                        continue
+                else:
+                    # Timeout occurred, continue loop
+                    continue
+                    
             except asyncio.QueueEmpty:
                 await asyncio.sleep(0.01)
                 continue
+            except Exception as e:
+                logger.error(f"Error getting audio message: {e}")
+                continue
 
+            if not message_audio:
+                continue
+
+            # Process audio chunks with cancellation checks
             max_chunk_size = 3200 * 2
             audio_chunks = [
                 message_audio.audio_int16_bytes[i : i + max_chunk_size]
@@ -387,6 +428,11 @@ class OjinPersonaClient(IOjinPersonaClient):
             )
 
             for i, chunk in enumerate(audio_chunks):
+                # Check for cancellation before each chunk
+                if self._cancelled or self._cancel_event.is_set():
+                    logger.info("Audio chunk sending cancelled")
+                    break
+                    
                 is_last = i == len(audio_chunks) - 1 and message_audio.is_last_input
 
                 interaction_input = InteractionInput(
@@ -399,8 +445,11 @@ class OjinPersonaClient(IOjinPersonaClient):
                 )
                 proxy_message = InteractionInputMessage(payload=interaction_input)
 
-                await self._ws.send(proxy_message.to_bytes())
-
+                try:
+                    await self._ws.send(proxy_message.to_bytes())
+                except Exception as e:
+                    logger.error(f"Failed to send audio chunk: {e}")
+                    break
 
     async def receive_message(self) -> BaseModel | None:
         """Receive the next message from the OJIN Persona service.

{ojin_client-0.1.7.dev8 → ojin_client-0.1.7.dev10}/ojin_client.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ojin-client
-Version: 0.1.7.dev8
+Version: 0.1.7.dev10
 Summary: Ojin platform services
 Author: Journee
 License: Apache-2.0

{ojin_client-0.1.7.dev8 → ojin_client-0.1.7.dev10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ojin-client"
-version = "0.1.7dev8"
+version = "0.1.7dev10"
 description = "Ojin platform services"
 readme = "README.md"
 requires-python = ">=3.10"