livekit-plugins-cartesia 0.2.0.dev7__py3-none-any.whl → 0.3.0__py3-none-any.whl

livekit/plugins/cartesia/__init__.py

@@ -19,13 +19,12 @@ __all__ = ["TTS", "ChunkedStream", "__version__"]
 
 from livekit.agents import Plugin
 
+from .log import logger
+
 
 class CartesiaPlugin(Plugin):
     def __init__(self):
-        super().__init__(__name__, __version__, __package__)
-
-    def download_files(self):
-        pass
+        super().__init__(__name__, __version__, __package__, logger)
 
 
 Plugin.register_plugin(CartesiaPlugin())

livekit/plugins/cartesia/models.py

@@ -11,4 +11,4 @@ TTSEncoding = Literal[
 
 TTSModels = Literal["sonic-english", "sonic-multilingual"]
 TTSLanguages = Literal["en", "es", "fr", "de", "pt", "zh", "ja"]
-TTSDefaultVoiceId = "248be419-c632-4f23-adf1-5324ed7dbf1d"
+TTSDefaultVoiceId = "b7d50908-b17c-442d-ad8d-810c63997ed9"

livekit/plugins/cartesia/tts.py

@@ -14,11 +14,14 @@
 
 from __future__ import annotations
 
+import asyncio
+import base64
+import json
 import os
 from dataclasses import dataclass
 
 import aiohttp
-from livekit.agents import tts, utils
+from livekit.agents import tokenize, tts, utils
 
 from .log import logger
 from .models import TTSDefaultVoiceId, TTSEncoding, TTSModels
@@ -27,6 +30,9 @@ API_AUTH_HEADER = "X-API-Key"
 API_VERSION_HEADER = "Cartesia-Version"
 API_VERSION = "2024-06-10"
 
+NUM_CHANNELS = 1
+BUFFERED_WORDS_COUNT = 8
+
 
 @dataclass
 class _TTSOptions:
@@ -36,6 +42,7 @@ class _TTSOptions:
     voice: str | list[float]
     api_key: str
     language: str
+    word_tokenizer: tokenize.WordTokenizer
 
 
 class TTS(tts.TTS):
@@ -49,11 +56,14 @@ class TTS(tts.TTS):
         sample_rate: int = 24000,
         api_key: str | None = None,
         http_session: aiohttp.ClientSession | None = None,
+        word_tokenizer: tokenize.WordTokenizer = tokenize.basic.WordTokenizer(
+            ignore_punctuation=False
+        ),
     ) -> None:
         super().__init__(
-            capabilities=tts.TTSCapabilities(streaming=False),
+            capabilities=tts.TTSCapabilities(streaming=True),
             sample_rate=sample_rate,
-            num_channels=1,
+            num_channels=NUM_CHANNELS,
         )
 
         api_key = api_key or os.environ.get("CARTESIA_API_KEY")
@@ -67,6 +77,7 @@ class TTS(tts.TTS):
             sample_rate=sample_rate,
             voice=voice,
             api_key=api_key,
+            word_tokenizer=word_tokenizer,
         )
         self._session = http_session
 
@@ -79,8 +90,13 @@ class TTS(tts.TTS):
     def synthesize(self, text: str) -> "ChunkedStream":
         return ChunkedStream(text, self._opts, self._ensure_session())
 
+    def stream(self) -> "SynthesizeStream":
+        return SynthesizeStream(self._opts, self._ensure_session())
+
 
 class ChunkedStream(tts.ChunkedStream):
+    """Synthesize chunked text using the bytes endpoint"""
+
     def __init__(
         self, text: str, opts: _TTSOptions, session: aiohttp.ClientSession
     ) -> None:
@@ -90,35 +106,17 @@ class ChunkedStream(tts.ChunkedStream):
     @utils.log_exceptions(logger=logger)
     async def _main_task(self):
         bstream = utils.audio.AudioByteStream(
-            sample_rate=self._opts.sample_rate, num_channels=1
+            sample_rate=self._opts.sample_rate, num_channels=NUM_CHANNELS
         )
-        request_id = utils.shortuuid()
-        segment_id = utils.shortuuid()
-
-        voice = {}
-        if isinstance(self._opts.voice, str):
-            voice["mode"] = "id"
-            voice["id"] = self._opts.voice
-        else:
-            voice["mode"] = "embedding"
-            voice["embedding"] = self._opts.voice
-
-        data = {
-            "model_id": self._opts.model,
-            "transcript": self._text,
-            "voice": voice,
-            "output_format": {
-                "container": "raw",
-                "encoding": self._opts.encoding,
-                "sample_rate": self._opts.sample_rate,
-            },
-            "language": self._opts.language,
-        }
+        request_id, segment_id = utils.shortuuid(), utils.shortuuid()
+
+        data = _to_cartesia_options(self._opts)
+        data["transcript"] = self._text
 
         async with self._session.post(
             "https://api.cartesia.ai/tts/bytes",
             headers={
-                API_AUTH_HEADER: f"{self._opts.api_key}",
+                API_AUTH_HEADER: self._opts.api_key,
                 API_VERSION_HEADER: API_VERSION,
             },
             json=data,
@@ -137,3 +135,169 @@ class ChunkedStream(tts.ChunkedStream):
                         request_id=request_id, segment_id=segment_id, frame=frame
                     )
                 )
+
+
+class SynthesizeStream(tts.SynthesizeStream):
+    def __init__(
+        self,
+        opts: _TTSOptions,
+        session: aiohttp.ClientSession,
+    ):
+        super().__init__()
+        self._opts, self._session = opts, session
+        self._buf = ""
+
+    @utils.log_exceptions(logger=logger)
+    async def _main_task(self) -> None:
+        retry_count = 0
+        max_retry = 3
+        while self._input_ch.qsize() or not self._input_ch.closed:
+            try:
+                url = f"wss://api.cartesia.ai/tts/websocket?api_key={self._opts.api_key}&cartesia_version={API_VERSION}"
+                ws = await self._session.ws_connect(url)
+                retry_count = 0  # connected successfully, reset the retry_count
+
+                await self._run_ws(ws)
+            except Exception as e:
+                if retry_count >= max_retry:
+                    logger.exception(
+                        f"failed to connect to Cartesia after {max_retry} tries"
+                    )
+                    break
+
+                retry_delay = min(retry_count * 2, 10)  # max 10s
+                retry_count += 1
+
+                logger.warning(
+                    f"Cartesia connection failed, retrying in {retry_delay}s",
+                    exc_info=e,
+                )
+                await asyncio.sleep(retry_delay)
+
+    async def _run_ws(self, ws: aiohttp.ClientWebSocketResponse) -> None:
+        request_id = utils.shortuuid()
+        pending_segments = []
+
+        async def send_task():
+            base_pkt = _to_cartesia_options(self._opts)
+
+            def _new_segment():
+                segment_id = utils.shortuuid()
+                pending_segments.append(segment_id)
+                return segment_id
+
+            current_segment_id: str | None = _new_segment()
+
+            async for data in self._input_ch:
+                if isinstance(data, self._FlushSentinel):
+                    if current_segment_id is None:
+                        continue
+
+                    end_pkt = base_pkt.copy()
+                    end_pkt["context_id"] = current_segment_id
+                    end_pkt["transcript"] = self._buf + " "
+                    end_pkt["continue"] = False
+                    await ws.send_str(json.dumps(end_pkt))
+
+                    current_segment_id = None
+                    self._buf = ""
+                elif data:
+                    if current_segment_id is None:
+                        current_segment_id = _new_segment()
+
+                    self._buf += data
+                    words = self._opts.word_tokenizer.tokenize(text=self._buf)
+                    if len(words) < BUFFERED_WORDS_COUNT + 1:
+                        continue
+
+                    data = self._opts.word_tokenizer.format_words(words[:-1]) + " "
+                    self._buf = words[-1]
+
+                    token_pkt = base_pkt.copy()
+                    token_pkt["context_id"] = current_segment_id
+                    token_pkt["transcript"] = data
+                    token_pkt["continue"] = True
+                    await ws.send_str(json.dumps(token_pkt))
+
+            if len(pending_segments) == 0:
+                await ws.close()
+
+        async def recv_task():
+            audio_bstream = utils.audio.AudioByteStream(
+                sample_rate=self._opts.sample_rate,
+                num_channels=NUM_CHANNELS,
+            )
+
+            while True:
+                msg = await ws.receive()
+                if msg.type in (
+                    aiohttp.WSMsgType.CLOSED,
+                    aiohttp.WSMsgType.CLOSE,
+                    aiohttp.WSMsgType.CLOSING,
+                ):
+                    raise Exception("Cartesia connection closed unexpectedly")
+
+                if msg.type != aiohttp.WSMsgType.TEXT:
+                    logger.warning("unexpected Cartesia message type %s", msg.type)
+                    continue
+
+                data = json.loads(msg.data)
+                segment_id = data.get("context_id")
+                if data.get("data"):
+                    b64data = base64.b64decode(data["data"])
+                    for frame in audio_bstream.write(b64data):
+                        self._event_ch.send_nowait(
+                            tts.SynthesizedAudio(
+                                request_id=request_id,
+                                segment_id=segment_id,
+                                frame=frame,
+                            )
+                        )
+                elif data.get("done"):
+                    for frame in audio_bstream.flush():
+                        self._event_ch.send_nowait(
+                            tts.SynthesizedAudio(
+                                request_id=request_id,
+                                segment_id=segment_id,
+                                frame=frame,
+                            )
+                        )
+
+                    pending_segments.remove(segment_id)
+                    if len(pending_segments) == 0 and self._input_ch.closed:
+                        # we're not going to receive more frames, close the connection
+                        await ws.close()
+                        break
+                else:
+                    logger.error("unexpected Cartesia message %s", data)
+
+        tasks = [
+            asyncio.create_task(send_task()),
+            asyncio.create_task(recv_task()),
+        ]
+
+        try:
+            await asyncio.gather(*tasks)
+        finally:
+            await utils.aio.gracefully_cancel(*tasks)
+
+
+def _to_cartesia_options(opts: _TTSOptions) -> dict:
+    voice: dict = {}
+    if isinstance(opts.voice, str):
+        voice["mode"] = "id"
+        voice["id"] = opts.voice
+    else:
+        voice["mode"] = "embedding"
+        voice["embedding"] = opts.voice
+
+    return {
+        "model_id": opts.model,
+        "voice": voice,
+        "output_format": {
+            "container": "raw",
+            "encoding": opts.encoding,
+            "sample_rate": opts.sample_rate,
+        },
+        "language": opts.language,
+    }

livekit/plugins/cartesia/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.2.0-dev.7"
+__version__ = "0.3.0"

{livekit_plugins_cartesia-0.2.0.dev7.dist-info → livekit_plugins_cartesia-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: livekit-plugins-cartesia
-Version: 0.2.0.dev7
+Version: 0.3.0
 Summary: LiveKit Agents Plugin for Cartesia
 Home-page: https://github.com/livekit/agents
 License: Apache-2.0

livekit_plugins_cartesia-0.3.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+livekit/plugins/cartesia/__init__.py,sha256=BUfWY_evL5dUHn9hBDQVor6ssctDKQfbQfZy5SWndN8,926
+livekit/plugins/cartesia/log.py,sha256=4Mnhjng_DU1dIWP9IWjIQGZ67EV3LnQhWMWCHVudJbo,71
+livekit/plugins/cartesia/models.py,sha256=T1iPQ18h4-o5rgSW236PDc73qp5zR9k4r_qNCl3XPWc,335
+livekit/plugins/cartesia/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+livekit/plugins/cartesia/tts.py,sha256=uklD9fIYL8QWUSiyypFDgflkie9VhTu1C-x4YwJcDCU,10283
+livekit/plugins/cartesia/version.py,sha256=G5iYozum4q7UpHwW43F7QfhzUfwcncPxBZ0gmUGsd5I,600
+livekit_plugins_cartesia-0.3.0.dist-info/METADATA,sha256=iJcOyrkQ-0yPK_lYtR-eEbIDav84xlN7DUvwncx7OpQ,1252
+livekit_plugins_cartesia-0.3.0.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+livekit_plugins_cartesia-0.3.0.dist-info/top_level.txt,sha256=OoDok3xUmXbZRvOrfvvXB-Juu4DX79dlq188E19YHoo,8
+livekit_plugins_cartesia-0.3.0.dist-info/RECORD,,

{livekit_plugins_cartesia-0.2.0.dev7.dist-info → livekit_plugins_cartesia-0.3.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (71.1.0)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

livekit_plugins_cartesia-0.2.0.dev7.dist-info/RECORD

@@ -1,10 +0,0 @@
-livekit/plugins/cartesia/__init__.py,sha256=_a8u7qqya1pjZTV19gNOpMKTO7ccAVZAeCukiDKAG-U,937
-livekit/plugins/cartesia/log.py,sha256=4Mnhjng_DU1dIWP9IWjIQGZ67EV3LnQhWMWCHVudJbo,71
-livekit/plugins/cartesia/models.py,sha256=06S-Z-M90kB-kEOQsQk70xfQUD-TztU4ZIU_AfAyUMc,335
-livekit/plugins/cartesia/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-livekit/plugins/cartesia/tts.py,sha256=S5BMSVtsbNI_c2PpgyFK6wvleudmJZLTUt3ZmGNKlRI,4319
-livekit/plugins/cartesia/version.py,sha256=cmcEldnzcpX8ulbgiuN3mLrZUz2H7FyBpXlMqC7B5K4,606
-livekit_plugins_cartesia-0.2.0.dev7.dist-info/METADATA,sha256=5tWoZTOirxenLlJEp8H0ZMMKK4TKNqkwG-kRizkQ1Mg,1257
-livekit_plugins_cartesia-0.2.0.dev7.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
-livekit_plugins_cartesia-0.2.0.dev7.dist-info/top_level.txt,sha256=OoDok3xUmXbZRvOrfvvXB-Juu4DX79dlq188E19YHoo,8
-livekit_plugins_cartesia-0.2.0.dev7.dist-info/RECORD,,