From adaac003e544dec6107b244b6a2b2db4ebf8420e Mon Sep 17 00:00:00 2001
From: Kunal Shah <kunal.shah@cartesia.ai>
Date: Mon, 16 Sep 2024 15:59:06 -0700
Subject: [PATCH] [Cartesia] Fix streaming truncation bug with Twilio Fast API
 WS

---
 src/pipecat/services/cartesia.py              | 52 +++++++++++--------
 .../transports/network/fastapi_websocket.py   | 10 ++--
 2 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/src/pipecat/services/cartesia.py b/src/pipecat/services/cartesia.py
index ac02ea469..6713a3825 100644
--- a/src/pipecat/services/cartesia.py
+++ b/src/pipecat/services/cartesia.py
@@ -4,33 +4,24 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
-import json
-import uuid
-import base64
 import asyncio
+import base64
+import io
+import json
 import time
-
+import uuid
 from typing import AsyncGenerator
 
-from pipecat.frames.frames import (
-    CancelFrame,
-    ErrorFrame,
-    Frame,
-    AudioRawFrame,
-    StartInterruptionFrame,
-    StartFrame,
-    EndFrame,
-    TTSStartedFrame,
-    TTSStoppedFrame,
-    TextFrame,
-    LLMFullResponseEndFrame
-)
-from pipecat.processors.frame_processor import FrameDirection
-from pipecat.transcriptions.language import Language
-from pipecat.services.ai_services import AsyncWordTTSService
-
 from loguru import logger
 
+from pipecat.frames.frames import (AudioRawFrame, CancelFrame, EndFrame,
+                                   ErrorFrame, Frame, LLMFullResponseEndFrame,
+                                   StartFrame, StartInterruptionFrame,
+                                   TextFrame, TTSStartedFrame, TTSStoppedFrame)
+from pipecat.processors.frame_processor import FrameDirection
+from pipecat.services.ai_services import AsyncWordTTSService
+from pipecat.transcriptions.language import Language
+
 # See .env.example for Cartesia configuration needed
 try:
     import websockets
@@ -161,6 +152,25 @@ class CartesiaTTSService(AsyncWordTTSService):
         await self.push_frame(LLMFullResponseEndFrame())
         self._context_id = None
 
+    async def flush_audio(self):
+        if not self._context_id or not self._websocket:
+            return
+        logger.debug("Flushing audio")
+        msg = {
+                "transcript": "",
+                "continue": False,
+                "context_id": self._context_id,
+                "model_id": self._model_id,
+                "voice": {
+                    "mode": "id",
+                    "id": self._voice_id
+                },
+                "output_format": self._output_format,
+                "language": self._language,
+                "add_timestamps": True,
+            }
+        await self._websocket.send(json.dumps(msg))
+
     async def _receive_task_handler(self):
         try:
             async for message in self._websocket:
diff --git a/src/pipecat/transports/network/fastapi_websocket.py b/src/pipecat/transports/network/fastapi_websocket.py
index 16c5e81fc..b5bcfe7ed 100644
--- a/src/pipecat/transports/network/fastapi_websocket.py
+++ b/src/pipecat/transports/network/fastapi_websocket.py
@@ -8,19 +8,19 @@
 import asyncio
 import io
 import wave
-
 from typing import Awaitable, Callable
+
+from loguru import logger
 from pydantic.main import BaseModel
 
-from pipecat.frames.frames import AudioRawFrame, CancelFrame, EndFrame, Frame, StartFrame, StartInterruptionFrame
+from pipecat.frames.frames import (AudioRawFrame, CancelFrame, EndFrame, Frame,
+                                   StartFrame, StartInterruptionFrame)
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
 from pipecat.serializers.base_serializer import FrameSerializer
 from pipecat.transports.base_input import BaseInputTransport
 from pipecat.transports.base_output import BaseOutputTransport
 from pipecat.transports.base_transport import BaseTransport, TransportParams
 
-from loguru import logger
-
 try:
     from fastapi import WebSocket
     from starlette.websockets import WebSocketState
@@ -101,7 +101,7 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
 
     async def write_raw_audio_frames(self, frames: bytes):
         self._websocket_audio_buffer += frames
-        while len(self._websocket_audio_buffer) >= self._params.audio_frame_size:
+        while len(self._websocket_audio_buffer):
             frame = AudioRawFrame(
                 audio=self._websocket_audio_buffer[:
                                                    self._params.audio_frame_size],